[PHI]Add non-bool mask cast (#72944)

Eddie-Wang1120 · web-flow · commit 41aa80a48a10 · 2025-05-27T18:59:38.000+08:00
* add non-bool mask cast

* fix space

* fix grad kernel
diff --git a/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu
@@ -41,6 +41,11 @@ __global__ void GPUMaskedFillGradKernel(const T* out_grad,
                                         const int64_t batch_size,
                                         T* x_grad) {
   int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x);
+
+  if (idx >= (input_len / VecSize)) {
+    return;
+  }
+
   int64_t vec_idx = idx * VecSize;
   int64_t mask_idx = vec_idx / batch_size;
   using VecType = kps::details::VectorType<T, VecSize>;
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -6417,6 +6417,9 @@ def masked_fill(
     if np.isscalar(value):
         value = paddle.full([], value, x.dtype)
 
+    if mask.dtype != "bool":
+        mask = paddle.cast(mask, "bool")
+
     if in_dynamic_or_pir_mode():
         out = _C_ops.masked_fill(x, mask, value)
         return out
@@ -6471,6 +6474,9 @@ def masked_fill_(
     if np.isscalar(value):
         value = paddle.full([], value, x.dtype)
 
+    if mask.dtype != "bool":
+        mask = paddle.cast(mask, "bool")
+
     x = _C_ops.masked_fill_(x, mask, value)
     return x
 
