eplb-fix-dev

Signed-off-by: Che Ruan <[email protected]>

Author: mercykid

vllm_ascend/eplb/core/eplb_device_transfer_loader.py

@@ -50,10 +50,6 @@ def generate_expert_d2d_transfer_task(self, expert_send_info,
             )
             return
 
-        # If neither send nor receive task is needed for this layer on this rank, return
-        if not (expert_send_info or expert_recv_info):
-            return
-
         self.updated_expert_map = updated_expert_map
 
         self.layer_id = layer_id

vllm_ascend/ops/expert_load_balancer.py

@@ -48,7 +48,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,
         )
@@ -71,7 +71,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/moe/moe_mlp.py

@@ -105,14 +105,17 @@ def quant_apply_mlp(hidden_states: torch.Tensor,
                 group_list=group_list,
                 output_dtype=torch.int32)[0]
             # act_fn: swiglu
+            group_diff = torch.diff(group_list)
+            new_group = torch.cat([group_list[0].unsqueeze(0), group_diff],
+                                  dim=0)
             hidden_states, swiglu_out_scale = torch_npu.npu_dequant_swiglu_quant(
                 x=hidden_states,
                 weight_scale=w1_scale,
                 activation_scale=pertoken_scale,
                 bias=None,
                 quant_scale=None,
                 quant_offset=None,
-                group_index=group_list,
+                group_index=new_group,
                 activate_left=True,
                 quant_mode=1,
             )

vllm_ascend/ops/moe/token_dispatcher.py

@@ -122,18 +122,14 @@ def get_dispatch_mc2_kwargs(
         expert_map: 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.moe__expert_num = 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.moe_expert_num,
             "global_bs": 0,
             "expert_token_nums_type": 0,
         }
@@ -229,15 +225,14 @@ def get_combine_mc_kwargs(self, hidden_states: torch.Tensor):
         assert self.topk_weights is not None
         assert self.topk_ids is not None
         assert self.output is not None
-        moe_expert_num = len(self.expert_map)
         # moeCombine
         kwargs_mc2 = {
             "expand_x": hidden_states,
             "expert_ids": self.topk_ids,
             "expert_scales": self.topk_weights.to(torch.float32),
             "expert_shard_type": 0,
             "shared_expert_rank_num": 0,
-            "moe_expert_num": moe_expert_num,
+            "moe_expert_num": self.moe_expert_num,
             "global_bs": 0,
         }
         if self.with_quant:
@@ -359,7 +354,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)
             self.topk_weights = topk_weights * mask
             first_expert_idx = get_ep_group(

vllm_ascend/quantization/w8a8_dynamic.py

@@ -250,7 +250,7 @@ def apply(
         return moe_comm_method.fused_experts(
             hidden_states=x,
             w1=layer.w13_weight,
-            w1_scale=layer.w13_weight_scale_fp32,
+            w1_scale=layer.w13_weight_scale.to(torch.float32),
             w2=layer.w2_weight,
             w2_scale=layer.w2_weight_scale,
             topk_weights=topk_weights,