[PHI] Fix masked_fill & masked_fill_grad kernel for big tensor (#73284)

lshpku · web-flow · commit e7b01654849c · 2025-06-13T16:47:46.000+08:00
diff --git a/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu
@@ -40,7 +40,7 @@ __global__ void GPUMaskedFillGradKernel(const T* out_grad,
                                         const int64_t input_len,
                                         const int64_t batch_size,
                                         T* x_grad) {
-  int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x);
+  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
 
   if (idx >= (input_len / VecSize)) {
     return;
diff --git a/paddle/phi/kernels/gpu/masked_fill_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_kernel.cu
@@ -37,7 +37,7 @@ __global__ void GPUMaskedFillKernel(const T* input,
                                     const int64_t input_len,
                                     const int64_t batch_size,
                                     T* output) {
-  int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x);
+  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
 
   if (idx >= (input_len / VecSize)) {
     return;
