support MERRouter

Signed-off-by: angazenn <zengyanjia@huawei.com>

Author: angazenn

vllm_ascend/models/pangu_moe.py

@@ -49,7 +49,7 @@
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import IntermediateTensors
 
-from vllm_ascend.distributed.parallel_state import get_ep_group
+from vllm_ascend.utils import is_310p
 
 logger = init_logger(__name__)
 
@@ -95,41 +95,69 @@ def forward(self, x):
         return x
 
 
-class PanguProMoESparseMoeBlock(nn.Module):
+def topk_wrapper(num_voted_experts):
 
-    @staticmethod
     def pangu_group8_topk(
         hidden_states: torch.Tensor,
         gating_output: torch.Tensor,
         topk: int,
-        renormalize: bool,
+        renormalize: bool = False,
         num_expert_group: int = 0,
         topk_group: int = 0,
         global_num_experts: int = 0,
     ):
-        ep_size = get_ep_group().world_size
-        local_num_experts = global_num_experts // ep_size
-        local_num_group = topk // ep_size
-        router_scale = _ROUTER_SCALE.squeeze()  # type: ignore
-        scores = F.softmax(gating_output, dim=1)
-        scores = scores[...,
-                        get_ep_group().rank_in_group *
-                        local_num_experts:(get_ep_group().rank_in_group + 1) *
-                        local_num_experts]
-
-        router_weights = router_scale[get_ep_group().rank_in_group *
-                                      local_num_experts:
-                                      (get_ep_group().rank_in_group + 1) *
-                                      local_num_experts]
-        topk_weights, topk_ids = torch.max(scores.view(scores.shape[0],
-                                                       local_num_group, -1),
-                                           dim=-1)
-        bias = torch.arange(0,
-                            local_num_experts,
-                            topk,
-                            device=scores.device,
-                            dtype=torch.int32).unsqueeze(0)
-        topk_ids = topk_ids.to(torch.int32) + bias
+        scores = F.softmax(gating_output, dim=1, dtype=torch.float16)
+        num_tokens = scores.shape[0]
+        router_weights = _ROUTER_SCALE.squeeze(  # type: ignore
+        ).to(torch.float16)
+
+        if num_voted_experts == 8:
+            # use original topk
+            topk_weights, topk_ids = torch.max(scores.view(
+                scores.shape[0], topk, -1),
+                                               dim=-1)
+            bias = torch.arange(0,
+                                global_num_experts,
+                                topk,
+                                device=scores.device,
+                                dtype=torch.int32).unsqueeze(0)
+            topk_ids = topk_ids.to(torch.int32) + bias
+
+        else:
+            experts_per_group = global_num_experts // topk
+            group_expert_indices = torch.arange(experts_per_group,
+                                                dtype=torch.int32,
+                                                device=scores.device).view(
+                                                    1, 1, -1)
+            group_expert_offset = (
+                torch.arange(topk, dtype=torch.int32, device=scores.device) *
+                experts_per_group).unsqueeze(0)
+            expert_index_range = torch.arange(experts_per_group,
+                                              dtype=torch.int32,
+                                              device=scores.device)
+
+            scores_grouped = scores.view(num_tokens, topk, experts_per_group)
+            best_expert_idx = torch.argmax(scores_grouped,
+                                           dim=2)  # (num_tokens, num_groups)
+            vote_mask = (best_expert_idx.unsqueeze(-1).to(
+                torch.int32) == group_expert_indices).to(torch.float16)
+
+            expert_vote_freq = vote_mask.sum(dim=0)
+
+            sorted_indices = torch.argsort(expert_vote_freq,
+                                           dim=1,
+                                           descending=True).to(torch.int32)
+            topk_experts = sorted_indices[:, :num_voted_experts]
+            keep_mask = ((
+                topk_experts.unsqueeze(-1) == expert_index_range).any(
+                    dim=1)).unsqueeze(0)
+
+            masked_scores = torch.where(keep_mask, scores_grouped, 0)
+
+            topk_weights, best_pos_in_group = masked_scores.max(dim=2)
+            best_pos_in_group = best_pos_in_group.to(torch.int32)
+            topk_ids = (best_pos_in_group + group_expert_offset).to(
+                torch.int32)
 
         flatten_topk_ids = topk_ids.view(-1)
         router_weights = router_weights.index_select(0, flatten_topk_ids).view(
@@ -138,6 +166,11 @@ def pangu_group8_topk(
 
         return topk_weights, topk_ids
 
+    return pangu_group8_topk
+
+
+class PanguProMoESparseMoeBlock(nn.Module):
+
     def __init__(
         self,
         config: PretrainedConfig,
@@ -153,23 +186,23 @@ def __init__(
                 f"Tensor parallel size {self.tp_size} is greater than "
                 f"the number of experts {config.num_experts}.")
 
-        self.local_num_group = config.num_experts_per_tok // get_ep_group(
-        ).world_size
         self.num_experts_per_tok = config.num_experts_per_tok
-        self.local_num_experts = config.num_experts // get_ep_group(
-        ).world_size
         self.router_scale = torch.nn.Parameter(
             torch.ones((1, self.num_experts)))
 
+        # on 300I Duo platform, we find that num_voted_experts set to 5 achieves
+        # good performance without sacrifice too much accuracy. for other platform,
+        # this is set to 8 to use original pangu grouped topk.
+        num_voted_experts = 5 if is_310p() else 8
+
         self.experts = FusedMoE(
             num_experts=config.num_experts,
             top_k=config.num_experts_per_tok,
             hidden_size=config.hidden_size,
             intermediate_size=config.moe_intermediate_size,
             reduce_results=False,
             quant_config=quant_config,
-            custom_routing_function=PanguProMoESparseMoeBlock.
-            pangu_group8_topk,
+            custom_routing_function=topk_wrapper(num_voted_experts),
             prefix=f"{prefix}.experts",
         )
 

vllm_ascend/ops/fused_moe.py

@@ -578,6 +578,13 @@ def fused_experts_310p(
     local_num_experts = global_num_experts // ep_size
     local_num_group = top_k // ep_size
 
+    if ep_size > 1:
+        ep_rank = get_ep_group().rank_in_group
+        local_group_start = ep_rank * local_num_experts
+        local_group_end = (ep_rank + 1) * local_num_experts
+        topk_ids = topk_ids[:, local_group_start:local_group_end]
+        topk_weights = topk_weights[:, local_group_start:local_group_end]
+
     if apply_router_weight_on_input:
         assert (topk_weights.dim() == 2
                 ), "`topk_weights` should be in shape (num_tokens, topk)"