PaddlePaddle
diff --git a/‎cmake/external/xpu.cmake
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎cmake/external/xpu.cmake
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
diff --git a/‎paddle/phi/backends/xpu/xpu3_op_list.cc
Copy file name to clipboard
+7-2Lines changed: 7 additions & 2 deletions b/‎paddle/phi/backends/xpu/xpu3_op_list.cc
Copy file name to clipboard
+7-2Lines changed: 7 additions & 2 deletions
diff --git a/‎paddle/phi/kernels/fusion/xpu/fused_rope_utils.h
Copy file name to clipboard
+73-69Lines changed: 73 additions & 69 deletions b/‎paddle/phi/kernels/fusion/xpu/fused_rope_utils.h
Copy file name to clipboard
+73-69Lines changed: 73 additions & 69 deletions
diff --git a/‎paddle/phi/kernels/xpu/activation_grad_kernel.cc
Copy file name to clipboard
+42-34Lines changed: 42 additions & 34 deletions b/‎paddle/phi/kernels/xpu/activation_grad_kernel.cc
Copy file name to clipboard
+42-34Lines changed: 42 additions & 34 deletions
@@ -30,16 +30,16 @@ 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/20241203")
+  set(XPU_XHPC_BASE_DATE "release/3.1.0")
 endif()
-set(XPU_XCCL_BASE_VERSION "3.0.1.1") # For XRE5
+set(XPU_XCCL_BASE_VERSION "3.0.1.6") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
   set(XPU_XFT_BASE_VERSION "20230602")
 endif()
 
 if(NOT DEFINED XPU_XRE_BASE_VERSION)
   if(WITH_XPU_XRE5)
-    set(XPU_XRE_BASE_VERSION "5.0.21.10")
+    set(XPU_XRE_BASE_VERSION "5.0.21.14")
   else()
     set(XPU_XRE_BASE_VERSION "4.32.0.1")
   endif()
 
@@ -428,7 +428,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32,
                      phi::DataType::BOOL})},
-      {"exp_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"exp_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"exp",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -1302,7 +1305,9 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
       {"tanh_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"tan", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"tanh", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"temporal_shift", XPUKernelSet({phi::DataType::FLOAT32})},
 
@@ -274,67 +274,55 @@ void XPUFusedRotaryEveryTwo(const Context& dev_ctx,
                             DenseTensor* out_q,
                             DenseTensor* out_k,
                             DenseTensor* out_v) {
-  auto single_func = &xpu::rotary_embedding_v3_single<XPUType, XPUSCType>;
-  auto fusion_func = &xpu::rotary_embedding_v3<XPUType, XPUSCType>;
-  const char* single_func_name = "rotary_embedding_v3_single";
-  const char* fusion_func_name = "rotary_embedding_v3";
+  auto everytwo_func = &xpu::rotary_embedding_everytwo<XPUType, XPUSCType>;
+  const char* single_func_name = "rotary_embedding_everytwo_unary_freqs";
+  const char* fusion_func_name = "rotary_embedding_everytwo_binary_freqs";
   if (is_bwd) {
-    single_func = &xpu::rotary_embedding_v3_single_grad<XPUType, XPUSCType>;
-    fusion_func = &xpu::rotary_embedding_v3_grad<XPUType, XPUSCType>;
-    single_func_name = "rotary_embedding_v3_single_grad";
-    fusion_func_name = "rotary_embedding_v3_grad";
+    everytwo_func = &xpu::rotary_embedding_everytwo_grad<XPUType, XPUSCType>;
+    single_func_name = "rotary_embedding_everytwo_unary_freqs_grad";
+    fusion_func_name = "rotary_embedding_everytwo_binary_freqs_grad";
   }
   if (!in_k) {
-    int ret = single_func(
-        dev_ctx.x_context(),
-        reinterpret_cast<const XPUType*>(in_q.data()),
-        cos_data,
-        sin_data,
-        reinterpret_cast<XPUType*>(out_q->data()),
-        batch_size,
-        seq_len,
-        num_heads,
-        head_dim,
-        {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
-        "BLHD",
-        true);
+    int ret = everytwo_func(dev_ctx.x_context(),
+                            reinterpret_cast<const XPUType*>(in_q.data()),
+                            nullptr,
+                            cos_data,
+                            sin_data,
+                            reinterpret_cast<XPUType*>(out_q->data()),
+                            nullptr,
+                            {batch_size, seq_len, num_heads, head_dim},
+                            {batch_size, seq_len, 1, head_dim},
+                            -1,
+                            10000.0f);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
   } else {
     int64_t num_heads_k = in_k->dims()[2];
-    int ret = fusion_func(
-        dev_ctx.x_context(),
-        reinterpret_cast<const XPUType*>(in_q.data()),
-        reinterpret_cast<const XPUType*>(in_k->data()),
-        cos_data,
-        sin_data,
-        reinterpret_cast<XPUType*>(out_q->data()),
-        reinterpret_cast<XPUType*>(out_k->data()),
-        batch_size,
-        seq_len,
-        num_heads,
-        head_dim,
-        {seq_len * num_heads * head_dim, num_heads * head_dim, head_dim, 1},
-        {seq_len * num_heads_k * head_dim, num_heads_k * head_dim, head_dim, 1},
-        num_heads_k,
-        "BLHD",
-        true);
+    int ret = everytwo_func(dev_ctx.x_context(),
+                            reinterpret_cast<const XPUType*>(in_q.data()),
+                            reinterpret_cast<const XPUType*>(in_k->data()),
+                            cos_data,
+                            sin_data,
+                            reinterpret_cast<XPUType*>(out_q->data()),
+                            reinterpret_cast<XPUType*>(out_k->data()),
+                            {batch_size, seq_len, num_heads, head_dim},
+                            {batch_size, seq_len, 1, head_dim},
+                            num_heads_k,
+                            10000.0f);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, fusion_func_name);
   }
   if (in_v) {
     int64_t num_heads_v = in_v->dims()[2];
-    int ret = single_func(
-        dev_ctx.x_context(),
-        reinterpret_cast<const XPUType*>(in_v->data()),
-        cos_data,
-        sin_data,
-        reinterpret_cast<XPUType*>(out_v->data()),
-        batch_size,
-        seq_len,
-        num_heads_v,
-        head_dim,
-        {seq_len * num_heads_v * head_dim, num_heads_v * head_dim, head_dim, 1},
-        "BLHD",
-        true);
+    int ret = everytwo_func(dev_ctx.x_context(),
+                            reinterpret_cast<const XPUType*>(in_v->data()),
+                            nullptr,
+                            cos_data,
+                            sin_data,
+                            reinterpret_cast<XPUType*>(out_v->data()),
+                            nullptr,
+                            {batch_size, seq_len, num_heads_v, head_dim},
+                            {batch_size, seq_len, 1, head_dim},
+                            -1,
+                            10000.0f);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
   }
 }
@@ -355,41 +343,57 @@ void XPUFusedRotaryHalf(const Context& dev_ctx,
                         DenseTensor* out_q,
                         DenseTensor* out_k,
                         DenseTensor* out_v) {
-  auto single_func =
-      &xpu::rotary_embedding_half_unary_freqs<XPUType, XPUSCType>;
-  auto fusion_func =
-      &xpu::rotary_embedding_half_binary_freqs<XPUType, XPUSCType>;
+  auto half_func = &xpu::rotary_embedding_half<XPUType, XPUSCType>;
   const char* single_func_name = "rotary_embedding_half_unary_freqs";
   const char* fusion_func_name = "xpu::rotary_embedding_half_binary_freqs";
   if (is_bwd) {
-    single_func =
-        &xpu::rotary_embedding_half_unary_freqs_grad<XPUType, XPUSCType>;
-    fusion_func =
-        &xpu::rotary_embedding_half_binary_freqs_grad<XPUType, XPUSCType>;
+    half_func = &xpu::rotary_embedding_half_grad<XPUType, XPUSCType>;
     single_func_name = "rotary_embedding_half_unary_freqs_grad";
     fusion_func_name = "xpu::rotary_embedding_half_binary_freqs_grad";
   }
 
-  int64_t num_heads_k = in_k->dims()[2];
-  int ret = fusion_func(dev_ctx.x_context(),
+  if (!in_k) {
+    int ret = half_func(dev_ctx.x_context(),
+                        reinterpret_cast<const XPUType*>(in_q.data()),
+                        nullptr,
+                        reinterpret_cast<const XPUSCType*>(sin_data),
+                        reinterpret_cast<const XPUSCType*>(cos_data),
+                        reinterpret_cast<XPUType*>(out_q->data()),
+                        nullptr,
+                        {batch_size, seq_len, num_heads, head_dim},
+                        {batch_size, seq_len, 1, head_dim},
+                        -1,
+                        10000.0f);
+    PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
+  } else {
+    int64_t num_heads_k = in_k->dims()[2];
+    int ret = half_func(dev_ctx.x_context(),
                         reinterpret_cast<const XPUType*>(in_q.data()),
                         reinterpret_cast<const XPUType*>(in_k->data()),
                         reinterpret_cast<const XPUSCType*>(sin_data),
                         reinterpret_cast<const XPUSCType*>(cos_data),
                         reinterpret_cast<XPUType*>(out_q->data()),
                         reinterpret_cast<XPUType*>(out_k->data()),
                         {batch_size, seq_len, num_heads, head_dim},
-                        num_heads_k);
-  PADDLE_ENFORCE_XDNN_SUCCESS(ret, fusion_func_name);
+                        {batch_size, seq_len, 1, head_dim},
+                        num_heads_k,
+                        10000.0f);
+    PADDLE_ENFORCE_XDNN_SUCCESS(ret, fusion_func_name);
+  }
 
   if (in_v) {
     int64_t num_heads_v = in_v->dims()[2];
-    int ret = single_func(dev_ctx.x_context(),
-                          reinterpret_cast<const XPUType*>(in_v->data()),
-                          reinterpret_cast<const XPUSCType*>(sin_data),
-                          reinterpret_cast<const XPUSCType*>(cos_data),
-                          reinterpret_cast<XPUType*>(out_v->data()),
-                          {batch_size, seq_len, num_heads_v, head_dim});
+    int ret = half_func(dev_ctx.x_context(),
+                        reinterpret_cast<const XPUType*>(in_v->data()),
+                        nullptr,
+                        reinterpret_cast<const XPUSCType*>(sin_data),
+                        reinterpret_cast<const XPUSCType*>(cos_data),
+                        reinterpret_cast<XPUType*>(out_v->data()),
+                        nullptr,
+                        {batch_size, seq_len, num_heads_v, head_dim},
+                        {batch_size, seq_len, 1, head_dim},
+                        -1,
+                        10000.0f);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, single_func_name);
   }
 }
 
@@ -160,21 +160,6 @@ int xpu_activation_backward(const Context& dev_ctx,
   return r;
 }
 
-template <typename T>
-struct XPUExpGradFunctor : public funcs::BaseActivationFunctor<T> {
-  using XPUType = typename XPUTypeTrait<T>::Type;
-  template <typename Context>
-  void operator()(const Context& dev_ctx,
-                  const DenseTensor* x,
-                  const DenseTensor* out,
-                  const DenseTensor* dout,
-                  DenseTensor* dx) const {
-    int r = xpu_activation_backward<Context, T, XPUType>(
-        dev_ctx, x, out, dout, dx, xpu::exp_grad<XPUType>);
-    PADDLE_ENFORCE_XDNN_SUCCESS(r, "exp_grad");
-  }
-};
-
 template <typename T>
 struct XPULogGradFunctor : public funcs::BaseActivationFunctor<T> {
   template <typename Context>
@@ -418,21 +403,6 @@ struct XPUSigmoidGradFunctor : public funcs::BaseActivationFunctor<T> {
   }
 };
 
-template <typename T>
-struct XPUTanhGradFunctor : public funcs::BaseActivationFunctor<T> {
-  using XPUType = typename XPUTypeTrait<T>::Type;
-  template <typename Context>
-  void operator()(const Context& dev_ctx,
-                  const DenseTensor* x,
-                  const DenseTensor* out,
-                  const DenseTensor* dout,
-                  DenseTensor* dx) const {
-    int r = xpu_activation_backward<Context, T, XPUType>(
-        dev_ctx, x, out, dout, dx, xpu::tanh_grad<XPUType>);
-    PADDLE_ENFORCE_XDNN_SUCCESS(r, "tanh_grad");
-  }
-};
-
 template <typename T>
 struct XPUSquareGradFunctor : public funcs::BaseActivationFunctor<T> {
   using XPUType = typename XPUTypeTrait<T>::Type;
@@ -642,11 +612,9 @@ struct XPURsqrtGradFunctor : public funcs::BaseActivationFunctor<T> {
   }
 };
 
-DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Exp, XPUExpGradFunctor);
 DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Reciprocal, XPUReciprocalGradFunctor);
 DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid, XPUSigmoidGradFunctor);
 DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Sqrt, XPUSqrtGradFunctor);
-DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Tanh, XPUTanhGradFunctor);
 DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu, XPUReluGradFunctor);
 DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu6, XPURelu6GradFunctor);
 DEFINE_XPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Rsqrt, XPURsqrtGradFunctor);
@@ -700,6 +668,38 @@ void SiluGradKernel(const Context& dev_ctx,
   ActivationGradXPUImpl<T, Context, XPUSiluGradFunctor<T>>(
       dev_ctx, &x, &out, &dout, dx, functor);
 }
+
+template <typename T, typename Context>
+void ExpGradKernel(const Context& dev_ctx,
+                   const DenseTensor& out,
+                   const DenseTensor& dout,
+                   DenseTensor* dx) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  dev_ctx.template Alloc<T>(dx);
+  const XPUType* y_data = reinterpret_cast<const XPUType*>(out.data<T>());
+  const XPUType* y_grad = reinterpret_cast<const XPUType*>(dout.data<T>());
+  XPUType* x_grad = reinterpret_cast<XPUType*>(dx->data<T>());
+
+  int r =
+      xpu::exp_grad(dev_ctx.x_context(), y_data, y_grad, x_grad, dx->numel());
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "exp_grad");
+}
+
+template <typename T, typename Context>
+void TanhGradKernel(const Context& dev_ctx,
+                    const DenseTensor& out,
+                    const DenseTensor& dout,
+                    DenseTensor* dx) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  dev_ctx.template Alloc<T>(dx);
+  const XPUType* y_data = reinterpret_cast<const XPUType*>(out.data<T>());
+  const XPUType* y_grad = reinterpret_cast<const XPUType*>(dout.data<T>());
+  XPUType* x_grad = reinterpret_cast<XPUType*>(dx->data<T>());
+
+  int r =
+      xpu::tanh_grad(dev_ctx.x_context(), y_data, y_grad, x_grad, dx->numel());
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "tanh_grad");
+}
 }  // namespace phi
 
 PD_REGISTER_KERNEL(relu_grad,
@@ -725,7 +725,16 @@ PD_REGISTER_KERNEL(tanh_grad,
                    ALL_LAYOUT,
                    phi::TanhGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
+PD_REGISTER_KERNEL(exp_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ExpGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(square_grad,
                    XPU,
@@ -772,7 +781,6 @@ PD_REGISTER_KERNEL(rsqrt_grad,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
-PD_REGISTER_ACTIVATION_GRAD_KERNEL(exp_grad, ExpGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(log_grad, LogGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(leaky_relu_grad, LeakyReluGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(hardsigmoid_grad, HardSigmoidGradKernel)