fix

1. use unum instead of string
2. avoid setting forward_context in AscendFusedMoE.forward()

Author: Pr0Wh1teGivee

vllm_ascend/ascend_forward_context.py

@@ -22,6 +22,16 @@ class FusedMoEState(Enum):
     All2AllSeq = 5
 
 
+class MoECommImpl(Enum):
+    ALLGATHER = "AllGather"
+    MC2 = "MC2"
+    ALLTOALL = "AlltoAll"
+    NAIVE_MULTICAST = "NaiveMulticast"
+
+    def __str__(self):
+        return self.value + "CommImpl"
+
+
 # TODO(zzzzwwjj): add soc_version to choose branch
 def _get_fused_moe_state(ep_size: int, with_prefill: bool,
                          is_deepseek_v3_r1: bool):
@@ -52,7 +62,7 @@ def set_ascend_forward_context(
         with_prefill: bool = True,
         in_profile_run: bool = False,
         reserved_mc2_mask: Optional[torch.Tensor] = None,
-        moe_comm_method: str = "",
+        moe_comm_method_type: Optional[MoECommImpl] = None,
         num_actual_tokens: Optional[int] = None,
         aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
         batch_descriptor: Optional[BatchDescriptor] = None,
@@ -72,7 +82,12 @@ def set_ascend_forward_context(
             batch_descriptor=batch_descriptor,
     ):
         forward_context = get_forward_context()
-        forward_context.moe_comm_method_name = moe_comm_method + "commimpl"
+
+        from vllm_ascend.ops.moe.moe_comm_method import get_moe_comm_method
+        forward_context.moe_comm_method_type = moe_comm_method_type
+        forward_context.moe_comm_method = get_moe_comm_method(
+            moe_comm_method_type)
+
         forward_context.with_prefill = with_prefill
         tp_world_size = get_tensor_model_parallel_world_size()
         ep_size = (get_ep_group().world_size if

vllm_ascend/ops/common_fused_moe.py

@@ -29,14 +29,13 @@
     FusedMoE, UnquantizedFusedMoEMethod, determine_expert_map)
 
 from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.ascend_forward_context import MoECommImpl
 from vllm_ascend.distributed.parallel_state import get_mc2_group
 from vllm_ascend.eplb.core.eplb_utils import (determine_default_expert_map,
                                               determine_default_log2phy_map)
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
 from vllm_ascend.ops.moe.experts_selector import select_experts
-from vllm_ascend.ops.moe.moe_comm_method import (AllGatherCommImpl,
-                                                 AlltoAllCommImpl, MC2CommImpl,
-                                                 NaiveMulticastCommImpl)
+from vllm_ascend.ops.moe.moe_comm_method import setup_moe_comm_method
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ,
                                get_all_reduce_merge_state,
                                get_rm_router_logits_state, is_310p)
@@ -145,6 +144,8 @@ def apply(self,
 
 
 class AscendFusedMoE(FusedMoE):
+    # The moe_counter parameter is required during the initialization of EPLB
+    # to identify the current layer index within the MOE model.
     moe_counter = -1
 
     def __init__(self, *args, **kwargs):
@@ -172,14 +173,11 @@ def __init__(self, *args, **kwargs):
 
         assert self.quant_method is not None
 
-        AscendFusedMoE.moe_counter += 1
-        self.moe_instance_id = AscendFusedMoE.moe_counter
         self.moe_config.tp_group = get_tp_group()
         self.moe_config.dp_group = get_dp_group()
         self.moe_config.ep_group = get_ep_group()
         self.moe_config.mc2_group = get_mc2_group()
         ascend_config = get_ascend_config()
-        self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp
         self.dynamic_eplb = ascend_config.dynamic_eplb
         self.expert_map_path = ascend_config.expert_map_path
         self.global_redundant_expert_num = ascend_config.init_redundancy_expert
@@ -215,13 +213,9 @@ def __init__(self, *args, **kwargs):
         if self.dynamic_eplb:
             self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64)
 
-        for method in {
-                AllGatherCommImpl, AlltoAllCommImpl, MC2CommImpl,
-                NaiveMulticastCommImpl
-        }:
-            setattr(
-                self, method.__name__.lower(),
-                method(moe_config=self.moe_config))  # type: ignore[abstract]
+        self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp
+
+        setup_moe_comm_method(self.moe_config)
 
     def update_expert_map(self, new_expert_map):
         self.expert_map = new_expert_map
@@ -245,8 +239,8 @@ def maybe_all_reduce_tensor_model_parallel(
         outputs since each rank only has partial outputs.
         """
         forward_context = get_forward_context()
-        moe_comm_method_name = forward_context.moe_comm_method_name
-        if moe_comm_method_name in {"alltoallcommimpl", "mc2commimpl"}:
+        moe_comm_method_type = forward_context.moe_comm_method_type
+        if moe_comm_method_type in {MoECommImpl.AllTOAll, MoECommImpl.MC2}:
             return final_hidden_states
         else:
             return tensor_model_parallel_all_reduce(final_hidden_states)
@@ -260,9 +254,6 @@ def forward_impl(self, hidden_states: torch.Tensor,
 
         forward_context = get_forward_context()
         enable_force_load_balance = forward_context.in_profile_run
-        moe_comm_method_name = forward_context.moe_comm_method_name
-
-        forward_context.moe_comm_method = getattr(self, moe_comm_method_name)
 
         hidden_states, router_logits = forward_context.moe_comm_method.prepare(
             hidden_states=hidden_states,
@@ -287,11 +278,13 @@ def forward_impl(self, hidden_states: torch.Tensor,
             e_score_correction_bias=self.e_score_correction_bias,
             activation=self.activation,
             apply_router_weight_on_input=self.apply_router_weight_on_input,
-            enable_eplb=self.enable_eplb,
-            expert_load_view=self.expert_load_view,
-            logical_to_physical_map=self.logical_to_physical_map,
-            logical_replica_count=self.logical_replica_count,
-        )
+            quantized_x_for_share=quantized_x_for_share,
+            dynamic_scale_for_share=dynamic_scale_for_share,
+            shared_experts=None,
+            enable_force_load_balance=enable_force_load_balance,
+            log2phy=self.log2phy,
+            global_redundant_expert_num=self.global_redundant_expert_num)
+
         if isinstance(final_hidden_states, tuple):
             final_hidden_states, group_list_type, expert_tokens = final_hidden_states
 
@@ -410,8 +403,8 @@ def forward(
 
         # NOTE: This is exactly the opposite of `maybe_all_reduce_tensor_model_parallel`
         forward_context = get_forward_context()
-        moe_comm_method_name = forward_context.moe_comm_method_name
-        if moe_comm_method_name in {"alltoallcommimpl", "mc2commimpl"}:
+        moe_comm_method_type = forward_context.moe_comm_method_type
+        if moe_comm_method_type in {MoECommImpl.AllTOAll, MoECommImpl.MC2}:
             shared_out = tensor_model_parallel_all_reduce(shared_out)
 
         fused_out = super().forward(

vllm_ascend/ops/moe/moe_comm_method.py

@@ -13,15 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
+from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import torch
 from vllm.config import get_current_vllm_config
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.fused_moe import FusedMoEConfig
 
+from vllm_ascend.ascend_forward_context import MoECommImpl
 from vllm_ascend.ops.moe.fused_moe_prepare_and_finalize import (
     FusedMoEPrepareAndFinalizeWithAll2All,
     FusedMoEPrepareAndFinalizeWithAllGather, FusedMoEPrepareAndFinalizeWithMC2,
@@ -32,6 +34,23 @@
                                                   TokenDispatcherWithMC2,
                                                   TokenDispatcherWithMoge)
 
+_MoECommMethods: Dict[str, MoECommMethod] = {}
+
+
+def _register_moe_comm_method(moe_comm_method: MoECommMethod):
+    _MoECommMethods[moe_comm_method.__class__.__name__] = moe_comm_method
+
+
+def get_moe_comm_method(name: MoECommImpl) -> Optional[MoECommMethod]:
+    return _MoECommMethods.get(str(name))
+
+
+def setup_moe_comm_method(moe_config):
+    _register_moe_comm_method(AlltoAllCommImpl(moe_config))
+    _register_moe_comm_method(AllGatherCommImpl(moe_config))
+    _register_moe_comm_method(MC2CommImpl(moe_config))
+    _register_moe_comm_method(NaiveMulticastCommImpl(moe_config))
+
 
 class MoECommMethod(ABC):
     """Base class for MoE communication methods."""

vllm_ascend/ops/moe/moe_mlp.py

@@ -21,6 +21,7 @@
 from torch.nn.functional import pad
 from vllm.forward_context import get_forward_context
 
+from vllm_ascend.ascend_forward_context import MoECommImpl
 from vllm_ascend.utils import dispose_tensor, is_310p
 
 
@@ -76,7 +77,7 @@ def quant_apply_mlp(hidden_states: torch.Tensor,
     bias1, bias2 = None, None
     _output_dtype = w2_scale.dtype
 
-    is_mc2 = get_forward_context().moe_comm_method_name == "mc2commimpl"
+    is_mc2 = get_forward_context().moe_comm_method_type == MoECommImpl.MC2
     if w1_scale_bias is None and is_mc2:
         if w1_scale.dtype != torch.float32:
             w1_scale = w1_scale.to(torch.float32)

vllm_ascend/torchair/models/torchair_pangu_moe.py

@@ -498,9 +498,6 @@ def forward(
         global _ROUTER_SCALE
         _ROUTER_SCALE = self.router_scale
 
-        # TODO(angazenn): Does not support MC2 currently
-        get_forward_context().moe_comm_method_name = "allgathercommimpl"
-
         if not use_h2p():
             final_hidden_states = self.experts.forward_impl(
                 hidden_states=hidden_states, router_logits=router_logits)

vllm_ascend/worker/model_runner_v1.py

@@ -88,7 +88,8 @@
                                   sanity_check_mm_encoder_outputs,
                                   scatter_mm_placeholders)
 
-from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.ascend_config import (MoECommImpl,
+                                       get_ascend_config)
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 from vllm_ascend.attention.attention_mask import AttentionMaskBuilder
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
@@ -1692,7 +1693,7 @@ def _pool(
         )
 
     def _select_moe_comm_method(self, num_tokens: int,
-                                with_prefill: bool) -> str:
+                                with_prefill: bool) -> MoECommImpl:
         """1. If expert parallel is not enabled, we use all-gather since MC2 and all-to-all
         are designed for expert parallelism.
         2. If expert parallel is enabled, we need to consider the soc version and the
@@ -1713,38 +1714,42 @@ def _select_moe_comm_method(self, num_tokens: int,
             ValueError: If the soc version is unsupported.
 
         Returns:
-            str: The selected MoE communication method, either "allgather", "mc2", or "alltoall".
+            MoECommImpl: The selected MoE communication method.
         """
         soc_version = get_ascend_soc_version()
         quant_type = getattr(self.vllm_config.model_config.hf_config,
                              'moe_quantize', None)
         model_type = self.vllm_config.model_config.hf_config.model_type
 
         if not self.parallel_config.enable_expert_parallel:
-            moe_comm_method = "allgather"
+            moe_comm_method = MoECommImpl.ALLGATHER
         elif soc_version in {AscendSocVersion.A2}:
-            if num_tokens <= self.mc2_tokens_capacity and self.parallel_config.world_size_across_dp >= 16:
-                moe_comm_method = "mc2"
+            if (num_tokens <= self.mc2_tokens_capacity
+                    and self.parallel_config.world_size_across_dp >= 16):
+                moe_comm_method = MoECommImpl.MC2
             else:
                 if quant_type == "w4a8_dynamic":
-                    moe_comm_method = "alltoall"
+                    moe_comm_method = MoECommImpl.ALLTOALL
                 else:
-                    moe_comm_method = "allgather"
+                    moe_comm_method = MoECommImpl.ALLGATHER
 
         elif soc_version in {AscendSocVersion.A3}:
-            moe_comm_method = "mc2" if num_tokens <= self.mc2_tokens_capacity else "alltoall"
+            moe_comm_method = (MoECommImpl.MC2
+                               if num_tokens <= self.mc2_tokens_capacity else
+                               MoECommImpl.ALLTOALL)
         else:
             raise ValueError(f"Unsupported soc_version: {soc_version}")
 
-        if moe_comm_method == "allgather" and with_prefill:
-            moe_comm_method = "naivemulticast"
-            
+        if moe_comm_method == MoECommImpl.ALLGATHER and with_prefill:
+            moe_comm_method = MoECommImpl.NAIVE_MULTICAST
+
         if model_type == "PanguProMoE":
-            moe_comm_method = "allgather"
+            moe_comm_method = MoECommImpl.ALLGATHER
 
         if is_global_first_rank():
             logger.debug(f"num_tokens: {num_tokens}, "
                          f"moe_comm_method: {moe_comm_method}")
+        print("moe_comm_method = ", moe_comm_method)
 
         return moe_comm_method
 
@@ -1777,8 +1782,8 @@ def execute_model(
             if self.dynamic_eplb:
                 self.eplb_updator.take_update_info_from_eplb_process()
 
-        moe_comm_method = self._select_moe_comm_method(num_input_tokens,
-                                                       self.with_prefill)
+        moe_comm_method_type = self._select_moe_comm_method(num_input_tokens,
+                                                            self.with_prefill)
 
         batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
                                            uniform_decode=False)
@@ -1794,7 +1799,7 @@ def execute_model(
                     num_tokens_across_dp=num_tokens_across_dp,
                     with_prefill=self.with_prefill,
                     reserved_mc2_mask=self.reserved_mc2_mask,
-                    moe_comm_method=moe_comm_method,
+                    moe_comm_method_type=moe_comm_method_type,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
                     batch_descriptor=batch_descriptor,
                     num_actual_tokens=scheduler_output.
@@ -2151,8 +2156,8 @@ def _dummy_run(
         (num_tokens, num_tokens_across_dp, with_prefill,
          _) = self._sync_metadata_across_dp(num_tokens, with_prefill, False)
 
-        moe_comm_method = self._select_moe_comm_method(num_tokens,
-                                                       with_prefill)
+        moe_comm_method_type = self._select_moe_comm_method(num_tokens,
+                                                            with_prefill)
 
         # If cudagraph_mode.decode_mode() == FULL and
         # cudagraph_mode.seperate_routine(). This means that we are using
@@ -2268,7 +2273,7 @@ def dummy_compute_logits(hidden_states):
                     with_prefill=with_prefill,
                     in_profile_run=self.in_profile_run,
                     reserved_mc2_mask=self.reserved_mc2_mask,
-                    moe_comm_method=moe_comm_method,
+                    moe_comm_method_type=moe_comm_method_type,
                     num_actual_tokens=0,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
                     batch_descriptor=batch_descriptor,