Simplify template param of fused_rope kernel func (#61446)

MarioLulab · web-flow · commit f03095bdd037 · 2024-02-04T10:48:20.000+08:00
diff --git a/paddle/phi/kernels/fusion/gpu/fused_rope_grad_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_rope_grad_kernel.cu
@@ -110,73 +110,57 @@ void FusedRopeGradKernel(const Context& dev_ctx,
   }
 
   int sign = -1;
+
+  VectorizedFusedRopeCudaKernelFunc<T, MPType, vec_size> kernel_func =
+      use_neox_rotary_style
+          ? VectorizedFusedRopeWithRotateEveryTwoKernel<T, MPType, vec_size>
+          : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, vec_size>;
+
   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 {
-    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>;
+    kernel_func<<<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 {
     // 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);
+
+    kernel_func<<<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,
+                                            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);
+    phi::Array<const T*, 3> input_kv{ins_data[1], ins_data[2], nullptr};
+    phi::Array<T*, 3> out_kv{outs_data[1], outs_data[2], nullptr};
+    kernel_func<<<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);
   }
 }
 
diff --git a/paddle/phi/kernels/fusion/gpu/fused_rope_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_rope_kernel.cu
@@ -181,27 +181,25 @@ void FusedRopeKernel(const Context& dev_ctx,
   }
 
   int sign = 1;
+  VectorizedFusedRopeCudaKernelFunc<T, MPType, vec_size> kernel_func =
+      use_neox_rotary_style
+          ? VectorizedFusedRopeWithRotateEveryTwoKernel<T, MPType, vec_size>
+          : VectorizedFusedRopeWithRotateHalfKernel<T, MPType, vec_size>;
+
   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<<<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);
 
-    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 {
     // Multi Query Attention (MQA) or Group Query Attention (GQA)
     PADDLE_ENFORCE_EQ(
@@ -227,53 +225,35 @@ void FusedRopeKernel(const Context& dev_ctx,
               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);
+    kernel_func<<<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,
+                                            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);
+    phi::Array<const T*, 3> input_kv{ins_data[1], ins_data[2], nullptr};
+    phi::Array<T*, 3> out_kv{outs_data[1], outs_data[2], nullptr};
+    kernel_func<<<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
diff --git a/paddle/phi/kernels/fusion/gpu/fused_rope_utils.h b/paddle/phi/kernels/fusion/gpu/fused_rope_utils.h
@@ -19,9 +19,9 @@
 namespace phi {
 namespace fusion {
 
-template <typename T, typename MPType, int NInputs, int VecSize>
+template <typename T, typename MPType, int VecSize>
 using VectorizedFusedRopeCudaKernelFunc =
-    void (*)(phi::Array<const T*, NInputs> ins_data,
+    void (*)(phi::Array<const T*, 3> ins_data,
              phi::Array<const T*, 2> sin_cos_data,
              const int64_t* position_ids_data,
              bool flag_sin_cos,
@@ -30,7 +30,7 @@ using VectorizedFusedRopeCudaKernelFunc =
              int64_t seq_len,
              int64_t num_heads,
              int64_t head_dim,
-             phi::Array<T*, NInputs> outs_data,
+             phi::Array<T*, 3> outs_data,
              int num_inputs,
              MPType div_c);
 
@@ -86,9 +86,9 @@ __device__ void VectorizedGetSinCos(phi::Array<const T*, 2> sin_cos_data,
   }
 }
 
-template <typename T, typename MPType, int NInputs, int VecSize = 2>
+template <typename T, typename MPType, int VecSize = 2>
 __global__ void VectorizedFusedRopeWithRotateEveryTwoKernel(
-    phi::Array<const T*, NInputs> ins_data,
+    phi::Array<const T*, 3> ins_data,
     phi::Array<const T*, 2> sin_cos_data,
     const int64_t* position_ids_data,
     bool flag_sin_cos,
@@ -97,7 +97,7 @@ __global__ void VectorizedFusedRopeWithRotateEveryTwoKernel(
     int64_t seq_len,
     int64_t num_heads,
     int64_t head_dim,
-    phi::Array<T*, NInputs> outs_data,
+    phi::Array<T*, 3> outs_data,
     int num_inputs,
     MPType div_c) {
   int64_t index =
@@ -127,7 +127,7 @@ __global__ void VectorizedFusedRopeWithRotateEveryTwoKernel(
                         div_c);
 
 #pragma unroll
-    for (int iter = 0; iter < NInputs; iter++) {
+    for (int iter = 0; iter < 3; iter++) {
       if (iter >= num_inputs) break;
       const T* input = ins_data[iter] + index;
       VecType* out = reinterpret_cast<VecType*>(outs_data[iter] + index);
@@ -161,9 +161,9 @@ __global__ void VectorizedFusedRopeWithRotateEveryTwoKernel(
   }
 }
 
-template <typename T, typename MPType, int NInputs, int VecSize = 2>
+template <typename T, typename MPType, int VecSize = 2>
 __global__ void VectorizedFusedRopeWithRotateHalfKernel(
-    phi::Array<const T*, NInputs> ins_data,
+    phi::Array<const T*, 3> ins_data,
     phi::Array<const T*, 2> sin_cos_data,
     const int64_t* position_ids_data,
     bool flag_sin_cos,
@@ -172,7 +172,7 @@ __global__ void VectorizedFusedRopeWithRotateHalfKernel(
     int64_t seq_len,
     int64_t num_heads,
     int64_t head_dim,
-    phi::Array<T*, NInputs> outs_data,
+    phi::Array<T*, 3> outs_data,
     int num_inputs,
     MPType div_c) {
   int64_t index =
@@ -204,7 +204,7 @@ __global__ void VectorizedFusedRopeWithRotateHalfKernel(
     // use rotate_half mode
     int stride_r = head_dim / 2;
 #pragma unroll
-    for (int iter = 0; iter < NInputs; iter++) {
+    for (int iter = 0; iter < 3; iter++) {
       if (iter >= num_inputs) break;
       // get value_index and rotate_half_index
       int index_v = index;
