[XPU] fix integral type of scale and update flash_attn_grad (#71776)

RuohengMa · web-flow · commit 86994e3f49b8 · 2025-03-24T10:20:47.000+08:00
* [XPU] fix integral type of scale

* modify xre date
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "dev/20250306")
+  set(XPU_XHPC_BASE_DATE "dev/20250318")
 endif()
 set(XPU_XCCL_BASE_VERSION "3.0.2.5") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
@@ -39,7 +39,7 @@ endif()
 
 if(NOT DEFINED XPU_XRE_BASE_VERSION)
   if(WITH_XPU_XRE5)
-    set(XPU_XRE_BASE_VERSION "5.0.21.15")
+    set(XPU_XRE_BASE_VERSION "5.0.21.18")
   else()
     set(XPU_XRE_BASE_VERSION "4.32.0.1")
   endif()
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -1237,7 +1237,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
-                     phi::DataType::INT32})},
+                     phi::DataType::INT32,
+                     phi::DataType::INT16,
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8})},
       {"scatter",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
diff --git a/paddle/phi/kernels/xpu/amp_kernel.cc b/paddle/phi/kernels/xpu/amp_kernel.cc
@@ -282,7 +282,7 @@ void CheckFiniteAndUnscaleKernel(const Context& dev_ctx,
                        x->numel(),
                        false,
                        inverse_scale,
-                       0.0);
+                       0.0f);
         PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
 
         r = xpu::cast(dev_ctx.x_context(),
@@ -297,7 +297,7 @@ void CheckFiniteAndUnscaleKernel(const Context& dev_ctx,
                            x->numel(),
                            false,
                            inverse_scale,
-                           0.0);
+                           0.0f);
         PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
       }
     }
diff --git a/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc
@@ -128,7 +128,13 @@ void FlashAttnGradKernelBase(const Context& ctx,
   // qkv_layout = AttnQKVLayout_t::ATTN_BLHD, const float* alibi_slopes =
   // nullptr, const std::vector<int64_t>& alibi_slopes_shape = {}, int
   // window_size_left = -1, int window_size_right = -1, int64_t v_head_dim =
-  // -1);
+  // -1, const int* downstart_row_indices_data = nullptr,
+  // const int* downend_row_indices_data = nullptr,
+  // const int* upstart_row_indices_data = nullptr,
+  // const int* upend_row_indices_data = nullptr,
+  // const int flash_mask_head_num = 0,
+  // int* flashmask_maxmin = nullptr,
+  // XPUStream side_stream = nullptr);
   int r = flash_attention_grad_kernel(
       ctx.x_context(),
       dout_data,                                  // dout
@@ -168,7 +174,14 @@ void FlashAttnGradKernelBase(const Context& ctx,
       {},                                         // alibi_slopes_shape
       -1,                                         // window_size_left
       -1,                                         // window_size_right
-      head_size_v                                 // v_head_dim
+      head_size_v,                                // v_head_dim
+      nullptr,                                    // downstart_row_indices_data
+      nullptr,                                    // downend_row_indices_data
+      nullptr,                                    // upstart_row_indices_data
+      nullptr,                                    // upend_row_indices_data
+      0,                                          // flash_mask_head_num
+      nullptr,                                    // flashmask_maxmin
+      nullptr                                     // side_stream
   );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
 }
diff --git a/paddle/phi/kernels/xpu/pool_grad_kernel.cc b/paddle/phi/kernels/xpu/pool_grad_kernel.cc
@@ -284,7 +284,7 @@ void Pool3dGradKernel(const Context& ctx,
                        n * c,
                        true,
                        scale,
-                       0.0);
+                       0.0f);
         PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
 
         r = xpu::broadcast(ctx.x_context(),
diff --git a/paddle/phi/kernels/xpu/scale_kernel.cc b/paddle/phi/kernels/xpu/scale_kernel.cc
@@ -15,6 +15,7 @@
 #include "paddle/phi/kernels/scale_kernel.h"
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/core/kernel_registry.h"
 
 namespace phi {
@@ -38,14 +39,16 @@ void ScaleKernel(const Context& dev_ctx,
   if (x.numel() == 0 || !x.IsInitialized()) {
     return;
   }
+
   using XPUType = typename XPUTypeTrait<T>::Type;
-  int r = xpu::scale(dev_ctx.x_context(),
-                     reinterpret_cast<const XPUType*>(x.data<T>()),
-                     reinterpret_cast<XPUType*>(out->data<T>()),
-                     x.numel(),
-                     bias_after_scale,
-                     scale.to<float>(),
-                     bias.to<float>());
+  using MT = typename phi::dtype::MPTypeTrait<T>::Type;
+  int r = xpu::scale<XPUType, MT>(dev_ctx.x_context(),
+                                  reinterpret_cast<const XPUType*>(x.data<T>()),
+                                  reinterpret_cast<XPUType*>(out->data<T>()),
+                                  x.numel(),
+                                  bias_after_scale,
+                                  scale.to<MT>(),
+                                  bias.to<MT>());
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
 }
 
@@ -58,5 +61,8 @@ PD_REGISTER_KERNEL(scale,
                    float,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
                    int,
                    int64_t) {}
diff --git a/paddle/phi/kernels/xpu/squared_l2_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/squared_l2_norm_grad_kernel.cc
@@ -52,7 +52,7 @@ void SquaredL2NormGradKernel(const Context& dev_ctx,
                      reinterpret_cast<XPUType*>(dx->data<T>()),
                      x.numel(),
                      false,
-                     dout_value_cpu * 2,
+                     static_cast<float>(dout_value_cpu * 2),
                      0.0f);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
 }

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@ void SquaredL2NormGradKernel(const Context& dev_ctx,`
`52`	`52`	`reinterpret_cast<XPUType*>(dx->data<T>()),`
`53`	`53`	`x.numel(),`
`54`	`54`	`false,`
`55`		`- dout_value_cpu * 2,`
	`55`	`+ static_cast<float>(dout_value_cpu * 2),`
`56`	`56`	`0.0f);`
`57`	`57`	`PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");`
`58`	`58`	`}`