Use DispatchGmmCombineDecode operator to replace MC2 when env variable VLLM_ASCEND_ENABLE_FUSED_MC2=1

Signed-off-by: wangqiankun <[email protected]>

Author: wangqiankun13

vllm_ascend/ascend_forward_context.py

@@ -24,6 +24,7 @@ class MoECommType(Enum):
     MC2 = 1
     ALLTOALL = 2
     FUSED_ALLTOALL = 3
+    FUSED_MC2 = 4
 
 
 @contextmanager

vllm_ascend/envs.py

@@ -132,6 +132,9 @@
     # Whether to anbale dynamic EPLB
     "DYNAMIC_EPLB":
     lambda: os.getenv("DYNAMIC_EPLB", "false").lower(),
+    # Whether to anbale fused mc2(dispatch_gmm_combine_decode operator)
+    "VLLM_ASCEND_ENABLE_FUSED_MC2":
+    lambda: bool(int(os.getenv("VLLM_ASCEND_ENABLE_FUSED_MC2", '0'))),
 }
 
 # end-env-vars-definition

vllm_ascend/ops/fused_moe/moe_comm_method.py

@@ -46,6 +46,8 @@ def setup_moe_comm_method(moe_config):
     _MoECommMethods[MoECommType.MC2] = MC2CommImpl(moe_config)
     _MoECommMethods[MoECommType.FUSED_ALLTOALL] = FusedAlltoAllCommImpl(
         moe_config)
+    _MoECommMethods[MoECommType.FUSED_MC2] = FusedMC2CommImpl(
+        moe_config)
 
 
 class MoECommMethod(ABC):
@@ -315,3 +317,71 @@ def fused_experts(
             out=out,
         )
         return out
+
+class FusedMC2CommImpl(MoECommMethod):
+    """This implementation is for the scenarios listed below:
+    1. `enable_expert_parallel=True`.
+    2. `npu_moe_distribute_dispatch` and `npu_moe_distribute_combine` are available.
+    3. `enable_expert_parallel=False` is not supported.
+    
+    This implementation uses the FusedMC2 communication method, which is optimized for
+    Communication and Computation parallelism on Ascend devices.
+    """
+
+    def _get_token_dispatcher(self):
+        return TokenDispatcherWithMC2()
+
+    def _get_prepare_finalize(self):
+        return PrepareAndFinalizeWithMC2(self.moe_config)
+
+    def fused_experts(
+            self,
+            hidden_states: torch.Tensor,
+            w1: torch.Tensor,
+            w2: torch.Tensor,
+            topk_weights: torch.Tensor,
+            topk_ids: torch.Tensor,
+            activation: str = "silu",
+            apply_router_weight_on_input: bool = False,
+            use_int8_w8a8: bool = False,
+            use_int4_w4a8: bool = False,
+            use_int4_w4a16: bool = False,
+            global_num_experts: Optional[int] = None,
+            expert_map: Optional[torch.Tensor] = None,
+            w1_scale: Optional[torch.Tensor] = None,
+            w2_scale: Optional[torch.Tensor] = None,
+            w1_scale_bias: torch.Tensor = None,
+            w2_scale_bias: torch.Tensor = None,
+            w1_offset: Optional[torch.Tensor] = None,
+            w2_offset: Optional[torch.Tensor] = None,
+            # For Cube/Vector parallel
+            shared_experts: Optional[Any] = None,
+            quantized_x_for_share: Optional[Any] = None,
+            dynamic_scale_for_share: Optional[Any] = None,
+            # For load balance
+            log2phy: torch.Tensor = None,
+            global_redundant_expert_num: int = 0,
+            need_trans: bool = False,
+            dynamic_eplb: bool = False,
+            mc2_mask: torch.Tensor = None,
+            pertoken_scale: Optional[torch.Tensor] = None):
+
+        output, _ = torch.ops._C_ascend.dispatch_gmm_combine_decode(
+            x=hidden_states,
+            expert_ids=topk_ids,
+            gmm1_permuted_weight=w1[0],
+            gmm1_permuted_weight_scale=w1_scale[0],
+            gmm2_weight=w2[0],
+            gmm2_weight_scale=w2_scale[0],
+            expert_smooth_scales=None,
+            expert_scales=topk_weights.to(torch.float32),
+            group_ep=self.token_dispatcher.moe_all_to_all_group_name,
+            ep_rank_size=self.token_dispatcher.ep_world_size,
+            ep_rank_id=self.token_dispatcher.ep_rank_id,
+            moe_expert_num=len(expert_map),
+            shared_expert_num=1,
+            shared_expert_rank_num=0,
+            quant_mode=0,
+            global_bs=self.token_dispatcher.global_bs
+        )
+        return output

vllm_ascend/ops/register_custom_ops.py

@@ -250,7 +250,7 @@ def _maybe_all_reduce_tensor_model_parallel_impl(
     forward_context = get_forward_context()
     moe_comm_type = forward_context.moe_comm_type
     if moe_comm_type in {
-            MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_ALLTOALL
+            MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_ALLTOALL, MoECommType.FUSED_MC2
     } or forward_context.sp_enabled:
         return final_hidden_states
     else:

vllm_ascend/quantization/w8a8_dynamic.py

@@ -235,6 +235,8 @@ def apply(
         topk_weights = topk_weights.to(self.in_dtype)
 
         moe_comm_method = get_forward_context().moe_comm_method
+        fused_mc2_flag = get_forward_context(
+        ).moe_comm_type == MoECommType.FUSED_MC2
         if self.dynamic_eplb:
             w1 = layer.w13_weight_list
             w1_scale = layer.w13_weight_scale_fp32_list
@@ -244,7 +246,7 @@ def apply(
             w1 = [layer.w13_weight]
             w1_scale = [layer.w13_weight_scale_fp32]
             w2 = [layer.w2_weight]
-            w2_scale = [layer.w2_weight_scale]
+            w2_scale = [layer.w2_weight_scale_fp32 if fused_mc2_flag else layer.w2_weight_scale]
 
         fused_flag = get_forward_context(
         ).moe_comm_type == MoECommType.FUSED_ALLTOALL
@@ -284,6 +286,8 @@ def process_weights_after_loading(self, layer):
             layer.w13_weight_offset.data.shape[0], -1)
         layer.w2_weight_scale.data = layer.w2_weight_scale.data.view(
             layer.w2_weight_scale.data.shape[0], -1)
+        layer.w2_weight_scale_fp32 = layer.w2_weight_scale.data.to(
+            torch.float32)
         layer.w2_weight_offset.data = layer.w2_weight_offset.data.view(
             layer.w2_weight_offset.data.shape[0], -1)
 

vllm_ascend/worker/model_runner_v1.py

@@ -1419,9 +1419,10 @@ def _select_moe_comm_method(self,
 
         elif soc_version in {AscendDeviceType._910_93}:
             moe_comm_type = (
-                MoECommType.MC2 if num_tokens <= mc2_tokens_capacity else
-                MoECommType.FUSED_ALLTOALL if quant_type == "w8a8_dynamic"
-                and get_ep_group().world_size <= 16 else MoECommType.ALLTOALL)
+                (MoECommType.FUSED_MC2 if envs_ascend.VLLM_ASCEND_ENABLE_FUSED_MC2 and quant_type == "w8a8_dynamic"
+                else MoECommType.MC2) if num_tokens <= mc2_tokens_capacity
+                else MoECommType.FUSED_ALLTOALL if quant_type == "w8a8_dynamic" and get_ep_group().world_size <= 16
+                else MoECommType.ALLTOALL)
         else:
             raise ValueError(f"Unsupported soc_version: {soc_version}")