From 4de9e57f83dccd2bfc52ba3dc8a517375bc2c07a Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Fri, 10 Jul 2026 08:07:22 +0000
Subject: [PATCH 1/2] Fix baddbmm_grad input_grad overflow for large tensors

---
 .../kernels/impl/baddbmm_grad_kernel_impl.h    | 18 +-----------------
 1 file changed, 1 insertion(+), 17 deletions(-)

diff --git a/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h b/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h
index 67bc4e35c8876d..3ad4664a41e1ba 100644
--- a/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h
+++ b/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h
@@ -180,26 +180,10 @@ void BaddbmmGradKernel(const Context& dev_ctx,
                                          .reshape(eigen_dinput_shape)
                                          .template cast<T>();
       }
-    } else {
-      // The VCOPY does not support the float16, bfloat16
-      if (!is_float16_or_bfloat16) {
-        mt_blas.VCOPY(
-            total_elems, out_grad.data<MPType>(), input_grad->data<MPType>());
-      } else {
-        funcs::ForRange<Context> for_range(dev_ctx, total_elems);
-        BCopyOrScaleFunctor<T> functor(
-            1, out_grad.data<T>(), input_grad->data<T>(), total_elems);
-        for_range(functor);
-      }
-    }
-
-    // The SCAL does not support the float16, bfloat16
-    if (!is_float16_or_bfloat16) {
-      mt_blas.SCAL(total_elems, beta, input_grad->data<MPType>());
     } else {
       funcs::ForRange<Context> for_range(dev_ctx, total_elems);
       BCopyOrScaleFunctor<T> functor(
-          beta, input_grad->data<T>(), input_grad->data<T>(), total_elems);
+          beta, out_grad.data<T>(), input_grad->data<T>(), total_elems);
       for_range(functor);
     }
   }

From 757e4feefb14d1c07dd5d71837995f20f180e266 Mon Sep 17 00:00:00 2001
From: feixi139 <18210811610@163.com>
Date: Fri, 10 Jul 2026 09:01:15 +0000
Subject: [PATCH 2/2] fix bugs

---
 paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h b/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h
index 3ad4664a41e1ba..f9cf417adaa46f 100644
--- a/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h
+++ b/paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h
@@ -183,9 +183,14 @@ void BaddbmmGradKernel(const Context& dev_ctx,
     } else {
       funcs::ForRange<Context> for_range(dev_ctx, total_elems);
       BCopyOrScaleFunctor<T> functor(
-          beta, out_grad.data<T>(), input_grad->data<T>(), total_elems);
+          1, out_grad.data<T>(), input_grad->data<T>(), total_elems);
       for_range(functor);
     }
+
+    funcs::ForRange<Context> for_range(dev_ctx, total_elems);
+    BCopyOrScaleFunctor<T> functor(
+        beta, input_grad->data<T>(), input_grad->data<T>(), total_elems);
+    for_range(functor);
   }
   if (x_grad) {
     dev_ctx.template Alloc<T>(x_grad);