PaddlePaddle
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_rope_grad_kernel.cu
Copy file name to clipboard
+89-36Lines changed: 89 additions & 36 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_rope_grad_kernel.cu
Copy file name to clipboard
+89-36Lines changed: 89 additions & 36 deletions
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_rope_kernel.cu
Copy file name to clipboard
+116-36Lines changed: 116 additions & 36 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_rope_kernel.cu
Copy file name to clipboard
+116-36Lines changed: 116 additions & 36 deletions
@@ -38,9 +38,11 @@ void FusedRopeGradKernel(const Context& dev_ctx,
   int64_t numel = dout_q.numel();
   if (numel <= 0) return;
   dev_ctx.template Alloc<T>(dq);
+
+  phi::Array<int64_t, 3> inputs_num_heads;
   // small size for broadcast
   auto batch_size = dout_q.dims()[0];
-  auto num_heads = dout_q.dims()[2];
+  inputs_num_heads[0] = dout_q.dims()[2];
   auto head_dim = dout_q.dims()[3];
   auto seq_len = dout_q.dims()[1];
   PADDLE_ENFORCE_NE(head_dim % 2,
@@ -64,19 +66,21 @@ void FusedRopeGradKernel(const Context& dev_ctx,
 
   ins_data[0] = dout_q.data<T>();
   outs_data[0] = dq->data<T>();
-  int num_inputs = 0;
+  int num_inputs = 1;
 
-  if (dout_k.get_ptr()) {
+  if (dout_k) {
     dev_ctx.template Alloc<T>(dk);
-    outs_data[1] = dk->data<T>();
-    ins_data[1] = dout_k->data<T>();
+    outs_data[num_inputs] = dk->data<T>();
+    ins_data[num_inputs] = dout_k->data<T>();
+    inputs_num_heads[num_inputs] = dk->dims()[2];
     num_inputs++;
   }
 
-  if (dout_v.get_ptr()) {
+  if (dout_v) {
     dev_ctx.template Alloc<T>(dv);
-    outs_data[2] = dv->data<T>();
-    ins_data[2] = dout_v->data<T>();
+    outs_data[num_inputs] = dv->data<T>();
+    ins_data[num_inputs] = dout_v->data<T>();
+    inputs_num_heads[num_inputs] = dv->dims()[2];
     num_inputs++;
   }
 
@@ -90,40 +94,89 @@ void FusedRopeGradKernel(const Context& dev_ctx,
 
     flag_sin_cos = true;
 
-    if (position_ids.get_ptr()) {
+    if (position_ids) {
       position_ids_data = position_ids->data<int64_t>();
     }
   }
 
+  bool is_same_num_heads = true;
+  auto prev_num_heads = inputs_num_heads[0];
+  for (int i = 1; i < num_inputs; ++i) {
+    if (prev_num_heads != inputs_num_heads[i]) {
+      is_same_num_heads = false;
+      break;
+    }
+    prev_num_heads = inputs_num_heads[i];
+  }
+
   int sign = -1;
-  if (use_neox_rotary_style) {
-    VectorizedFusedRopeWithRotateEveryTwoKernel<T, MPType, vec_size>
-        <<<grid, block, 0, stream>>>(ins_data,
-                                     sin_cos_data,
-                                     position_ids_data,
-                                     flag_sin_cos,
-                                     sign,
-                                     batch_size,
-                                     seq_len,
-                                     num_heads,
-                                     head_dim,
-                                     outs_data,
-                                     num_inputs,
-                                     div_c);
+  if (is_same_num_heads) {
+    VectorizedFusedRopeCudaKernelFunc<T, MPType, 3, vec_size> kernel_func_qkv =
+        use_neox_rotary_style
+            ? VectorizedFusedRopeWithRotateEveryTwoKernel<T,
+                                                          MPType,
+                                                          3,
+                                                          vec_size>
+            : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, 3, vec_size>;
+    kernel_func_qkv<<<grid, block, 0, stream>>>(ins_data,
+                                                sin_cos_data,
+                                                position_ids_data,
+                                                flag_sin_cos,
+                                                sign,
+                                                batch_size,
+                                                seq_len,
+                                                inputs_num_heads[0],
+                                                head_dim,
+                                                outs_data,
+                                                num_inputs,
+                                                div_c);
   } else {
-    VectorizedFusedRopeWithRotateHalfKernel<T, MPType, vec_size>
-        <<<grid, block, 0, stream>>>(ins_data,
-                                     sin_cos_data,
-                                     position_ids_data,
-                                     flag_sin_cos,
-                                     sign,
-                                     batch_size,
-                                     seq_len,
-                                     num_heads,
-                                     head_dim,
-                                     outs_data,
-                                     num_inputs,
-                                     div_c);
+    VectorizedFusedRopeCudaKernelFunc<T, MPType, 1, vec_size> kernel_func_q =
+        use_neox_rotary_style
+            ? VectorizedFusedRopeWithRotateEveryTwoKernel<T,
+                                                          MPType,
+                                                          1,
+                                                          vec_size>
+            : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, 1, vec_size>;
+    VectorizedFusedRopeCudaKernelFunc<T, MPType, 2, vec_size> kernel_func_kv =
+        use_neox_rotary_style
+            ? VectorizedFusedRopeWithRotateEveryTwoKernel<T,
+                                                          MPType,
+                                                          2,
+                                                          vec_size>
+            : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, 2, vec_size>;
+
+    // rotary position embedding Q
+    phi::Array<const T*, 1> input_q{ins_data[0]};
+    phi::Array<T*, 1> out_q{outs_data[0]};
+    kernel_func_q<<<grid, block, 0, stream>>>(input_q,
+                                              sin_cos_data,
+                                              position_ids_data,
+                                              flag_sin_cos,
+                                              sign,
+                                              batch_size,
+                                              seq_len,
+                                              inputs_num_heads[0],
+                                              head_dim,
+                                              out_q,
+                                              1,
+                                              div_c);
+
+    // rotary position embedding K,V
+    phi::Array<const T*, 2> input_kv{ins_data[1], ins_data[2]};
+    phi::Array<T*, 2> out_kv{outs_data[1], outs_data[2]};
+    kernel_func_kv<<<grid, block, 0, stream>>>(input_kv,
+                                               sin_cos_data,
+                                               position_ids_data,
+                                               flag_sin_cos,
+                                               sign,
+                                               batch_size,
+                                               seq_len,
+                                               inputs_num_heads[1],
+                                               head_dim,
+                                               out_kv,
+                                               num_inputs - 1,
+                                               div_c);
   }
 }
 
 
@@ -39,11 +39,14 @@ void FusedRopeKernel(const Context& dev_ctx,
   if (numel <= 0) return;
   dev_ctx.template Alloc<T>(out_q);
 
+  phi::Array<int64_t, 3> inputs_num_heads;
+
   // q.shape: [batch_size, seq_len, num_heads, head_dim]
   auto batch_size = q.dims()[0];
   auto seq_len = q.dims()[1];
-  auto num_heads = q.dims()[2];
+  inputs_num_heads[0] = q.dims()[2];
   auto head_dim = q.dims()[3];
+
   PADDLE_ENFORCE_EQ(head_dim % 2,
                     0,
                     phi::errors::InvalidArgument(
@@ -65,19 +68,21 @@ void FusedRopeKernel(const Context& dev_ctx,
 
   ins_data[0] = q.data<T>();
   outs_data[0] = out_q->data<T>();
-  int num_inputs = 0;
+  int num_inputs = 1;
 
-  if (k.get_ptr()) {
+  if (k) {
     dev_ctx.template Alloc<T>(out_k);
-    ins_data[1] = k->data<T>();
-    outs_data[1] = out_k->data<T>();
+    ins_data[num_inputs] = k->data<T>();
+    outs_data[num_inputs] = out_k->data<T>();
+    inputs_num_heads[num_inputs] = k->dims()[2];
     num_inputs++;
   }
 
-  if (v.get_ptr()) {
+  if (v) {
     dev_ctx.template Alloc<T>(out_v);
-    ins_data[2] = v->data<T>();
-    outs_data[2] = out_v->data<T>();
+    ins_data[num_inputs] = v->data<T>();
+    outs_data[num_inputs] = out_v->data<T>();
+    inputs_num_heads[num_inputs] = v->dims()[2];
     num_inputs++;
   }
 
@@ -113,7 +118,7 @@ void FusedRopeKernel(const Context& dev_ctx,
     }
     int sin_seq_len_dim = (dims_size) == 4 ? 1 : 0;
 
-    if (position_ids.get_ptr()) {
+    if (position_ids) {
       PADDLE_ENFORCE_EQ(
           (sin_dims[dims_size - 1] == head_dim &&
            sin_dims[sin_seq_len_dim] >= seq_len),
@@ -165,35 +170,110 @@ void FusedRopeKernel(const Context& dev_ctx,
     flag_sin_cos = true;
   }
 
+  bool is_same_num_heads = true;
+  auto prev_num_heads = inputs_num_heads[0];
+  for (int i = 1; i < num_inputs; ++i) {
+    if (prev_num_heads != inputs_num_heads[i]) {
+      is_same_num_heads = false;
+      break;
+    }
+    prev_num_heads = inputs_num_heads[i];
+  }
+
   int sign = 1;
-  if (use_neox_rotary_style) {
-    VectorizedFusedRopeWithRotateEveryTwoKernel<T, MPType, vec_size>
-        <<<grid, block, 0, stream>>>(ins_data,
-                                     sin_cos_data,
-                                     position_ids_data,
-                                     flag_sin_cos,
-                                     sign,
-                                     batch_size,
-                                     seq_len,
-                                     num_heads,
-                                     head_dim,
-                                     outs_data,
-                                     num_inputs,
-                                     div_c);
+  if (is_same_num_heads) {
+    VectorizedFusedRopeCudaKernelFunc<T, MPType, 3, vec_size> kernel_func_qkv =
+        use_neox_rotary_style
+            ? VectorizedFusedRopeWithRotateEveryTwoKernel<T,
+                                                          MPType,
+                                                          3,
+                                                          vec_size>
+            : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, 3, vec_size>;
+
+    kernel_func_qkv<<<grid, block, 0, stream>>>(ins_data,
+                                                sin_cos_data,
+                                                position_ids_data,
+                                                flag_sin_cos,
+                                                sign,
+                                                batch_size,
+                                                seq_len,
+                                                inputs_num_heads[0],
+                                                head_dim,
+                                                outs_data,
+                                                num_inputs,
+                                                div_c);
   } else {
-    VectorizedFusedRopeWithRotateHalfKernel<T, MPType, vec_size>
-        <<<grid, block, 0, stream>>>(ins_data,
-                                     sin_cos_data,
-                                     position_ids_data,
-                                     flag_sin_cos,
-                                     sign,
-                                     batch_size,
-                                     seq_len,
-                                     num_heads,
-                                     head_dim,
-                                     outs_data,
-                                     num_inputs,
-                                     div_c);
+    // Multi Query Attention (MQA) or Group Query Attention (GQA)
+    PADDLE_ENFORCE_EQ(
+        (inputs_num_heads[0] != inputs_num_heads[num_inputs - 1]) &&
+            (inputs_num_heads[0] % inputs_num_heads[num_inputs - 1] == 0),
+        true,
+        phi::errors::InvalidArgument(
+            "The MQA or GQA mode is entered, when the number of heads of qkv "
+            "is not exactly the same two by two. This mode requires "
+            "num_heads of q to be divisible by k,v."
+            "But recieved num_heads of q is %d, num_heads of k,v is %d",
+            inputs_num_heads[0],
+            inputs_num_heads[num_inputs - 1]));
+
+    if (k.get_ptr() && v.get_ptr()) {
+      PADDLE_ENFORCE_EQ(
+          inputs_num_heads[1] == inputs_num_heads[2],
+          true,
+          phi::errors::InvalidArgument(
+              "The num_heads of k must be equal to the num_heads of v when v "
+              "is not none."
+              "But recieved num_heads of k is %d, num_heads of v is %d",
+              inputs_num_heads[1],
+              inputs_num_heads[2]));
+    }
+
+    VectorizedFusedRopeCudaKernelFunc<T, MPType, 1, vec_size> kernel_func_q =
+        use_neox_rotary_style
+            ? VectorizedFusedRopeWithRotateEveryTwoKernel<T,
+                                                          MPType,
+                                                          1,
+                                                          vec_size>
+            : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, 1, vec_size>;
+    VectorizedFusedRopeCudaKernelFunc<T, MPType, 2, vec_size> kernel_func_kv =
+        use_neox_rotary_style
+            ? VectorizedFusedRopeWithRotateEveryTwoKernel<T,
+                                                          MPType,
+                                                          2,
+                                                          vec_size>
+            : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, 2, vec_size>;
+
+    // rotary position embedding Q
+    phi::Array<const T*, 1> input_q{ins_data[0]};
+    phi::Array<T*, 1> out_q{outs_data[0]};
+    kernel_func_q<<<grid, block, 0, stream>>>(input_q,
+                                              sin_cos_data,
+                                              position_ids_data,
+                                              flag_sin_cos,
+                                              sign,
+                                              batch_size,
+                                              seq_len,
+                                              inputs_num_heads[0],
+                                              head_dim,
+                                              out_q,
+                                              1,
+                                              div_c);
+
+    // rotary position embedding K,V
+    phi::Array<const T*, 2> input_kv{ins_data[1], ins_data[2]};
+    phi::Array<T*, 2> out_kv{outs_data[1], outs_data[2]};
+    kernel_func_kv<<<grid, block, 0, stream>>>(input_kv,
+                                               sin_cos_data,
+                                               position_ids_data,
+                                               flag_sin_cos,
+                                               sign,
+                                               batch_size,
+                                               seq_len,
+                                               inputs_num_heads[1],
+                                               head_dim,
+                                               out_kv,
+                                               num_inputs - 1,
+                                               div_c);
   }
 }
 }  // namespace fusion