From b0efc764fbe70d793ea2ae542ae1664d845a8da4 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Wed, 8 Jul 2026 06:46:21 +0000
Subject: [PATCH 01/12] fix bugs

fix bugs
---
 .../kernels/primitive/datamover_primitives.h  | 20 ++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index c357f5142dcd1..2e8f15d36bb02 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -40,7 +40,7 @@ struct alignas(sizeof(T) * VecSize) VectorType {
  * must be [dim1, dim0].
  */
 struct BroadcastConfig {
-  funcs::FastDivMod<int> divmoders[DDim::kMaxRank];
+  funcs::FastDivMod<int64_t> divmoders[DDim::kMaxRank];
   uint64_t strides[DDim::kMaxRank];
   int rank{0};
 
@@ -51,7 +51,7 @@ struct BroadcastConfig {
                   const std::vector<int64_t>& in_dims,
                   int dim_size) {
     for (int i = 0; i < dim_size; ++i) {
-      divmoders[i] = funcs::FastDivMod<int>(out_dims[i]);
+      divmoders[i] = funcs::FastDivMod<int64_t>(out_dims[i]);
     }
 
     for (int i = 0; i < dim_size; ++i) {
@@ -59,7 +59,7 @@ struct BroadcastConfig {
       strides[i] = (i != 0 && strides[i] != 0)
                        ? std::accumulate(in_dims.begin(),
                                          in_dims.begin() + i,
-                                         1,
+                                         int64_t{1},
                                          std::multiplies<int64_t>())
                        : strides[i];
     }
@@ -418,13 +418,15 @@ __device__ __forceinline__ void ReadDataBc(
     int stride_nx,
     int stride_ny) {
   uint32_t thread_offset = block_offset + threadIdx.x;
-  uint32_t index_src = 0;
+  int64_t index_src = 0;
 
 #pragma unroll
   for (int ny = 0; ny < NY; ++ny) {
 #pragma unroll
     for (uint32_t nx = 0; nx < NX; ++nx) {
-      uint32_t index_output = thread_offset + ny * stride_ny + nx * stride_nx;
+      int64_t index_output = thread_offset +
+                             static_cast<int64_t>(ny) * stride_ny +
+                             static_cast<int64_t>(nx) * stride_nx;
       index_src = 0;
       if (IsBoundary) {
         if (index_output >= total_num_output) {
@@ -753,11 +755,11 @@ __device__ __forceinline__ void ReadDataBc(
     int total_num_output,
     int read_lens = NX) {
   uint32_t thread_offset = block_offset + threadIdx.x * NX;
-  uint32_t index_src = 0;
+  int64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    uint32_t index_output = thread_offset + nx;
+    int64_t index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {
@@ -814,11 +816,11 @@ __device__ __forceinline__ void ReadDataBc(
     int total_num_output,
     int read_lens = NX) {
   uint32_t thread_offset = block_offset + threadIdx.x * NX;
-  uint32_t index_src = 0;
+  int64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    uint32_t index_output = thread_offset + nx;
+    int64_t index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {

From bc16d0038e3e5a956aa1d6e8922f2024eea91580 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Thu, 9 Jul 2026 07:12:39 +0000
Subject: [PATCH 02/12] fix bugs

---
 .../phi/kernels/primitive/datamover_primitives.h | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index 2e8f15d36bb02..1c260b4ebe614 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -418,15 +418,15 @@ __device__ __forceinline__ void ReadDataBc(
     int stride_nx,
     int stride_ny) {
   uint32_t thread_offset = block_offset + threadIdx.x;
-  int64_t index_src = 0;
+  uint64_t index_src = 0;
 
 #pragma unroll
   for (int ny = 0; ny < NY; ++ny) {
 #pragma unroll
     for (uint32_t nx = 0; nx < NX; ++nx) {
-      int64_t index_output = thread_offset +
-                             static_cast<int64_t>(ny) * stride_ny +
-                             static_cast<int64_t>(nx) * stride_nx;
+      uint64_t index_output = thread_offset +
+                              static_cast<uint64_t>(ny) * stride_ny +
+                              static_cast<uint64_t>(nx) * stride_nx;
       index_src = 0;
       if (IsBoundary) {
         if (index_output >= total_num_output) {
@@ -755,11 +755,11 @@ __device__ __forceinline__ void ReadDataBc(
     int total_num_output,
     int read_lens = NX) {
   uint32_t thread_offset = block_offset + threadIdx.x * NX;
-  int64_t index_src = 0;
+  uint64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    int64_t index_output = thread_offset + nx;
+    uint64_t index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {
@@ -816,11 +816,11 @@ __device__ __forceinline__ void ReadDataBc(
     int total_num_output,
     int read_lens = NX) {
   uint32_t thread_offset = block_offset + threadIdx.x * NX;
-  int64_t index_src = 0;
+  uint64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    int64_t index_output = thread_offset + nx;
+    uint64_t index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {

From 95f21f79b2f15ce6e0d48551d1b236c70e1d929e Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Thu, 9 Jul 2026 08:03:29 +0000
Subject: [PATCH 03/12] fix bugs

---
 paddle/phi/kernels/primitive/datamover_primitives.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index 1c260b4ebe614..030114425c748 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -414,7 +414,7 @@ __device__ __forceinline__ void ReadDataBc(
     const T* __restrict__ src,
     uint32_t block_offset,
     const details::BroadcastConfig& config,
-    int total_num_output,
+    uint64_t total_num_output,
     int stride_nx,
     int stride_ny) {
   uint32_t thread_offset = block_offset + threadIdx.x;
@@ -752,7 +752,7 @@ __device__ __forceinline__ void ReadDataBc(
     const T* __restrict__ src,
     uint32_t block_offset,
     const details::BroadcastConfig& config,
-    int total_num_output,
+    uint64_t total_num_output,
     int read_lens = NX) {
   uint32_t thread_offset = block_offset + threadIdx.x * NX;
   uint64_t index_src = 0;
@@ -813,7 +813,7 @@ __device__ __forceinline__ void ReadDataBc(
     const T* __restrict__ src,
     uint32_t block_offset,
     const details::BroadcastConfig& config,
-    int total_num_output,
+    uint64_t total_num_output,
     int read_lens = NX) {
   uint32_t thread_offset = block_offset + threadIdx.x * NX;
   uint64_t index_src = 0;

From 54c3fee15a9450a393de485ca3ddeef1213bd302 Mon Sep 17 00:00:00 2001
From: feixi <66367509+feixi139@users.noreply.github.com>
Date: Mon, 13 Jul 2026 14:22:58 +0800
Subject: [PATCH 04/12] Update datamover_primitives.h

---
 paddle/phi/kernels/primitive/datamover_primitives.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index 030114425c748..5e6781da52fcc 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -1,5 +1,5 @@
 // Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
-//
+// 
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at

From 1333dd1a3cc2a89811376fdfb19c577094ebd339 Mon Sep 17 00:00:00 2001
From: feixi <66367509+feixi139@users.noreply.github.com>
Date: Mon, 13 Jul 2026 14:37:09 +0800
Subject: [PATCH 05/12] Update datamover_primitives.h

---
 paddle/phi/kernels/primitive/datamover_primitives.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index 5e6781da52fcc..030114425c748 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -1,5 +1,5 @@
 // Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
-// 
+//
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at

From ebebd101517ce89879ebb204b510d5fd43c185e0 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Mon, 13 Jul 2026 08:08:31 +0000
Subject: [PATCH 06/12] fix bugs

---
 paddle/phi/kernels/funcs/broadcast_function.h | 59 ++++++++++---------
 paddle/phi/kernels/funcs/elementwise_base.h   |  4 +-
 .../kernels/primitive/datamover_primitives.h  | 12 ++--
 3 files changed, 40 insertions(+), 35 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 1dabfeb425793..a6c8841870194 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -120,9 +120,9 @@ struct BroadcastDataLoader {
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               const int block_offset,
+                                               uint64_t block_offset,
                                                const int num,
-                                               const uint32_t numel,
+                                               uint64_t numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
 #ifdef PADDLE_WITH_XPU_KP
@@ -178,16 +178,17 @@ struct BroadcastDataLoader<Index, VecSize, true, kElementwise> {
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               const int block_offset,
+                                               uint64_t block_offset,
                                                const int num,
-                                               const uint32_t numel,
+                                               uint64_t numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
-    int thread_offset = threadIdx.x * VecSize + block_offset;
+    uint64_t thread_offset =
+        static_cast<uint64_t>(threadIdx.x) * VecSize + block_offset;
 #pragma unroll
     for (int idx = 0; idx < VecSize; ++idx) {
       std::get<Index>(args[idx]) = static_cast<Type>(1);
-      int index = thread_offset + idx;
+      uint64_t index = thread_offset + idx;
       if (index < numel) {
         std::get<Index>(args[idx]) =
             reinterpret_cast<const _ptr_ Type *>(ins[Index])[index];
@@ -204,9 +205,9 @@ struct BroadcastDataLoader<Index, VecSize, false, kElementwise> {
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               const int block_offset,
+                                               uint64_t block_offset,
                                                const int num,
-                                               const uint32_t numel,
+                                               uint64_t numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
     using VecType = phi::kps::details::VectorType<Type, VecSize>;
@@ -242,7 +243,7 @@ struct BroadcastDataSetter {
   template <typename Array, typename ArgsT>
   static __device__ __forceinline__ void Apply(const Array &ins,
                                                ArgsT *args,
-                                               uint32_t index_bc[][VecSize]) {
+                                               uint64_t index_bc[][VecSize]) {
     using Type = std::tuple_element_t<Index, ArgsT>;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
@@ -294,10 +295,10 @@ __device__ void VectorizedBroadcastKernelImpl(
     const Array<const _ptr_ char *__restrict__, Arity> &ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     const Array<bool, Arity> &use_broadcast,
-    const uint32_t numel,
+    const uint64_t numel,
     const Array<kps::details::BroadcastConfig, Arity> &configs,
     uint32_t num,
-    uint32_t block_offset,
+    uint64_t block_offset,
     int read_lens,
     Functor func) {
   using Traits = funcs::FunctionTraits<Functor>;
@@ -310,12 +311,12 @@ __device__ void VectorizedBroadcastKernelImpl(
       ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
 #else
   if (LoadType == kBroadcast) {
-    uint32_t index_bc[Arity][VecSize] = {0};
+    uint64_t index_bc[Arity][VecSize] = {0};
     Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
-    uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
+    uint64_t thread_offset = block_offset + threadIdx.x * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
-      uint32_t idx = thread_offset + k;
+      uint64_t idx = thread_offset + k;
       if (IsBoundary && idx == numel) break;
 #pragma unroll
       for (int i = 0; i < DDim::kMaxRank; ++i) {
@@ -353,15 +354,16 @@ __global__ void VectorizedBroadcastKernel(
     Array<const _ptr_ char *__restrict__, Arity> ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     Array<bool, Arity> use_broadcast,
-    uint32_t numel,
+    uint64_t numel,
     Array<kps::details::BroadcastConfig, Arity> configs,
-    uint32_t main_offset,
-    uint32_t tail_tid,
+    uint64_t main_offset,
+    uint64_t tail_tid,
     int read_lens,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int64_t block_offset = BLOCK_ID_X * BLOCK_NUM_X * read_lens;
-  int64_t stride = BLOCK_NUM_X * GRID_NUM_X * read_lens;
+  uint64_t block_offset =
+      static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
+  uint64_t stride = static_cast<uint64_t>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
   for (; block_offset < main_offset; block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -379,7 +381,7 @@ __global__ void VectorizedBroadcastKernel(
                                             read_lens,
                                             func);
   }
-  int64_t num = numel - block_offset;
+  uint64_t num = numel - block_offset;
   if (num > 0) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -398,7 +400,8 @@ __global__ void VectorizedBroadcastKernel(
                                             func);
   }
 #else
-  int64_t block_offset = BLOCK_ID_X * BLOCK_NUM_X * VecSize;
+  uint64_t block_offset =
+      static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
   if (block_offset < main_offset) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -441,13 +444,14 @@ void LaunchBroadcastKernel(
     const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int numel = classifier.numel;
+  const int64_t numel = classifier.numel;
   const int threads = 64;
   const int blocks = 8;
   int read_lens = configs[0].buf_len;
   auto stream = dev_ctx.x_context()->xpu_stream;
-  uint32_t main_offset = (numel / (read_lens * threads)) * read_lens * threads;
-  uint32_t tail_tid = numel % (read_lens * threads);
+  const int64_t block_len = static_cast<int64_t>(read_lens) * threads;
+  const int64_t main_offset = (numel / block_len) * block_len;
+  const int64_t tail_tid = numel % block_len;
 
   VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
       <<<blocks, threads, 0, stream>>>(classifier.ins_data,
@@ -460,14 +464,15 @@ void LaunchBroadcastKernel(
                                        read_lens,
                                        func);
 #else
-  const int64_t &numel = classifier.numel;
+  const int64_t numel = classifier.numel;
   auto gpu_config =
       phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, VecSize);
   auto stream = dev_ctx.stream();
   auto threads = gpu_config.GetBlockSize();
   auto blocks = gpu_config.block_per_grid;
-  uint32_t main_offset = (numel / (VecSize * threads)) * VecSize * threads;
-  uint32_t tail_tid = numel % (VecSize * threads);
+  const int64_t block_len = static_cast<int64_t>(VecSize) * threads;
+  const int64_t main_offset = (numel / block_len) * block_len;
+  const int64_t tail_tid = numel % block_len;
 
   if (classifier.all_elementwise) {
     VectorizedBroadcastKernel<Functor,
diff --git a/paddle/phi/kernels/funcs/elementwise_base.h b/paddle/phi/kernels/funcs/elementwise_base.h
index d54fa602a840c..acab9b7cb02d0 100644
--- a/paddle/phi/kernels/funcs/elementwise_base.h
+++ b/paddle/phi/kernels/funcs/elementwise_base.h
@@ -679,7 +679,7 @@ struct ElementwiseWriteDataCallerBc {
   __device__ __forceinline__ void operator()(
       Array<_ptr_ OutT *, NumOuts> outs,
       ConditionalT<OutT, NumOuts> src[VecSize],
-      kps::IndexType block_offset,
+      int64_t block_offset,
       int num,
       int read_lens) {
     OutT dst[NumOuts][VecSize];
@@ -702,7 +702,7 @@ template <typename OutT, int VecSize, bool IsBoundary>
 struct ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, 1> {
   __device__ __forceinline__ void operator()(Array<_ptr_ OutT *, 1> outs,
                                              OutT src[VecSize],
-                                             kps::IndexType block_offset,
+                                             int64_t block_offset,
                                              int num,
                                              int read_lens) {
     kps::WriteData<OutT, VecSize, 1, IsBoundary>(
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index 030114425c748..c0d7f1ae7bb88 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -412,12 +412,12 @@ template <typename T, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
-    uint32_t block_offset,
+    uint64_t block_offset,
     const details::BroadcastConfig& config,
     uint64_t total_num_output,
     int stride_nx,
     int stride_ny) {
-  uint32_t thread_offset = block_offset + threadIdx.x;
+  uint64_t thread_offset = block_offset + threadIdx.x;
   uint64_t index_src = 0;
 
 #pragma unroll
@@ -750,11 +750,11 @@ template <typename T, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
-    uint32_t block_offset,
+    uint64_t block_offset,
     const details::BroadcastConfig& config,
     uint64_t total_num_output,
     int read_lens = NX) {
-  uint32_t thread_offset = block_offset + threadIdx.x * NX;
+  uint64_t thread_offset = block_offset + threadIdx.x * NX;
   uint64_t index_src = 0;
 
 #pragma unroll
@@ -811,11 +811,11 @@ template <typename T,
 __device__ __forceinline__ void ReadDataBc(
     ArgsT* dst,
     const T* __restrict__ src,
-    uint32_t block_offset,
+    uint64_t block_offset,
     const details::BroadcastConfig& config,
     uint64_t total_num_output,
     int read_lens = NX) {
-  uint32_t thread_offset = block_offset + threadIdx.x * NX;
+  uint64_t thread_offset = block_offset + threadIdx.x * NX;
   uint64_t index_src = 0;
 
 #pragma unroll

From 4ef6525f41782f2f257755fa7ecc541024011a2d Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Mon, 13 Jul 2026 08:30:18 +0000
Subject: [PATCH 07/12] fix bugs

---
 paddle/phi/kernels/funcs/broadcast_function.h | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index a6c8841870194..2f21f653a9c60 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include <sstream>
+#include "paddle/common/enforce.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 
 #if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
@@ -381,8 +382,8 @@ __global__ void VectorizedBroadcastKernel(
                                             read_lens,
                                             func);
   }
-  uint64_t num = numel - block_offset;
-  if (num > 0) {
+  if (block_offset < numel) {
+    uint64_t num = numel - block_offset;
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   Arity,
@@ -445,6 +446,7 @@ void LaunchBroadcastKernel(
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
   const int64_t numel = classifier.numel;
+  PADDLE_ENFORCE_LE_INT_MAX(numel, "BroadcastKernel numel (XPU)");
   const int threads = 64;
   const int blocks = 8;
   int read_lens = configs[0].buf_len;

From e243bfa375f50388a4ca0ba97bf037610b78d3e5 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Mon, 13 Jul 2026 09:24:10 +0000
Subject: [PATCH 08/12] fix bugs

---
 paddle/phi/kernels/funcs/broadcast_function.h       | 11 ++++++++---
 paddle/phi/kernels/primitive/datamover_primitives.h |  8 +++++---
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 2f21f653a9c60..525b52f3551b0 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -48,7 +48,12 @@ struct BroadcastTypeClassifier {
                           int axis) {
     numel = (*outs)[0]->numel();
 
-#ifndef PADDLE_WITH_XPU_KP
+#ifdef PADDLE_WITH_XPU_KP
+    // datamover_primitives_xpu2.h::BroadcastConfig (built inside
+    // InitBroadcastConfigs below) still computes strides/numel with 32-bit
+    // int, so the INT_MAX check must happen before that call, not after.
+    PADDLE_ENFORCE_LE_INT_MAX(numel, "BroadcastKernel numel (XPU)");
+#else
     for (size_t i = 0; i < ins.size(); ++i) {
       bool is_same_dim = ins[i]->numel() == numel;
       if (is_same_dim) {
@@ -314,7 +319,8 @@ __device__ void VectorizedBroadcastKernelImpl(
   if (LoadType == kBroadcast) {
     uint64_t index_bc[Arity][VecSize] = {0};
     Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
-    uint64_t thread_offset = block_offset + threadIdx.x * VecSize;
+    uint64_t thread_offset =
+        block_offset + static_cast<uint64_t>(threadIdx.x) * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
       uint64_t idx = thread_offset + k;
@@ -446,7 +452,6 @@ void LaunchBroadcastKernel(
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
   const int64_t numel = classifier.numel;
-  PADDLE_ENFORCE_LE_INT_MAX(numel, "BroadcastKernel numel (XPU)");
   const int threads = 64;
   const int blocks = 8;
   int read_lens = configs[0].buf_len;
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index c0d7f1ae7bb88..c8bfd42f11cf0 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -417,7 +417,7 @@ __device__ __forceinline__ void ReadDataBc(
     uint64_t total_num_output,
     int stride_nx,
     int stride_ny) {
-  uint64_t thread_offset = block_offset + threadIdx.x;
+  uint64_t thread_offset = block_offset + static_cast<uint64_t>(threadIdx.x);
   uint64_t index_src = 0;
 
 #pragma unroll
@@ -754,7 +754,8 @@ __device__ __forceinline__ void ReadDataBc(
     const details::BroadcastConfig& config,
     uint64_t total_num_output,
     int read_lens = NX) {
-  uint64_t thread_offset = block_offset + threadIdx.x * NX;
+  uint64_t thread_offset =
+      block_offset + static_cast<uint64_t>(threadIdx.x) * NX;
   uint64_t index_src = 0;
 
 #pragma unroll
@@ -815,7 +816,8 @@ __device__ __forceinline__ void ReadDataBc(
     const details::BroadcastConfig& config,
     uint64_t total_num_output,
     int read_lens = NX) {
-  uint64_t thread_offset = block_offset + threadIdx.x * NX;
+  uint64_t thread_offset =
+      block_offset + static_cast<uint64_t>(threadIdx.x) * NX;
   uint64_t index_src = 0;
 
 #pragma unroll

From 6848e24d263172dbae7a424462b4a4b8f5dd15d8 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Wed, 15 Jul 2026 06:55:05 +0000
Subject: [PATCH 09/12] fix bugs

---
 paddle/phi/kernels/funcs/broadcast_function.h | 284 +++++++++++-------
 paddle/phi/kernels/funcs/elementwise_base.h   |   4 +-
 .../kernels/primitive/datamover_primitives.h  |  35 ++-
 3 files changed, 202 insertions(+), 121 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 525b52f3551b0..ef23bf789b946 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 
+#include <limits>
 #include <sstream>
 #include "paddle/common/enforce.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
@@ -119,16 +120,20 @@ struct BroadcastTypeClassifier {
 };
 
 // Common broadcast/elementwise Loader.
-template <int Index, int VecSize, bool IsBoundary, int LoadType>
+template <typename IndexT,
+          int Index,
+          int VecSize,
+          bool IsBoundary,
+          int LoadType>
 struct BroadcastDataLoader {
   template <typename Array1, typename Array2, typename Array3, typename ArgsT>
   static __device__ __forceinline__ void Apply(const Array1 &ins,
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               uint64_t block_offset,
+                                               IndexT block_offset,
                                                const int num,
-                                               uint64_t numel,
+                                               IndexT numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
 #ifdef PADDLE_WITH_XPU_KP
@@ -152,7 +157,7 @@ struct BroadcastDataLoader {
 #else
     kps::Init<Type, ArgsT, Index, VecSize>(args, static_cast<Type>(1.0f));
     if (use_broadcast[Index]) {
-      kps::ReadDataBc<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
+      kps::ReadDataBc<Type, IndexT, VecSize, 1, ArgsT, Index, IsBoundary>(
           args,
           reinterpret_cast<const _ptr_ Type *>(ins[Index]),
           block_offset,
@@ -177,24 +182,24 @@ struct BroadcastDataLoader {
 /* BroadcastDataLoaders Partial specialization */
 #ifndef PADDLE_WITH_XPU_KP
 // Scalar elementwise Loader with consideration of IsBoundary.
-template <int Index, int VecSize>
-struct BroadcastDataLoader<Index, VecSize, true, kElementwise> {
+template <typename IndexT, int Index, int VecSize>
+struct BroadcastDataLoader<IndexT, Index, VecSize, true, kElementwise> {
   template <typename Array1, typename Array2, typename Array3, typename ArgsT>
   static __device__ __forceinline__ void Apply(const Array1 &ins,
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               uint64_t block_offset,
+                                               IndexT block_offset,
                                                const int num,
-                                               uint64_t numel,
+                                               IndexT numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
-    uint64_t thread_offset =
-        static_cast<uint64_t>(threadIdx.x) * VecSize + block_offset;
+    IndexT thread_offset =
+        static_cast<IndexT>(threadIdx.x) * VecSize + block_offset;
 #pragma unroll
     for (int idx = 0; idx < VecSize; ++idx) {
       std::get<Index>(args[idx]) = static_cast<Type>(1);
-      uint64_t index = thread_offset + idx;
+      IndexT index = thread_offset + idx;
       if (index < numel) {
         std::get<Index>(args[idx]) =
             reinterpret_cast<const _ptr_ Type *>(ins[Index])[index];
@@ -204,24 +209,23 @@ struct BroadcastDataLoader<Index, VecSize, true, kElementwise> {
 };
 
 // Vectorized elementwise Loader without consideration of IsBoundary.
-template <int Index, int VecSize>
-struct BroadcastDataLoader<Index, VecSize, false, kElementwise> {
+template <typename IndexT, int Index, int VecSize>
+struct BroadcastDataLoader<IndexT, Index, VecSize, false, kElementwise> {
   template <typename Array1, typename Array2, typename Array3, typename ArgsT>
   static __device__ __forceinline__ void Apply(const Array1 &ins,
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               uint64_t block_offset,
+                                               IndexT block_offset,
                                                const int num,
-                                               uint64_t numel,
+                                               IndexT numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
     using VecType = phi::kps::details::VectorType<Type, VecSize>;
     VecType vec_temp;
 
-    int64_t thread_offset =
-        static_cast<int64_t>(threadIdx.x) +
-        static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
+    IndexT thread_offset =
+        block_offset / VecSize + static_cast<IndexT>(threadIdx.x);
     const VecType *__restrict__ vec_input =
         reinterpret_cast<const VecType *__restrict__>(ins[Index]);
     vec_temp = vec_input[thread_offset];
@@ -262,8 +266,13 @@ struct BroadcastDataSetter {
 #endif
 
 // static broadcast unroller
-template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
+template <template <typename IndexT,
+                    int Index,
+                    int VecSize,
+                    bool IsBoundary,
+                    int LoadType>
           typename Func,
+          typename IndexT,
           bool IsBoundary,
           int LoadType,
           int VecSize,
@@ -272,26 +281,32 @@ template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
 struct BcUnroller {
   template <typename... Args>
   static HOSTDEVICE inline void step(Args &&...args) {
-    Func<Begin, VecSize, IsBoundary, LoadType>::Apply(
+    Func<IndexT, Begin, VecSize, IsBoundary, LoadType>::Apply(
         std::forward<Args>(args)...);
-    BcUnroller<Func, IsBoundary, LoadType, VecSize, End, Begin + 1>::step(
-        args...);
+    BcUnroller<Func, IndexT, IsBoundary, LoadType, VecSize, End, Begin + 1>::
+        step(args...);
   }
 };
 
-template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
+template <template <typename IndexT,
+                    int Index,
+                    int VecSize,
+                    bool IsBoundary,
+                    int LoadType>
           typename Func,
+          typename IndexT,
           bool IsBoundary,
           int LoadType,
           int VecSize,
           int End>
-struct BcUnroller<Func, IsBoundary, LoadType, VecSize, End, End> {
+struct BcUnroller<Func, IndexT, IsBoundary, LoadType, VecSize, End, End> {
   template <typename... Args>
   static HOSTDEVICE inline void step(Args &&...args) {}
 };
 
 template <typename OutT,
           typename Functor,
+          typename IndexT,
           int Arity,
           int NumOuts,
           int VecSize,
@@ -301,10 +316,10 @@ __device__ void VectorizedBroadcastKernelImpl(
     const Array<const _ptr_ char *__restrict__, Arity> &ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     const Array<bool, Arity> &use_broadcast,
-    const uint64_t numel,
+    const IndexT numel,
     const Array<kps::details::BroadcastConfig, Arity> &configs,
     uint32_t num,
-    uint64_t block_offset,
+    IndexT block_offset,
     int read_lens,
     Functor func) {
   using Traits = funcs::FunctionTraits<Functor>;
@@ -313,14 +328,25 @@ __device__ void VectorizedBroadcastKernelImpl(
   __simd__ ConditionalT<OutT, NumOuts> result[VecSize];
 
 #ifdef PADDLE_WITH_XPU_KP
-  BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
-      ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
+  BcUnroller<BroadcastDataLoader,
+             IndexT,
+             IsBoundary,
+             LoadType,
+             VecSize,
+             Arity>::step(ins,
+                          args,
+                          configs,
+                          use_broadcast,
+                          block_offset,
+                          num,
+                          numel,
+                          read_lens);
 #else
   if (LoadType == kBroadcast) {
     uint64_t index_bc[Arity][VecSize] = {0};
     Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
-    uint64_t thread_offset =
-        block_offset + static_cast<uint64_t>(threadIdx.x) * VecSize;
+    IndexT thread_offset =
+        block_offset + static_cast<IndexT>(threadIdx.x) * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
       uint64_t idx = thread_offset + k;
@@ -338,8 +364,19 @@ __device__ void VectorizedBroadcastKernelImpl(
     }
     Unroller<BroadcastDataSetter, VecSize, Arity>::step(ins, args, index_bc);
   } else {
-    BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
-        ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
+    BcUnroller<BroadcastDataLoader,
+               IndexT,
+               IsBoundary,
+               LoadType,
+               VecSize,
+               Arity>::step(ins,
+                            args,
+                            configs,
+                            use_broadcast,
+                            block_offset,
+                            num,
+                            numel,
+                            read_lens);
   }
 #endif
   SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
@@ -353,6 +390,7 @@ __device__ void VectorizedBroadcastKernelImpl(
 
 template <typename Functor,
           typename OutT,
+          typename IndexT,
           int Arity,
           int NumOuts,
           int VecSize,
@@ -361,19 +399,21 @@ __global__ void VectorizedBroadcastKernel(
     Array<const _ptr_ char *__restrict__, Arity> ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     Array<bool, Arity> use_broadcast,
-    uint64_t numel,
+    IndexT numel,
     Array<kps::details::BroadcastConfig, Arity> configs,
-    uint64_t main_offset,
-    uint64_t tail_tid,
+    IndexT main_offset,
+    IndexT tail_tid,
     int read_lens,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
   uint64_t block_offset =
       static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
   uint64_t stride = static_cast<uint64_t>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
-  for (; block_offset < main_offset; block_offset += stride) {
+  for (; block_offset < static_cast<uint64_t>(main_offset);
+       block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
+                                  IndexT,
                                   Arity,
                                   NumOuts,
                                   VecSize,
@@ -384,14 +424,16 @@ __global__ void VectorizedBroadcastKernel(
                                             numel,
                                             configs,
                                             BLOCK_NUM_X * read_lens,
-                                            block_offset,
+                                            static_cast<IndexT>(block_offset),
                                             read_lens,
                                             func);
   }
-  if (block_offset < numel) {
-    uint64_t num = numel - block_offset;
+  int64_t num =
+      static_cast<int64_t>(numel) - static_cast<int64_t>(block_offset);
+  if (num > 0) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
+                                  IndexT,
                                   Arity,
                                   NumOuts,
                                   VecSize,
@@ -401,17 +443,20 @@ __global__ void VectorizedBroadcastKernel(
                                             use_broadcast,
                                             numel,
                                             configs,
-                                            num,
-                                            block_offset,
+                                            static_cast<uint32_t>(num),
+                                            static_cast<IndexT>(block_offset),
                                             read_lens,
                                             func);
   }
 #else
   uint64_t block_offset =
       static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
-  if (block_offset < main_offset) {
+  uint64_t stride = static_cast<uint64_t>(BLOCK_NUM_X) * GRID_NUM_X * VecSize;
+  for (; block_offset < static_cast<uint64_t>(main_offset);
+       block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
+                                  IndexT,
                                   Arity,
                                   NumOuts,
                                   VecSize,
@@ -422,12 +467,16 @@ __global__ void VectorizedBroadcastKernel(
                                             numel,
                                             configs,
                                             BLOCK_NUM_X * VecSize,
-                                            block_offset,
+                                            static_cast<IndexT>(block_offset),
                                             read_lens,
                                             func);
-  } else {
+  }
+  int64_t num =
+      static_cast<int64_t>(numel) - static_cast<int64_t>(block_offset);
+  if (num > 0) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
+                                  IndexT,
                                   Arity,
                                   NumOuts,
                                   VecSize,
@@ -437,8 +486,8 @@ __global__ void VectorizedBroadcastKernel(
                                             use_broadcast,
                                             numel,
                                             configs,
-                                            tail_tid,
-                                            block_offset,
+                                            static_cast<uint32_t>(num),
+                                            static_cast<IndexT>(block_offset),
                                             read_lens,
                                             func);
   }
@@ -450,44 +499,29 @@ void LaunchBroadcastKernel(
     const KPDevice &dev_ctx,
     const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
     Functor func) {
+#ifndef PADDLE_WITH_XPU_KP
+  auto gpu_config = phi::backends::gpu::GetGpuLaunchConfig1D(
+      dev_ctx, classifier.numel, VecSize);
+#endif
+  auto launch_with_index_t = [&](auto index_tag) {
+    using IndexT = decltype(index_tag);
 #ifdef PADDLE_WITH_XPU_KP
-  const int64_t numel = classifier.numel;
-  const int threads = 64;
-  const int blocks = 8;
-  int read_lens = configs[0].buf_len;
-  auto stream = dev_ctx.x_context()->xpu_stream;
-  const int64_t block_len = static_cast<int64_t>(read_lens) * threads;
-  const int64_t main_offset = (numel / block_len) * block_len;
-  const int64_t tail_tid = numel % block_len;
-
-  VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
-      <<<blocks, threads, 0, stream>>>(classifier.ins_data,
-                                       classifier.outs_data,
-                                       classifier.use_broadcast,
-                                       numel,
-                                       classifier.configs,
-                                       main_offset,
-                                       tail_tid,
-                                       read_lens,
-                                       func);
-#else
-  const int64_t numel = classifier.numel;
-  auto gpu_config =
-      phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, VecSize);
-  auto stream = dev_ctx.stream();
-  auto threads = gpu_config.GetBlockSize();
-  auto blocks = gpu_config.block_per_grid;
-  const int64_t block_len = static_cast<int64_t>(VecSize) * threads;
-  const int64_t main_offset = (numel / block_len) * block_len;
-  const int64_t tail_tid = numel % block_len;
-
-  if (classifier.all_elementwise) {
+    const IndexT numel = static_cast<IndexT>(classifier.numel);
+    const int threads = 64;
+    const int blocks = 8;
+    int read_lens = classifier.configs[0].buf_len;
+    auto stream = dev_ctx.x_context()->xpu_stream;
+    const IndexT block_len = static_cast<IndexT>(read_lens) * threads;
+    const IndexT main_offset = (numel / block_len) * block_len;
+    const IndexT tail_tid = numel % block_len;
+
     VectorizedBroadcastKernel<Functor,
                               OutT,
+                              IndexT,
                               Arity,
                               NumOuts,
                               VecSize,
-                              kElementwise>
+                              false>
         <<<blocks, threads, 0, stream>>>(classifier.ins_data,
                                          classifier.outs_data,
                                          classifier.use_broadcast,
@@ -495,31 +529,79 @@ void LaunchBroadcastKernel(
                                          classifier.configs,
                                          main_offset,
                                          tail_tid,
-                                         VecSize,
-                                         func);
-  } else if (classifier.broadcast_num > (Arity >> 1)) {
-    constexpr BroadcastType type_ = (Arity > 1) ? kBroadcast : kMixed;
-    VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, type_>
-        <<<blocks, threads, 0, stream>>>(classifier.ins_data,
-                                         classifier.outs_data,
-                                         classifier.use_broadcast,
-                                         numel,
-                                         classifier.configs,
-                                         main_offset,
-                                         tail_tid,
-                                         VecSize,
+                                         read_lens,
                                          func);
+#else
+    const IndexT numel = static_cast<IndexT>(classifier.numel);
+    auto stream = dev_ctx.stream();
+    auto threads = gpu_config.GetBlockSize();
+    auto blocks = gpu_config.block_per_grid;
+    const IndexT block_len = static_cast<IndexT>(VecSize) * threads;
+    const IndexT main_offset = (numel / block_len) * block_len;
+    const IndexT tail_tid = numel % block_len;
+
+    if (classifier.all_elementwise) {
+      VectorizedBroadcastKernel<Functor,
+                                OutT,
+                                IndexT,
+                                Arity,
+                                NumOuts,
+                                VecSize,
+                                kElementwise>
+          <<<blocks, threads, 0, stream>>>(classifier.ins_data,
+                                           classifier.outs_data,
+                                           classifier.use_broadcast,
+                                           numel,
+                                           classifier.configs,
+                                           main_offset,
+                                           tail_tid,
+                                           VecSize,
+                                           func);
+    } else if (classifier.broadcast_num > (Arity >> 1)) {
+      constexpr BroadcastType type_ = (Arity > 1) ? kBroadcast : kMixed;
+      VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, type_>
+          <<<blocks, threads, 0, stream>>>(classifier.ins_data,
+                                           classifier.outs_data,
+                                           classifier.use_broadcast,
+                                           numel,
+                                           classifier.configs,
+                                           main_offset,
+                                           tail_tid,
+                                           VecSize,
+                                           func);
+    } else {
+      VectorizedBroadcastKernel<Functor,
+                                OutT,
+                                IndexT,
+                                Arity,
+                                NumOuts,
+                                VecSize,
+                                kMixed>
+          <<<blocks, threads, 0, stream>>>(classifier.ins_data,
+                                           classifier.outs_data,
+                                           classifier.use_broadcast,
+                                           numel,
+                                           classifier.configs,
+                                           main_offset,
+                                           tail_tid,
+                                           VecSize,
+                                           func);
+    }
+#endif
+  };
+
+#ifdef PADDLE_WITH_XPU_KP
+  launch_with_index_t(uint32_t{0});
+#else
+  const uint64_t index_limit = std::numeric_limits<uint32_t>::max();
+  const uint64_t grid_stride =
+      static_cast<uint64_t>(gpu_config.block_per_grid.x) *
+      static_cast<uint64_t>(gpu_config.GetBlockSize()) * VecSize;
+  if (static_cast<uint64_t>(classifier.numel) <= index_limit &&
+      grid_stride <= index_limit) {
+    launch_with_index_t(uint32_t{0});
   } else {
-    VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, kMixed>
-        <<<blocks, threads, 0, stream>>>(classifier.ins_data,
-                                         classifier.outs_data,
-                                         classifier.use_broadcast,
-                                         numel,
-                                         classifier.configs,
-                                         main_offset,
-                                         tail_tid,
-                                         VecSize,
-                                         func);
+    launch_with_index_t(uint64_t{0});
   }
 #endif
 }
diff --git a/paddle/phi/kernels/funcs/elementwise_base.h b/paddle/phi/kernels/funcs/elementwise_base.h
index acab9b7cb02d0..d54fa602a840c 100644
--- a/paddle/phi/kernels/funcs/elementwise_base.h
+++ b/paddle/phi/kernels/funcs/elementwise_base.h
@@ -679,7 +679,7 @@ struct ElementwiseWriteDataCallerBc {
   __device__ __forceinline__ void operator()(
       Array<_ptr_ OutT *, NumOuts> outs,
       ConditionalT<OutT, NumOuts> src[VecSize],
-      int64_t block_offset,
+      kps::IndexType block_offset,
       int num,
       int read_lens) {
     OutT dst[NumOuts][VecSize];
@@ -702,7 +702,7 @@ template <typename OutT, int VecSize, bool IsBoundary>
 struct ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, 1> {
   __device__ __forceinline__ void operator()(Array<_ptr_ OutT *, 1> outs,
                                              OutT src[VecSize],
-                                             int64_t block_offset,
+                                             kps::IndexType block_offset,
                                              int num,
                                              int read_lens) {
     kps::WriteData<OutT, VecSize, 1, IsBoundary>(
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index c8bfd42f11cf0..55bc0e3fc71f4 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -408,25 +408,25 @@ __device__ __forceinline__ void ReadData(ArgsT* dst,
  * stride_nx: Each read one element stride stride_nx elements in the last dim.
  * stride_ny: Each read one element stride stride_ny elements in the first dim.
  */
-template <typename T, int NX, int NY, bool IsBoundary = false>
+template <typename T, typename IndexT, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
-    uint64_t block_offset,
+    IndexT block_offset,
     const details::BroadcastConfig& config,
-    uint64_t total_num_output,
+    IndexT total_num_output,
     int stride_nx,
     int stride_ny) {
-  uint64_t thread_offset = block_offset + static_cast<uint64_t>(threadIdx.x);
+  IndexT thread_offset = block_offset + static_cast<IndexT>(threadIdx.x);
   uint64_t index_src = 0;
 
 #pragma unroll
   for (int ny = 0; ny < NY; ++ny) {
 #pragma unroll
     for (uint32_t nx = 0; nx < NX; ++nx) {
-      uint64_t index_output = thread_offset +
-                              static_cast<uint64_t>(ny) * stride_ny +
-                              static_cast<uint64_t>(nx) * stride_nx;
+      IndexT index_output = thread_offset +
+                            static_cast<IndexT>(ny) * stride_ny +
+                            static_cast<IndexT>(nx) * stride_nx;
       index_src = 0;
       if (IsBoundary) {
         if (index_output >= total_num_output) {
@@ -746,21 +746,20 @@ __device__ __forceinline__ void Init(T* dst, T* init_data, int num) {
  * coordinate mapping relationship between output data and input data.
  * total_num_output: Total number of original output.
  */
-template <typename T, int NX, int NY, bool IsBoundary = false>
+template <typename T, typename IndexT, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
-    uint64_t block_offset,
+    IndexT block_offset,
     const details::BroadcastConfig& config,
-    uint64_t total_num_output,
+    IndexT total_num_output,
     int read_lens = NX) {
-  uint64_t thread_offset =
-      block_offset + static_cast<uint64_t>(threadIdx.x) * NX;
+  IndexT thread_offset = block_offset + static_cast<IndexT>(threadIdx.x) * NX;
   uint64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    uint64_t index_output = thread_offset + nx;
+    IndexT index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {
@@ -804,6 +803,7 @@ __device__ __forceinline__ void ReadDataBc(
  */
 
 template <typename T,
+          typename IndexT,
           int NX,
           int NY,
           typename ArgsT,
@@ -812,17 +812,16 @@ template <typename T,
 __device__ __forceinline__ void ReadDataBc(
     ArgsT* dst,
     const T* __restrict__ src,
-    uint64_t block_offset,
+    IndexT block_offset,
     const details::BroadcastConfig& config,
-    uint64_t total_num_output,
+    IndexT total_num_output,
     int read_lens = NX) {
-  uint64_t thread_offset =
-      block_offset + static_cast<uint64_t>(threadIdx.x) * NX;
+  IndexT thread_offset = block_offset + static_cast<IndexT>(threadIdx.x) * NX;
   uint64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    uint64_t index_output = thread_offset + nx;
+    IndexT index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {

From e8ece3906f505712f9554b2c40b169fc471fb4fa Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Wed, 15 Jul 2026 07:43:58 +0000
Subject: [PATCH 10/12] fix bugs

---
 paddle/phi/kernels/funcs/broadcast_function.h | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index ef23bf789b946..0dc40322eb2f1 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -559,7 +559,13 @@ void LaunchBroadcastKernel(
                                            func);
     } else if (classifier.broadcast_num > (Arity >> 1)) {
       constexpr BroadcastType type_ = (Arity > 1) ? kBroadcast : kMixed;
-      VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, type_>
+      VectorizedBroadcastKernel<Functor,
+                                OutT,
+                                IndexT,
+                                Arity,
+                                NumOuts,
+                                VecSize,
+                                type_>
           <<<blocks, threads, 0, stream>>>(classifier.ins_data,
                                            classifier.outs_data,
                                            classifier.use_broadcast,

From 21b076b4dc618e2cec4de074cf5e5fcbbb5e7392 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Wed, 15 Jul 2026 10:35:00 +0000
Subject: [PATCH 11/12] Align BroadcastConfig index type with kernel IndexT

Make the GPU BroadcastConfig a template on IndexT so normal-sized tensors
use 32-bit FastDivMod again instead of always paying for 64-bit divmod.
Build the typed configs at launch time (once IndexT is chosen) from the
simplified dims stored on the classifier, and tighten the uint32 fast-path
gate to INT_MAX for 32-bit FastDivMod ctor safety.
---
 paddle/phi/kernels/funcs/broadcast_function.h | 70 +++++++++++++++----
 .../kernels/primitive/datamover_primitives.h  | 30 +++++---
 2 files changed, 75 insertions(+), 25 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index 0dc40322eb2f1..fb85952f919a7 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -33,13 +33,33 @@ namespace funcs {
 
 enum BroadcastType { kMixed = 1, kBroadcast = 2, kElementwise = 3 };
 
+// On XPU the BroadcastConfig lives in datamover_primitives_xpu2.h and is not
+// templated; on GPU it is templated on the index type. This alias lets the
+// shared kernel signatures name the right type on each backend.
+#ifdef PADDLE_WITH_XPU_KP
+template <typename IndexT>
+using BroadcastConfigType = kps::details::BroadcastConfig;
+#else
+template <typename IndexT>
+using BroadcastConfigType = kps::details::BroadcastConfig<IndexT>;
+#endif
+
 template <typename OutT, typename Functor, int Arity, int NumOuts>
 struct BroadcastTypeClassifier {
   int64_t numel{0};
   int broadcast_num{0};              // Not used for XPU
   bool all_elementwise{true};        // Not used for XPU
   Array<bool, Arity> use_broadcast;  // Not used for XPU
+#ifdef PADDLE_WITH_XPU_KP
   Array<kps::details::BroadcastConfig, Arity> configs;
+#else
+  // The GPU BroadcastConfig is templated on the index type, which is only
+  // decided at launch time (uint32 vs uint64). Keep the simplified dims here
+  // and build the typed configs later in LaunchBroadcastKernel.
+  std::vector<int64_t> out_dims;
+  std::vector<std::vector<int64_t>> in_dims;
+  int rank{0};
+#endif
   Array<const _ptr_ char *__restrict__, Arity> ins_data;
   Array<_ptr_ OutT *, NumOuts> outs_data;
 
@@ -104,14 +124,17 @@ struct BroadcastTypeClassifier {
         DimsSimplifiedLogger<int64_t>::Log(
             ins, outs, dims_simplifier, "BroadcastKernel");
       }
+      // Store the simplified dims; the typed BroadcastConfig is built once the
+      // index type is chosen in LaunchBroadcastKernel. An empty in_dims[i]
+      // means "no config needed" (mirrors the old `ins[i]->numel()` guard).
+      out_dims = dims_simplifier.out_dims;
+      rank = dims_simplifier.rank;
+      in_dims.resize(Arity);
       for (int i = 0; i < Arity; ++i) {
         // if data shape is[m, n], then you should set data_dim = {n, m}
         // eg: out's shape [3, 45, 1]. then out_dims = {1, 45, 3}
-        // if (ins[i]->numel() != (*outs)[0]->numel()) {
         if (ins[i]->numel()) {
-          configs[i] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
-                                                     dims_simplifier.in_dims[i],
-                                                     dims_simplifier.rank);
+          in_dims[i] = dims_simplifier.in_dims[i];
         }
       }
     }
@@ -250,10 +273,10 @@ struct BroadcastDataInit {
 
 template <int Index, int VecSize>
 struct BroadcastDataSetter {
-  template <typename Array, typename ArgsT>
+  template <typename Array, typename ArgsT, typename IndexT>
   static __device__ __forceinline__ void Apply(const Array &ins,
                                                ArgsT *args,
-                                               uint64_t index_bc[][VecSize]) {
+                                               IndexT index_bc[][VecSize]) {
     using Type = std::tuple_element_t<Index, ArgsT>;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
@@ -317,7 +340,7 @@ __device__ void VectorizedBroadcastKernelImpl(
     Array<_ptr_ OutT *, NumOuts> outs,
     const Array<bool, Arity> &use_broadcast,
     const IndexT numel,
-    const Array<kps::details::BroadcastConfig, Arity> &configs,
+    const Array<BroadcastConfigType<IndexT>, Arity> &configs,
     uint32_t num,
     IndexT block_offset,
     int read_lens,
@@ -343,13 +366,13 @@ __device__ void VectorizedBroadcastKernelImpl(
                           read_lens);
 #else
   if (LoadType == kBroadcast) {
-    uint64_t index_bc[Arity][VecSize] = {0};
+    IndexT index_bc[Arity][VecSize] = {0};
     Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
     IndexT thread_offset =
         block_offset + static_cast<IndexT>(threadIdx.x) * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
-      uint64_t idx = thread_offset + k;
+      IndexT idx = thread_offset + k;
       if (IsBoundary && idx == numel) break;
 #pragma unroll
       for (int i = 0; i < DDim::kMaxRank; ++i) {
@@ -400,7 +423,7 @@ __global__ void VectorizedBroadcastKernel(
     Array<_ptr_ OutT *, NumOuts> outs,
     Array<bool, Arity> use_broadcast,
     IndexT numel,
-    Array<kps::details::BroadcastConfig, Arity> configs,
+    Array<BroadcastConfigType<IndexT>, Arity> configs,
     IndexT main_offset,
     IndexT tail_tid,
     int read_lens,
@@ -540,6 +563,19 @@ void LaunchBroadcastKernel(
     const IndexT main_offset = (numel / block_len) * block_len;
     const IndexT tail_tid = numel % block_len;
 
+    // Build the index-typed configs now that IndexT is known. An empty
+    // in_dims[i] means the input needs no broadcast config (see
+    // InitBroadcastConfigs).
+    Array<BroadcastConfigType<IndexT>, Arity> configs;
+    if (!classifier.all_elementwise) {
+      for (int i = 0; i < Arity; ++i) {
+        if (!classifier.in_dims[i].empty()) {
+          configs[i] = kps::details::BroadcastConfig<IndexT>(
+              classifier.out_dims, classifier.in_dims[i], classifier.rank);
+        }
+      }
+    }
+
     if (classifier.all_elementwise) {
       VectorizedBroadcastKernel<Functor,
                                 OutT,
@@ -552,7 +588,7 @@ void LaunchBroadcastKernel(
                                            classifier.outs_data,
                                            classifier.use_broadcast,
                                            numel,
-                                           classifier.configs,
+                                           configs,
                                            main_offset,
                                            tail_tid,
                                            VecSize,
@@ -570,7 +606,7 @@ void LaunchBroadcastKernel(
                                            classifier.outs_data,
                                            classifier.use_broadcast,
                                            numel,
-                                           classifier.configs,
+                                           configs,
                                            main_offset,
                                            tail_tid,
                                            VecSize,
@@ -587,7 +623,7 @@ void LaunchBroadcastKernel(
                                            classifier.outs_data,
                                            classifier.use_broadcast,
                                            numel,
-                                           classifier.configs,
+                                           configs,
                                            main_offset,
                                            tail_tid,
                                            VecSize,
@@ -599,7 +635,13 @@ void LaunchBroadcastKernel(
 #ifdef PADDLE_WITH_XPU_KP
   launch_with_index_t(uint32_t{0});
 #else
-  const uint64_t index_limit = std::numeric_limits<uint32_t>::max();
+  // Use the cheaper uint32 indexing (and 32-bit FastDivMod inside
+  // BroadcastConfig) on GPU only when the element count and the grid stride
+  // both fit. The bound is INT_MAX rather than UINT32_MAX because the 32-bit
+  // FastDivMod constructor computes `1 << shift_val` in signed 32-bit and would
+  // overflow for a divisor (an out dim) larger than INT_MAX. The grid_stride
+  // bound mirrors `stride` inside the kernel.
+  const uint64_t index_limit = std::numeric_limits<int32_t>::max();
   const uint64_t grid_stride =
       static_cast<uint64_t>(gpu_config.block_per_grid.x) *
       static_cast<uint64_t>(gpu_config.GetBlockSize()) * VecSize;
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
index 55bc0e3fc71f4..b257cfc74fa9e 100644
--- a/paddle/phi/kernels/primitive/datamover_primitives.h
+++ b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -20,6 +20,8 @@
 #ifdef PADDLE_WITH_HIP
 #include <hip/hip_fp16.h>
 #endif
+#include <type_traits>
+
 #include "paddle/common/ddim.h"
 #include "paddle/phi/kernels/funcs/fast_divmod.h"
 
@@ -39,9 +41,15 @@ struct alignas(sizeof(T) * VecSize) VectorType {
  * index of the output data. if input or output shape is [dim0, dim1] then dims
  * must be [dim1, dim0].
  */
+template <typename IndexT = int64_t>
 struct BroadcastConfig {
-  funcs::FastDivMod<int64_t> divmoders[DDim::kMaxRank];
-  uint64_t strides[DDim::kMaxRank];
+  // FastDivMod has a 32-bit primary template and a 64-bit int64_t
+  // specialization (there is no unsigned specialization), so pick the matching
+  // signed divmod type according to the index width.
+  using DivModT = std::conditional_t<(sizeof(IndexT) > 4), int64_t, int>;
+
+  funcs::FastDivMod<DivModT> divmoders[DDim::kMaxRank];
+  IndexT strides[DDim::kMaxRank];
   int rank{0};
 
   // BroadcastConfig should be defined on host used on device.
@@ -51,7 +59,7 @@ struct BroadcastConfig {
                   const std::vector<int64_t>& in_dims,
                   int dim_size) {
     for (int i = 0; i < dim_size; ++i) {
-      divmoders[i] = funcs::FastDivMod<int64_t>(out_dims[i]);
+      divmoders[i] = funcs::FastDivMod<DivModT>(out_dims[i]);
     }
 
     for (int i = 0; i < dim_size; ++i) {
@@ -59,8 +67,8 @@ struct BroadcastConfig {
       strides[i] = (i != 0 && strides[i] != 0)
                        ? std::accumulate(in_dims.begin(),
                                          in_dims.begin() + i,
-                                         int64_t{1},
-                                         std::multiplies<int64_t>())
+                                         IndexT{1},
+                                         std::multiplies<IndexT>())
                        : strides[i];
     }
     rank = dim_size;
@@ -413,12 +421,12 @@ __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
     IndexT block_offset,
-    const details::BroadcastConfig& config,
+    const details::BroadcastConfig<IndexT>& config,
     IndexT total_num_output,
     int stride_nx,
     int stride_ny) {
   IndexT thread_offset = block_offset + static_cast<IndexT>(threadIdx.x);
-  uint64_t index_src = 0;
+  IndexT index_src = 0;
 
 #pragma unroll
   for (int ny = 0; ny < NY; ++ny) {
@@ -751,11 +759,11 @@ __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
     IndexT block_offset,
-    const details::BroadcastConfig& config,
+    const details::BroadcastConfig<IndexT>& config,
     IndexT total_num_output,
     int read_lens = NX) {
   IndexT thread_offset = block_offset + static_cast<IndexT>(threadIdx.x) * NX;
-  uint64_t index_src = 0;
+  IndexT index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
@@ -813,11 +821,11 @@ __device__ __forceinline__ void ReadDataBc(
     ArgsT* dst,
     const T* __restrict__ src,
     IndexT block_offset,
-    const details::BroadcastConfig& config,
+    const details::BroadcastConfig<IndexT>& config,
     IndexT total_num_output,
     int read_lens = NX) {
   IndexT thread_offset = block_offset + static_cast<IndexT>(threadIdx.x) * NX;
-  uint64_t index_src = 0;
+  IndexT index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {

From f65527e494213f9b8d95e0072ec811c62a3c44b1 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Wed, 15 Jul 2026 10:57:18 +0000
Subject: [PATCH 12/12] fix bugs

---
 paddle/phi/kernels/funcs/broadcast_function.h | 45 ++++++++++---------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
index fb85952f919a7..794e689dc5f92 100644
--- a/paddle/phi/kernels/funcs/broadcast_function.h
+++ b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -429,11 +429,10 @@ __global__ void VectorizedBroadcastKernel(
     int read_lens,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  uint64_t block_offset =
-      static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
-  uint64_t stride = static_cast<uint64_t>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
-  for (; block_offset < static_cast<uint64_t>(main_offset);
-       block_offset += stride) {
+  IndexT block_offset =
+      static_cast<IndexT>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
+  IndexT stride = static_cast<IndexT>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
+  for (; block_offset < main_offset; block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   IndexT,
@@ -447,13 +446,16 @@ __global__ void VectorizedBroadcastKernel(
                                             numel,
                                             configs,
                                             BLOCK_NUM_X * read_lens,
-                                            static_cast<IndexT>(block_offset),
+                                            block_offset,
                                             read_lens,
                                             func);
   }
-  int64_t num =
-      static_cast<int64_t>(numel) - static_cast<int64_t>(block_offset);
-  if (num > 0) {
+  // `num` is the tail element count handled by a single block, always bounded
+  // by one block's span, so casting it to uint32 for the Impl is safe. Guard on
+  // `block_offset < numel` (not `num > 0`) so the subtraction never underflows
+  // when IndexT is unsigned.
+  if (block_offset < numel) {
+    IndexT num = numel - block_offset;
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   IndexT,
@@ -467,16 +469,14 @@ __global__ void VectorizedBroadcastKernel(
                                             numel,
                                             configs,
                                             static_cast<uint32_t>(num),
-                                            static_cast<IndexT>(block_offset),
+                                            block_offset,
                                             read_lens,
                                             func);
   }
 #else
-  uint64_t block_offset =
-      static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
-  uint64_t stride = static_cast<uint64_t>(BLOCK_NUM_X) * GRID_NUM_X * VecSize;
-  for (; block_offset < static_cast<uint64_t>(main_offset);
-       block_offset += stride) {
+  IndexT block_offset = static_cast<IndexT>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
+  IndexT stride = static_cast<IndexT>(BLOCK_NUM_X) * GRID_NUM_X * VecSize;
+  for (; block_offset < main_offset; block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   IndexT,
@@ -490,13 +490,15 @@ __global__ void VectorizedBroadcastKernel(
                                             numel,
                                             configs,
                                             BLOCK_NUM_X * VecSize,
-                                            static_cast<IndexT>(block_offset),
+                                            block_offset,
                                             read_lens,
                                             func);
   }
-  int64_t num =
-      static_cast<int64_t>(numel) - static_cast<int64_t>(block_offset);
-  if (num > 0) {
+  // `num` is one block's tail element count, so casting to uint32 for the Impl
+  // is safe. Guard on `block_offset < numel` (not `num > 0`) so the subtraction
+  // never underflows when IndexT is unsigned.
+  if (block_offset < numel) {
+    IndexT num = numel - block_offset;
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   IndexT,
@@ -510,7 +512,7 @@ __global__ void VectorizedBroadcastKernel(
                                             numel,
                                             configs,
                                             static_cast<uint32_t>(num),
-                                            static_cast<IndexT>(block_offset),
+                                            block_offset,
                                             read_lens,
                                             func);
   }
@@ -528,6 +530,9 @@ void LaunchBroadcastKernel(
 #endif
   auto launch_with_index_t = [&](auto index_tag) {
     using IndexT = decltype(index_tag);
+    // Although this runs on the host and int64 would be numerically fine, we
+    // keep IndexT here so the offsets/main_offset/tail_tid passed to the kernel
+    // match the IndexT the device kernel is instantiated with.
 #ifdef PADDLE_WITH_XPU_KP
     const IndexT numel = static_cast<IndexT>(classifier.numel);
     const int threads = 64;