fix fused_bias_act big tensor bug (#73093)

wanghuancoder · web-flow · commit 517262633b2b · 2025-06-06T17:26:11.000+08:00
* fix fused_bias_act big tensor bug

* refine

* refine

* refine

* refine
diff --git a/paddle/phi/kernels/funcs/load_store_util.h b/paddle/phi/kernels/funcs/load_store_util.h
@@ -81,7 +81,7 @@ struct Store {
 
 template <typename T>
 struct Store<T, true> {
-  Store(T *dst, const T *shift, const T *smooth, const int cols)
+  Store(T *dst, const T *shift, const T *smooth, const int64_t cols)
       : dst_(dst), shift_(shift), smooth_(smooth), cols_(cols) {}
 
   template <int VecSize>
@@ -102,12 +102,14 @@ struct Store<T, true> {
   T *dst_;
   const T *shift_;
   const T *smooth_;
-  const int cols_;
+  const int64_t cols_;
 };
 
 template <typename T>
 struct DequantLoad {
-  DequantLoad(const int32_t *src, const float *dequant_scales, const int cols)
+  DequantLoad(const int32_t *src,
+              const float *dequant_scales,
+              const int64_t cols)
       : src_(src), dequant_scales_(dequant_scales), cols_(cols) {}
 
   template <int VecSize>
@@ -132,7 +134,7 @@ struct DequantLoad {
 
   const int32_t *src_;
   const float *dequant_scales_;
-  const int cols_;
+  const int64_t cols_;
 };
 
 template <typename T, typename OutT, bool Smooth = false>
@@ -186,7 +188,7 @@ struct QuantStore<T, int8_t, true> {
   QuantStore(int8_t *dst,
              const T *shift,
              const T *smooth,
-             const int cols,
+             const int64_t cols,
              const int quant_round_type,
              const float quant_scale,
              const float quant_max_bound,
@@ -232,7 +234,7 @@ struct QuantStore<T, int8_t, true> {
   const float quant_min_bound_;
   const T *shift_;
   const T *smooth_;
-  const int cols_;
+  const int64_t cols_;
 };
 
 }  // namespace funcs
diff --git a/paddle/phi/kernels/fusion/gpu/fused_bias_act_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_bias_act_kernel.cu
@@ -236,8 +236,8 @@ void DispatchComputeImpl(const Context &dev_ctx,
                          const DenseTensor &x,
                          const DenseTensor *bias,
                          const std::string &act_method,
-                         int rows,
-                         int cols,
+                         int64_t rows,
+                         int64_t cols,
                          const float quant_scale,
                          const int quant_round_type,
                          const float quant_max_bound,
@@ -398,6 +398,83 @@ void DispatchWithDtype(const Context &dev_ctx,
                        float quant_min_bound,
                        DenseTensor *out,
                        NormalVersion) {
+  const auto &x_dims = x.dims();
+  bool use_glu = (act_method == "geglu" || act_method == "swiglu");
+  if (bias.get_ptr() != nullptr) {
+    const auto &bias_dims = bias->dims();
+    PADDLE_ENFORCE_EQ(bias_dims.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "The bias must be a 1D tensor, but got %dD tensor.",
+                          bias_dims.size()));
+    PADDLE_ENFORCE_EQ(
+        bias_dims[0],
+        x_dims[x_dims.size() - 1],
+        common::errors::InvalidArgument(
+            "The bias length must be equal to the last dimension of input x. "
+            "Expected %d, but got %d.",
+            x_dims[x_dims.size() - 1],
+            bias_dims[0]));
+  }
+
+  if (dequant_scales.get_ptr() != nullptr) {
+    const auto &scales_dims = dequant_scales->dims();
+    PADDLE_ENFORCE_EQ(
+        scales_dims.size(),
+        1,
+        common::errors::InvalidArgument(
+            "The dequant_scales must be a 1D tensor, but got %dD tensor.",
+            scales_dims.size()));
+    PADDLE_ENFORCE_EQ(scales_dims[0],
+                      x_dims[x_dims.size() - 1],
+                      common::errors::InvalidArgument(
+                          "The dequant_scales length must be equal to the last "
+                          "dimension of input x. "
+                          "Expected %d, but got %d.",
+                          x_dims[x_dims.size() - 1],
+                          scales_dims[0]));
+  }
+  if (shift.get_ptr() != nullptr) {
+    const auto &shift_dims = shift->dims();
+    PADDLE_ENFORCE_EQ(shift_dims.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "The shift must be a 1D tensor, but got %dD tensor.",
+                          shift_dims.size()));
+    int64_t shift_dim =
+        use_glu ? std::div(static_cast<int64_t>(x_dims[x_dims.size() - 1]),
+                           static_cast<int64_t>(2))
+                      .quot
+                : x_dims[x_dims.size() - 1];
+    PADDLE_ENFORCE_EQ(
+        shift_dims[0],
+        shift_dim,
+        common::errors::InvalidArgument("The shift length invalid. "
+                                        "Expected %d, but got %d.",
+                                        shift_dim,
+                                        shift_dims[0]));
+  }
+  if (smooth.get_ptr() != nullptr) {
+    const auto &smooth_dims = smooth->dims();
+    PADDLE_ENFORCE_EQ(smooth_dims.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "The smooth must be a 1D tensor, but got %dD tensor.",
+                          smooth_dims.size()));
+    int64_t smooth_dim =
+        use_glu ? std::div(static_cast<int64_t>(x_dims[x_dims.size() - 1]),
+                           static_cast<int64_t>(2))
+                      .quot
+                : x_dims[x_dims.size() - 1];
+    PADDLE_ENFORCE_EQ(
+        smooth_dims[0],
+        smooth_dim,
+        common::errors::InvalidArgument("The smooth length invalid. "
+                                        "Expected %d, but got %d.",
+                                        smooth_dim,
+                                        smooth_dims[0]));
+  }
+
   auto *bias_p = bias.get_ptr();
   auto *dequant_scales_p = dequant_scales.get_ptr();
   auto *shift_p = shift.get_ptr();
@@ -499,7 +576,6 @@ void FusedBiasActKernel(const Context &dev_ctx,
           quant_min_bound,
           out,
           typename DispatchDtypeTrait<phi::dtype::bfloat16>::FuncVersion{});
-
     } else if (compute_dtype == "fp16") {
       DispatchWithDtype<phi::dtype::float16, Context>(
           dev_ctx,
