[Bugfix] fix bias optional for fused_moe op (#67783)

penPenf28 · web-flow · commit 0867262a92c3 · 2024-08-28T20:30:29.000-07:00
diff --git a/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h b/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h
@@ -1016,17 +1016,31 @@ void MoeGemmRunner<T, WeightType>::moe_gemm_bias_act(
     std::string activation_type,
     cudaStream_t stream) {
   if (activation_type == "none") {
-    run_gemm<EpilogueOpBias>(A,
-                             B,
-                             weight_scales,
-                             biases,
-                             C,
-                             total_rows_before_expert,
-                             total_rows,
-                             gemm_n,
-                             gemm_k,
-                             num_experts,
-                             stream);
+    if (biases) {
+      run_gemm<EpilogueOpBias>(A,
+                               B,
+                               weight_scales,
+                               biases,
+                               C,
+                               total_rows_before_expert,
+                               total_rows,
+                               gemm_n,
+                               gemm_k,
+                               num_experts,
+                               stream);
+    } else {
+      run_gemm<EpilogueOpNoBias>(A,
+                                 B,
+                                 weight_scales,
+                                 nullptr,
+                                 C,
+                                 total_rows_before_expert,
+                                 total_rows,
+                                 gemm_n,
+                                 gemm_k,
+                                 num_experts,
+                                 stream);
+    }
   }
 }
 
diff --git a/paddle/phi/kernels/fusion/cutlass/fused_moe_kernel.cu b/paddle/phi/kernels/fusion/cutlass/fused_moe_kernel.cu
@@ -79,11 +79,11 @@ void FusedMoeKernel(const Context& ctx,
   moe_compute.ComputeFFN(&X,
                          &gate_weight,
                          &ffn1_weight,
-                         ffn1_scale,
-                         ffn1_bias,
+                         ffn1_scale ? ffn1_scale.get_ptr() : nullptr,
+                         ffn1_bias ? ffn1_bias.get_ptr() : nullptr,
                          &ffn2_weight,
-                         ffn2_scale,
-                         ffn2_bias,
+                         ffn2_scale ? ffn2_scale.get_ptr() : nullptr,
+                         ffn2_bias ? ffn2_bias.get_ptr() : nullptr,
                          nullptr,
                          moe_topk,
                          norm_topk_prob,
diff --git a/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h b/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h
@@ -125,20 +125,20 @@ class MoeHelper {
   void ComputeFFN(const DenseTensor *X,
                   const DenseTensor *gate_weight,
                   const DenseTensor *ffn1_weight,
-                  const paddle::optional<DenseTensor> &ffn1_scale,
-                  const paddle::optional<DenseTensor> &ffn1_bias,
+                  const DenseTensor *ffn1_scale,
+                  const DenseTensor *ffn1_bias,
                   const DenseTensor *ffn2_weight,
-                  const paddle::optional<DenseTensor> &ffn2_scale,
-                  const paddle::optional<DenseTensor> &ffn2_bias,
+                  const DenseTensor *ffn2_scale,
+                  const DenseTensor *ffn2_bias,
                   const DenseTensor *moe_token_type_ids,
                   const int moe_topk,
                   const bool norm_topk_prob,
                   const std::string moe_type,
                   DenseTensor *output) {
     auto *input_activations = X->data<T>();
     auto *gating_weights = gate_weight->data<float>();
-    auto *fc1_expert_biases = ffn1_bias.get_ptr()->data<T>();
-    auto *fc2_expert_biases = ffn2_bias.get_ptr()->data<T>();
+    const T *fc1_expert_biases = ffn1_bias ? ffn1_bias->data<T>() : nullptr;
+    const T *fc2_expert_biases = ffn2_bias ? ffn2_bias->data<T>() : nullptr;
 
     auto *output_ = output->data<T>();
     auto stream = ctx.stream();
@@ -312,7 +312,7 @@ class MoeHelper {
       int8_moe_gemm_runner_->moe_gemm_bias_act(
           reinterpret_cast<NvType *>(permuted_data_),
           reinterpret_cast<const uint8_t *>(ffn1_weight->data<int8_t>()),
-          reinterpret_cast<const NvType *>(ffn1_scale.get_ptr()->data<T>()),
+          reinterpret_cast<const NvType *>(ffn1_scale->data<T>()),
           reinterpret_cast<const NvType *>(fc1_expert_biases),
           reinterpret_cast<NvType *>(fc1_out),
           total_rows_before_expert_,
@@ -327,7 +327,7 @@ class MoeHelper {
           reinterpret_cast<NvType *>(permuted_data_),
           reinterpret_cast<const cutlass::uint4b_t *>(
               ffn1_weight->data<int8_t>()),
-          reinterpret_cast<const NvType *>(ffn1_scale.get_ptr()->data<T>()),
+          reinterpret_cast<const NvType *>(ffn1_scale->data<T>()),
           reinterpret_cast<const NvType *>(fc1_expert_biases),
           reinterpret_cast<NvType *>(fc1_out),
           total_rows_before_expert_,
@@ -372,7 +372,7 @@ class MoeHelper {
         int8_moe_gemm_runner_->moe_gemm(
             reinterpret_cast<NvType *>(act_out),
             reinterpret_cast<const uint8_t *>(ffn2_weight->data<int8_t>()),
-            reinterpret_cast<const NvType *>(ffn2_scale.get_ptr()->data<T>()),
+            reinterpret_cast<const NvType *>(ffn2_scale->data<T>()),
             reinterpret_cast<NvType *>(fc2_result),
             total_rows_before_expert_,
             expanded_active_expert_rows,
@@ -385,7 +385,7 @@ class MoeHelper {
             reinterpret_cast<NvType *>(act_out),
             reinterpret_cast<const cutlass::uint4b_t *>(
                 ffn2_weight->data<int8_t>()),
-            reinterpret_cast<const NvType *>(ffn2_scale.get_ptr()->data<T>()),
+            reinterpret_cast<const NvType *>(ffn2_scale->data<T>()),
             reinterpret_cast<NvType *>(fc2_result),
             total_rows_before_expert_,
             expanded_active_expert_rows,
diff --git a/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_op.h b/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_op.h
@@ -745,13 +745,14 @@ __global__ void finalize_moe_routing_kernel(
           expanded_permuted_rows + expanded_permuted_row * cols;
 
       const int expert_idx = expert_for_source_row[k_offset];
-      const T* bias_ptr = bias + expert_idx * cols;
+      const T* bias_ptr = bias ? bias + expert_idx * cols : nullptr;
+      const T bias_value = bias_ptr ? bias_ptr[tid] : T{0.f};
 
       thread_output =
           static_cast<float>(thread_output) +
           row_scale * static_cast<float>(
                           expanded_permuted_rows_row_ptr[tid] +
-                          bias_ptr[tid] *
+                          bias_value *
                               static_cast<T>(static_cast<float>(compute_bias)));
     }
     thread_output = static_cast<float>(thread_output) /
