xpu rotary support use_neox_rotary_style == True (#64090)

HarperCy · web-flow · commit c1f6e81b5317 · 2024-05-11T11:25:25.000+08:00
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_grad_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_rope_grad_kernel.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
-
 #include "paddle/phi/kernels/fusion/xpu/fused_rope_utils.h"
 
 namespace phi {
@@ -37,7 +36,6 @@ void FusedRopeGradKernel(const Context& dev_ctx,
   if (dout_q.numel() <= 0) {
     return;
   }
-
   int64_t batch_size = dout_q.dims()[0];
   int64_t seq_len = dout_q.dims()[1];
   int64_t num_heads = dout_q.dims()[2];
@@ -46,33 +44,68 @@ void FusedRopeGradKernel(const Context& dev_ctx,
                     0,
                     phi::errors::InvalidArgument(
                         "The head_dim of input must be a multiple of 2."));
-
   xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
-  int64_t sin_cos_len = batch_size * seq_len * head_dim;
-  auto* sin_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
-  auto* cos_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
-
-  if (sin.get_ptr() && cos.get_ptr()) {
-    PADDLE_ENFORCE_EQ(sin.get_ptr()->dims(),
-                      cos.get_ptr()->dims(),
-                      phi::errors::InvalidArgument(
-                          "The dims of sin and cos must be the same. But "
-                          "received sin's dims is {%s}, cos's dims is {%s}.",
-                          sin.get_ptr()->dims(),
-                          cos.get_ptr()->dims()));
-  }
-
-  XPUGetSinCosData<XPUType, Context>(
-      dev_ctx, sin, position_ids, sin_data, batch_size, seq_len, head_dim);
-  XPUGetSinCosData<XPUType, Context>(
-      dev_ctx, cos, position_ids, cos_data, batch_size, seq_len, head_dim);
 
+  int64_t sin_cos_len = batch_size * seq_len * head_dim;
   if (use_neox_rotary_style) {
-    // TODO(lijin23): support rotary_embedding every_two.
-    PADDLE_THROW(
-        phi::errors::Unimplemented("XPU do not support rotary_embedding_grad "
-                                   "with use_neox_rotary_style set."));
+    auto* sin_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    auto* cos_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    if (sin.get_ptr() && cos.get_ptr()) {
+      PADDLE_ENFORCE_EQ(sin.get_ptr()->dims(),
+                        cos.get_ptr()->dims(),
+                        phi::errors::InvalidArgument(
+                            "The dims of sin and cos must be the same. But "
+                            "received sin's dims is {%s}, cos's dims is {%s}.",
+                            sin.get_ptr()->dims(),
+                            cos.get_ptr()->dims()));
+    }
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, sin, position_ids, sin_data, batch_size, seq_len, head_dim);
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, cos, position_ids, cos_data, batch_size, seq_len, head_dim);
+    if (!dout_k.get_ptr()) {
+      PADDLE_THROW(phi::errors::Unimplemented(
+          "XPU do not support rotary_embedding_grad without qk when "
+          "use_neox_rotary_style == True."));
+    }
+    auto* dq_data = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(dq));
+    auto* dk_data = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(dk));
+    int ret = xpu::rotary_embedding_v3_grad<XPUType, XPUType>(
+        dev_ctx.x_context(),
+        reinterpret_cast<const XPUType*>(dout_q.data<T>()),
+        reinterpret_cast<const XPUType*>(dout_k->data<T>()),
+        cos_data,
+        sin_data,
+        dq_data,
+        dk_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 * head_dim, num_heads * head_dim, head_dim, 1});
+    PADDLE_ENFORCE_XDNN_SUCCESS(ret, "rotary_embedding_v3_grad");
+    if (dout_v.get_ptr()) {
+      PADDLE_THROW(phi::errors::Unimplemented(
+          "XPU do not support rotary_embedding_grad with v when "
+          "use_neox_rotary_style == True."));
+    }
   } else {
+    auto* sin_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    auto* cos_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    if (sin.get_ptr() && cos.get_ptr()) {
+      PADDLE_ENFORCE_EQ(sin.get_ptr()->dims(),
+                        cos.get_ptr()->dims(),
+                        phi::errors::InvalidArgument(
+                            "The dims of sin and cos must be the same. But "
+                            "received sin's dims is {%s}, cos's dims is {%s}.",
+                            sin.get_ptr()->dims(),
+                            cos.get_ptr()->dims()));
+    }
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, sin, position_ids, sin_data, batch_size, seq_len, head_dim);
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, cos, position_ids, cos_data, batch_size, seq_len, head_dim);
     if (head_dim * sizeof(T) <= 1024 && head_dim % 64 == 0 && dout_k) {
       int64_t num_heads_k = dout_k->dims()[2];
       auto* dq_data = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(dq));
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_rope_kernel.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
-
 #include "paddle/phi/kernels/fusion/xpu/fused_rope_utils.h"
 
 namespace phi {
@@ -38,7 +37,6 @@ void FusedRopeKernel(const Context& dev_ctx,
   if (q.numel() <= 0) {
     return;
   }
-
   PADDLE_ENFORCE_EQ(
       time_major,
       false,
@@ -51,33 +49,70 @@ void FusedRopeKernel(const Context& dev_ctx,
                     0,
                     phi::errors::InvalidArgument(
                         "The head_dim of input must be a multiple of 2."));
-
   xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
 
   int64_t sin_cos_len = batch_size * seq_len * head_dim;
-  auto* sin_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
-  auto* cos_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
-
-  if (sin.get_ptr() && cos.get_ptr()) {
-    PADDLE_ENFORCE_EQ(sin.get_ptr()->dims(),
-                      cos.get_ptr()->dims(),
-                      phi::errors::InvalidArgument(
-                          "The dims of sin and cos must be the same. But "
-                          "received sin's dims is {%s}, cos's dims is {%s}.",
-                          sin.get_ptr()->dims(),
-                          cos.get_ptr()->dims()));
-  }
-
-  XPUGetSinCosData<XPUType, Context>(
-      dev_ctx, sin, position_ids, sin_data, batch_size, seq_len, head_dim);
-  XPUGetSinCosData<XPUType, Context>(
-      dev_ctx, cos, position_ids, cos_data, batch_size, seq_len, head_dim);
-
   if (use_neox_rotary_style) {
-    // TODO(lijin23): support rotary_embedding every_two.
-    PADDLE_THROW(phi::errors::Unimplemented(
-        "XPU do not support rotary_embedding with use_neox_rotary_style set."));
+    auto* sin_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    auto* cos_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    if (sin.get_ptr() && cos.get_ptr()) {
+      PADDLE_ENFORCE_EQ(sin.get_ptr()->dims(),
+                        cos.get_ptr()->dims(),
+                        phi::errors::InvalidArgument(
+                            "The dims of sin and cos must be the same. But "
+                            "received sin's dims is {%s}, cos's dims is {%s}.",
+                            sin.get_ptr()->dims(),
+                            cos.get_ptr()->dims()));
+    }
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, sin, position_ids, sin_data, batch_size, seq_len, head_dim);
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, cos, position_ids, cos_data, batch_size, seq_len, head_dim);
+    if (!k) {
+      PADDLE_THROW(phi::errors::Unimplemented(
+          "XPU do not support rotary_embedding without qk when "
+          "use_neox_rotary_style == True."));
+    }
+    auto* outq_data =
+        reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(out_q));
+    auto* outk_data =
+        reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(out_k));
+    int ret = xpu::rotary_embedding_v3<XPUType, XPUType>(
+        dev_ctx.x_context(),
+        reinterpret_cast<const XPUType*>(q.data<T>()),
+        reinterpret_cast<const XPUType*>(k->data<T>()),
+        cos_data,
+        sin_data,
+        outq_data,
+        outk_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 * head_dim, num_heads * head_dim, head_dim, 1});
+    PADDLE_ENFORCE_XDNN_SUCCESS(ret, "rotary_embedding_v3");
+    if (v) {
+      PADDLE_THROW(
+          phi::errors::Unimplemented("XPU do not support rotary_embedding with "
+                                     "v when use_neox_rotary_style == True."));
+    }
   } else {
+    auto* sin_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    auto* cos_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(sin_cos_len);
+    if (sin.get_ptr() && cos.get_ptr()) {
+      PADDLE_ENFORCE_EQ(sin.get_ptr()->dims(),
+                        cos.get_ptr()->dims(),
+                        phi::errors::InvalidArgument(
+                            "The dims of sin and cos must be the same. But "
+                            "received sin's dims is {%s}, cos's dims is {%s}.",
+                            sin.get_ptr()->dims(),
+                            cos.get_ptr()->dims()));
+    }
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, sin, position_ids, sin_data, batch_size, seq_len, head_dim);
+    XPUGetSinCosData<XPUType, Context>(
+        dev_ctx, cos, position_ids, cos_data, batch_size, seq_len, head_dim);
     if (head_dim * sizeof(T) <= 1024 && head_dim % 64 == 0 && k) {
       int64_t num_heads_k = k->dims()[2];
       auto* outq_data =
diff --git a/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h b/paddle/phi/kernels/fusion/xpu/fused_rope_utils.h
@@ -36,11 +36,10 @@ void XPUGetSinCosData(const Context& dev_ctx,
                                      dims_size));
     if (dims_size == 4) {
       // sin.shape: [1, seq_len, 1, head_dim]
-      PADDLE_ENFORCE_EQ(
-          (sin_cos_dims[0] == 1 && sin_cos_dims[2] == 1),
-          true,
-          phi::errors::InvalidArgument(
-              "The batch_size and num_heads of sin and cos must be 1."));
+      PADDLE_ENFORCE_EQ((sin_cos_dims[2] == 1),
+                        true,
+                        phi::errors::InvalidArgument(
+                            "The num_heads of sin and cos must be 1."));
     }
     int sin_seq_len_dim = (dims_size) == 4 ? 1 : 0;
     if (position_ids.get_ptr()) {
@@ -92,11 +91,12 @@ void XPUGetSinCosData(const Context& dev_ctx,
         PADDLE_ENFORCE_XDNN_SUCCESS(ret, "gather");
       }
     } else {
+      int sin_cos_batch_size = (dims_size) == 4 ? sin_cos_dims[0] : 1;
       int ret = xpu::broadcast<XPUType>(
           dev_ctx.x_context(),
           reinterpret_cast<const XPUType*>(sin_cos->data()),
           sin_cos_data,
-          {1, seq_len, head_dim},
+          {sin_cos_batch_size, seq_len, head_dim},
           {batch_size, seq_len, head_dim});
       PADDLE_ENFORCE_XDNN_SUCCESS(ret, "broadcast");
     }
