feat(moe): Refactor MoE communication method framework

This commit refactors the MoE communication method framework to improve modularity, clarity, and extensibility.

Key changes include:

- **Revised `MoECommMethod` Interface:**
  - Renamed `_pre_process` to `permute` and `_post_process` to `unpermute` for better clarity.
  - Introduced `prepare` and `finalize` methods to encapsulate logic that happens before and after the core MoE computation, such as tensor padding/splitting for MC2 and the final AllReduce.

- **Simplified `AscendFusedMoE`:**
  - The `forward_impl` is significantly simplified by delegating pre- and post-processing logic (padding, splitting, reduction) to the specific `MoECommMethod` implementation.
  - `AscendFusedMoE` now instantiates all communication method objects at initialization and selects the appropriate one at runtime based on a string identifier.

- **Centralized Expert Logic:**
  - Removed `unified_fused_experts` and introduced a new `fused_experts` function in `common_fused_moe.py`.
  - This new function utilizes the `permute`/`unpermute` methods from the `MoECommMethod` abstraction, decoupling the core expert logic from specific communication implementations.

- **Configuration and Invocation:**
  - The communication method is now selected and passed around as a string (e.g., "mc2", "allgather") instead of a class type, simplifying the invocation in the model runner.

These changes result in a cleaner separation of concerns, making the MoE implementation easier to understand, maintain, and extend with new communication strategies.

Signed-off-by: Yizhou Liu <[email protected]>

Author: yiz-liu

tests/e2e/multicard/moe/test_moe_comm.py

@@ -103,11 +103,11 @@ def test_all_gather_comm_impl(
         native_permuted_hidden,
         native_expert_tokens,
         _,
-    ) = native_impl._pre_process(hidden_states, topk_ids, topk_weights,
-                                 expert_map, num_experts)
+    ) = native_impl.permute(hidden_states, topk_ids, topk_weights, expert_map,
+                            num_experts)
     # Simulate MLP output
     native_mlp_output = torch.randn_like(native_permuted_hidden)
-    native_impl._post_process(native_mlp_output, native_hidden_states_out)
+    native_impl.unpermute(native_mlp_output, native_hidden_states_out)
 
     # --- Run AllGather Implementation ---
     all_gather_hidden_states_out = hidden_states.clone()

vllm_ascend/ascend_forward_context.py

@@ -11,7 +11,6 @@
                                   set_forward_context)
 
 import vllm_ascend.envs as envs_ascend
-from vllm_ascend.distributed.moe_comm_method import MoECommMethod
 
 
 class FusedMoEState(Enum):
@@ -57,7 +56,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: Optional[MoECommMethod] = None,
+        moe_comm_method: str = "",
         num_actual_tokens: Optional[int] = None,
         aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
         batch_descriptor: Optional[BatchDescriptor] = None):
@@ -75,7 +74,7 @@ def set_ascend_forward_context(
             batch_descriptor=batch_descriptor,
     ):
         forward_context = get_forward_context()
-        forward_context.moe_comm_method = moe_comm_method
+        forward_context.moe_comm_method_name = moe_comm_method + "commimpl"
         forward_context.with_prefill = with_prefill
         ep_size = (get_ep_group().world_size if
                    vllm_config.parallel_config.enable_expert_parallel else 1)

vllm_ascend/distributed/moe_comm_method.py

@@ -1,12 +1,14 @@
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import torch
+import torch.distributed as dist
+import torch.nn as nn
 import torch_npu
-from transformers.configuration_utils import PretrainedConfig
-from vllm.distributed.parallel_state import (
-    get_ep_group, get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size, get_tp_group)
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.distributed.parallel_state import get_ep_group, get_tp_group
 from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.model_executor.layers.fused_moe import FusedMoEConfig
 from vllm.utils import direct_register_custom_op
 
 from vllm_ascend.distributed.parallel_state import get_mc2_group
@@ -16,26 +18,36 @@
 class MoECommMethod(ABC):
     """Base class for MoE communication methods."""
 
-    def __init__(
-        self,
-        device: torch.device,
-        dtype: torch.dtype,
-        hf_config: PretrainedConfig,
-    ):
-        self.device = device
-        self.dtype = dtype
-        self.top_k_num = getattr(hf_config, "num_experts_per_tok", 0)
-        # global_num_experts may be called num_experts or n_routed_experts in different models.
-        possible_keys = ["num_experts", "n_routed_experts"]
-        for key in possible_keys:
-            if hasattr(hf_config, key):
-                self.global_num_experts = getattr(hf_config, key)
-                break
-        else:
-            self.global_num_experts = 0
+    moe_config: FusedMoEConfig = None
+
+    def __init__(self, moe_config: Optional[FusedMoEConfig]):
+        self.moe_config = moe_config
+
+    @abstractmethod
+    def prepare(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare the MoE communication method.
+
+        This method is called before quant_method.apply to prepare the
+        communication method. It can be used to initialize any necessary
+        resources or configurations.
+        """
+        pass
+
+    @abstractmethod
+    def finalize(self, hidden_states: torch.Tensor,
+                 reduce_results: bool) -> torch.Tensor:
+        """Finalize the MoE communication method.
+
+        This method is called after quant_method.apply to finalize the
+        communication method. It can be used to clean up any resources or
+        configurations.
+        """
+        pass
 
     @abstractmethod
-    def _pre_process(
+    def permute(
         self,
         hidden_states: torch.Tensor,
         topk_ids: torch.Tensor,
@@ -69,8 +81,8 @@ def _pre_process(
         pass
 
     @abstractmethod
-    def _post_process(self, mlp_output: torch.Tensor,
-                      hidden_states: torch.Tensor) -> None:
+    def unpermute(self, mlp_output: torch.Tensor,
+                  hidden_states: torch.Tensor) -> None:
         """Post-process after MLP.
 
         Args:
@@ -84,7 +96,18 @@ def _post_process(self, mlp_output: torch.Tensor,
 
 class DummyCommImpl(MoECommMethod):
 
-    def _pre_process(
+    def prepare(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Dummy prepare method that does nothing."""
+        return hidden_states, router_logits
+
+    def finalize(self, hidden_states: torch.Tensor,
+                 reduce_results: bool) -> torch.Tensor:
+        """Dummy finalize method that does nothing."""
+        return hidden_states
+
+    def permute(
         self,
         hidden_states: torch.Tensor,
         topk_ids: torch.Tensor,
@@ -96,8 +119,8 @@ def _pre_process(
         return moe_comm_pre_process_fake(hidden_states, topk_ids, topk_weights,
                                          expert_map, num_experts)
 
-    def _post_process(self, mlp_output: torch.Tensor,
-                      hidden_states: torch.Tensor) -> None:
+    def unpermute(self, mlp_output: torch.Tensor,
+                  hidden_states: torch.Tensor) -> None:
         """Dummy implementation that does nothing."""
         pass
 
@@ -110,7 +133,22 @@ class NativeAllGatherCommImpl(MoECommMethod):
     But it is a good fallback for scenarios where NPU-specific ops are not available.
     """
 
-    def _pre_process(
+    def prepare(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Dummy prepare method that does nothing."""
+        return hidden_states, router_logits
+
+    def finalize(self, hidden_states: torch.Tensor,
+                 reduce_results: bool) -> torch.Tensor:
+        if reduce_results and (self.moe_config.tp_size > 1
+                               or self.moe_config.ep_size > 1):
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                hidden_states)
+
+        return final_hidden_states
+
+    def permute(
         self,
         hidden_states: torch.Tensor,
         topk_ids: torch.Tensor,
@@ -122,10 +160,10 @@ def _pre_process(
 
         # Generate token indices and flatten
         token_indices = torch.arange(num_tokens,
-                                     device=self.device,
+                                     device=hidden_states.device,
                                      dtype=torch.int64)
         token_indices = (token_indices.unsqueeze(1).expand(
-            -1, self.top_k_num).reshape(-1))
+            -1, self.moe_config.experts_per_token).reshape(-1))
 
         # Flatten token-to-expert mappings and map to local experts
         weights_flat = topk_weights.view(-1)
@@ -140,7 +178,7 @@ def _pre_process(
         # This is a workaround and should be removed after the issue is fixed
         filtered_weights = torch.where(mask, weights_flat,
                                        torch.zeros_like(weights_flat)).to(
-                                           self.dtype)
+                                           topk_weights.dtype)
         filtered_experts = torch.where(
             mask,
             local_experts_flat,
@@ -156,7 +194,7 @@ def _pre_process(
         # This is equivalent to but faster than:
         # >>> token_counts = torch.bincount(filtered_experts, minlength=num_experts)[:-1]
         token_counts = torch.zeros(num_experts + 1,
-                                   device=self.device,
+                                   device=hidden_states.device,
                                    dtype=torch.int64)
         ones = torch.ones_like(filtered_experts, dtype=torch.int64)
         token_counts.scatter_add_(0, filtered_experts.to(torch.int64), ones)
@@ -169,8 +207,8 @@ def _pre_process(
 
         return permuted_hidden_states, expert_tokens, group_list_type
 
-    def _post_process(self, mlp_output: torch.Tensor,
-                      hidden_states: torch.Tensor) -> None:
+    def unpermute(self, mlp_output: torch.Tensor,
+                  hidden_states: torch.Tensor) -> None:
         mlp_output = mlp_output * self.sorted_weights.unsqueeze(1)
 
         final_hidden_states = torch.zeros_like(hidden_states)
@@ -199,7 +237,22 @@ class AllGatherCommImpl(MoECommMethod):
     This is a workaround and should be removed after the issue is fixed.
     """
 
-    def _pre_process(
+    def prepare(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Dummy prepare method that does nothing."""
+        return hidden_states, router_logits
+
+    def finalize(self, hidden_states: torch.Tensor,
+                 reduce_results: bool) -> torch.Tensor:
+        if reduce_results and (self.moe_config.tp_size > 1
+                               or self.moe_config.ep_size > 1):
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                hidden_states)
+
+        return final_hidden_states
+
+    def permute(
         self,
         hidden_states: torch.Tensor,
         topk_ids: torch.Tensor,
@@ -229,8 +282,8 @@ def _pre_process(
             torch_npu.npu_moe_init_routing_v2(
                 hidden_states,
                 topk_ids,
-                active_num=num_tokens * self.top_k_num,
-                expert_num=self.global_num_experts,
+                active_num=num_tokens * self.moe_config.experts_per_token,
+                expert_num=self.moe_config.num_experts,
                 expert_tokens_num_type=1,  # Only support `count` mode now
                 expert_tokens_num_flag=True,  # Output `expert_tokens`
                 active_expert_range=[first_expert_idx, last_expert_idx],
@@ -243,8 +296,8 @@ def _pre_process(
 
         return permuted_hidden_states, expert_tokens, group_list_type
 
-    def _post_process(self, mlp_output: torch.Tensor,
-                      hidden_states: torch.Tensor) -> None:
+    def unpermute(self, mlp_output: torch.Tensor,
+                  hidden_states: torch.Tensor) -> None:
         hidden_states[:] = torch_npu.npu_moe_token_unpermute(
             permuted_tokens=mlp_output,
             sorted_indices=self.expanded_row_idx,
@@ -261,19 +314,13 @@ class MC2CommImpl(MoECommMethod):
     Communication and Computation parallelism on Ascend devices.
     """
 
-    def __init__(
-        self,
-        device: torch.device,
-        dtype: torch.dtype,
-        hf_config: PretrainedConfig,
-    ):
-        super().__init__(device, dtype, hf_config)
+    def __init__(self, moe_config: Optional[FusedMoEConfig]):
+        super().__init__(moe_config)
 
         # Shared communication configurations
         ep_group = get_mc2_group()
         self.ep_rank_id = ep_group.rank_in_group
         self.ep_world_size = ep_group.world_size
-        self.tp_world_size = get_tp_group().world_size
 
         device_group = ep_group.device_group
         local_rank = torch.distributed.get_rank(group=device_group)
@@ -286,15 +333,51 @@ def __init__(
         self.is_ascend_a3 = get_ascend_soc_version() == AscendSocVersion.A3
         self.need_extra_args = self.is_ascend_a3  # or is_torchair
 
-        # Intermediate tensors to be passed from pre_process to post_process
-        self.topk_ids = None
-        self.topk_weights = None
-        self.mc2_mask = None
-        self.assist_info_for_combine = None
-        self.ep_recv_counts = None
-        self.tp_recv_counts = None
+    def prepare(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        num_tokens, _ = hidden_states.shape
+        self.mc2_mask = get_forward_context().mc2_mask
 
-    def _pre_process(
+        if num_tokens % self.moe_config.ep_size != 0:
+            pad_size = self.moe_config.ep_size - (num_tokens %
+                                                  self.moe_config.ep_size)
+            hidden_states = nn.functional.pad(hidden_states,
+                                              (0, 0, 0, pad_size))
+            router_logits = nn.functional.pad(router_logits,
+                                              (0, 0, 0, pad_size))
+
+        split_hidden_states = torch.tensor_split(hidden_states,
+                                                 self.moe_config.ep_size,
+                                                 dim=0)
+        split_router_logits = torch.tensor_split(router_logits,
+                                                 self.moe_config.ep_size,
+                                                 dim=0)
+        split_mc2_mask = torch.tensor_split(self.mc2_mask,
+                                            self.moe_config.ep_size,
+                                            dim=0)
+        self.num_tokens = num_tokens
+        self.split_hidden_states = split_hidden_states
+
+        hidden_states = split_hidden_states[self.moe_config.ep_rank]
+        router_logits = split_router_logits[self.moe_config.ep_rank]
+        self.mc2_mask = split_mc2_mask[self.moe_config.ep_rank]
+
+        return hidden_states, router_logits
+
+    def finalize(self, hidden_states: torch.Tensor,
+                 reduce_results: bool) -> torch.Tensor:
+        """Dummy finalize method that does nothing."""
+        tp_group = get_tp_group().device_group
+        dist.all_gather(list(self.split_hidden_states), hidden_states,
+                        tp_group)
+        final_hidden_states = torch.cat(self.split_hidden_states, dim=0)
+        if self.num_tokens % self.moe_config.ep_size != 0:
+            final_hidden_states = final_hidden_states[:self.num_tokens]
+
+        return final_hidden_states
+
+    def permute(
         self,
         hidden_states: torch.Tensor,
         topk_ids: torch.Tensor,
@@ -305,19 +388,13 @@ def _pre_process(
         # Store tensors needed for post_process
         self.topk_ids = topk_ids
         self.topk_weights = topk_weights.to(torch.float32)
-        self.mc2_mask = get_forward_context().mc2_mask
-
-        tp_size = get_tensor_model_parallel_world_size()
-        split_mc2_mask = torch.tensor_split(self.mc2_mask, tp_size, dim=0)
-        tp_rank = get_tensor_model_parallel_rank()
-        self.mc2_mask = split_mc2_mask[tp_rank]
 
         dispatch_kwargs = {
             "x": hidden_states,
             "expert_ids": self.topk_ids,
             "expert_shard_type": 0,
             "shared_expert_rank_num": 0,
-            "moe_expert_num": self.global_num_experts,
+            "moe_expert_num": self.moe_config.num_experts,
             "global_bs": 0,
             "scales": None,
             "quant_mode": 0,
@@ -352,15 +429,15 @@ def _pre_process(
 
         return permuted_hidden_states, expert_tokens, group_list_type
 
-    def _post_process(self, mlp_output: torch.Tensor,
-                      hidden_states: torch.Tensor) -> None:
+    def unpermute(self, mlp_output: torch.Tensor,
+                  hidden_states: torch.Tensor) -> None:
         combine_kwargs = {
             "expand_x": mlp_output,
             "expert_ids": self.topk_ids,
             "expert_scales": self.topk_weights,
             "expert_shard_type": 0,
             "shared_expert_rank_num": 0,
-            "moe_expert_num": self.global_num_experts,
+            "moe_expert_num": self.moe_config.num_experts,
             "global_bs": 0,
             "ep_send_counts": self.ep_recv_counts,
             "group_ep": self.moe_all_to_all_group_name,