bugfix

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

Author: shenchuxiaofugui

vllm_ascend/ops/expert_load_balancer.py

@@ -14,8 +14,6 @@ def __init__(self, expert_map_path, num_experts):
         self.tensor_data = []
         self.expert_map_tensor, self.layers_num, self.ranks_num = (
             self._expert_file_to_tensor())
-        self.global_expert_num = num_experts + self.get_global_redundant_expert_num(
-        )
         self.expert_placement_map = self.generate_expert_placement_map()
 
     def _expert_file_to_tensor(self):
@@ -47,7 +45,7 @@ def generate_index_dicts(self, tensor_2d):
 
     def generate_expert_placement_map(self):
         expert_placement_map = torch.full(
-            (self.layers_num, self.ranks_num, self.global_expert_num),
+            (self.layers_num, self.ranks_num, self.num_experts),
             -1,
             dtype=torch.int32,
         )
@@ -70,7 +68,7 @@ def generate_log2phy_expert_map(self, layer_id):
                 result_dict[key] = []
             result_dict[key].append(idx)
 
-        log2phy_map = torch.full((self.ranks_num, self.global_expert_num),
+        log2phy_map = torch.full((self.ranks_num, self.num_experts),
                                  -1,
                                  dtype=torch.int32)
         for rank in range(self.ranks_num):

vllm_ascend/ops/fused_moe/token_dispatcher.py

@@ -118,18 +118,15 @@ def get_dispatch_mc2_kwargs(
         mc2_mask: torch.Tensor,
         global_redundant_expert_num: int = 0,
     ):
-        if self.with_quant:
-            quant_mode = 2
-            moe_expert_num = len(expert_map)
-        else:
-            quant_mode = 0
-            moe_expert_num = len(expert_map)
+        quant_mode = 2 if self.with_quant else 0
+        self.physics_num_experts = len(expert_map) + global_redundant_expert_num
+
         kwargs_mc2 = {
             "x": hidden_states,
             "expert_ids": topk_ids,
             "expert_shard_type": 0,
             "shared_expert_rank_num": 0,
-            "moe_expert_num": moe_expert_num,
+            "moe_expert_num": self.physics_num_experts,
             "global_bs": 0,
             "expert_token_nums_type": 0,
         }
@@ -247,15 +244,14 @@ def get_combine_mc_kwargs(self, hidden_states: torch.Tensor,
         expand_scales = context_metadata["expand_scales"]
 
         assert expert_map is not None
-        moe_expert_num = len(expert_map)
 
         kwargs_mc2 = {
             "expand_x": hidden_states,
             "expert_ids": topk_ids,
             "expert_scales": topk_weights.to(torch.float32),
             "expert_shard_type": 0,
             "shared_expert_rank_num": 0,
-            "moe_expert_num": moe_expert_num,
+            "moe_expert_num": self.physics_num_experts,
             "global_bs": 0,
         }
 
@@ -360,7 +356,7 @@ def token_dispatch(self,
             hidden_states = hidden_states * \
                 topk_weights.to(hidden_states.dtype)
         if expert_map is not None:
-            global_num_experts = len(expert_map)
+            global_num_experts = len(expert_map) + global_redundant_expert_num
             mask = (expert_map[topk_ids] != -1)
             topk_weights = topk_weights * mask
             first_expert_idx = get_ep_group(