[EPLB][BugFix] generate experts map with redundant experts

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

Author: shenchuxiaofugui

vllm_ascend/eplb/core/eplb_utils.py

@@ -25,6 +25,29 @@
 import vllm_ascend.envs as envs_ascend
 
 
+def generate_experts_map(ep_rank, ep_size, n_expert, n_redundant):
+    def split_and_insert(n, k, m):
+        all_experts = torch.arange(n)
+        groups = torch.array_split(all_experts, k)
+        for i in range(m):
+            j = i % k + 1
+            if len(groups[-j]) == 0:
+                groups[-j] = torch.append(groups[-j], j)
+            else:
+                groups[-j] = torch.append(groups[-j], (groups[-j][-1] + 1) % n_expert)
+        return torch.concatenate(groups)
+
+    random_placement = split_and_insert(n_expert, ep_size, n_redundant)
+    global_num_experts = random_placement.shape[0]
+    local_num_experts = global_num_experts // ep_size
+
+    expert_map = torch.full((random_placement.shape[0]), -1, dtype=torch.int32)
+    expert_map[ep_rank * local_num_experts: (ep_rank + 1) * local_num_experts] = \
+        random_placement[ep_rank * local_num_experts: (ep_rank + 1) * local_num_experts]
+
+    return expert_map
+
+
 def generate_log2phy_map(expert_map):
     num_local_experts = expert_map.max() + 1
     log2phy_map = expert_map.clone()

vllm_ascend/ops/fused_moe/fused_moe.py

@@ -35,7 +35,8 @@
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import MoECommType
 from vllm_ascend.distributed.parallel_state import get_mc2_group
-from vllm_ascend.eplb.core.eplb_utils import determine_default_log2phy_map
+from vllm_ascend.eplb.core.eplb_utils import (determine_default_log2phy_map,
+                                              generate_experts_map)
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
 from vllm_ascend.ops.fused_moe.experts_selector import select_experts
 from vllm_ascend.ops.fused_moe.moe_comm_method import setup_moe_comm_method
@@ -182,8 +183,8 @@ def __init__(self, *args, **kwargs):
                 dtype=vllm_config.model_config.dtype)
 
         # init moe.
-        self.local_num_experts, self.expert_map, _ = determine_expert_map(
-            self.ep_size, self.ep_rank, self.global_num_experts)
+        self.local_num_experts, self.expert_map = generate_experts_map(
+            self.ep_size, self.ep_rank, num_experts, self.global_redundant_expert_num)
         # TODO: Temporary flag to indicate if static EPLB is enabled. This is a
         # workaround to bypass a quantization check that fails with float weights.
         init_eplb_enable = False