fix

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

Author: shenchuxiaofugui

vllm_ascend/ops/common_fused_moe.py

@@ -279,6 +279,12 @@ def forward_impl(self, hidden_states: torch.Tensor,
         quantized_x_for_share, dynamic_scale_for_share = None, None
 
         forward_context = get_forward_context()
+
+        # Load balancing for token distribution among experts in dummy_run
+        # TODO: The community only considers load balancing when DP > 1.
+        # This approach may overlook some extreme scenarios.
+        enable_force_load_balance = forward_context.in_profile_run
+
         hidden_states, router_logits = forward_context.moe_comm_method.prepare(
             hidden_states=hidden_states,
             router_logits=router_logits,
@@ -305,6 +311,7 @@ def forward_impl(self, hidden_states: torch.Tensor,
             quantized_x_for_share=quantized_x_for_share,
             dynamic_scale_for_share=dynamic_scale_for_share,
             shared_experts=None,
+            enable_force_load_balance=enable_force_load_balance,
             log2phy=self.log2phy,
             global_redundant_expert_num=self.global_redundant_expert_num)
 

vllm_ascend/quantization/w4a8_dynamic.py

@@ -371,7 +371,8 @@ def apply(
         # to avoid accumulating too much tokens on a single rank.
         # currently it is only activated when doing profile runs.
         if enable_force_load_balance:
-            topk_ids = torch.randint_like(topk_ids, 0, global_num_experts)
+            topk_ids = torch.randint_like(
+                topk_ids, 0, global_num_experts - global_redundant_expert_num)
 
         topk_weights = topk_weights.to(x.dtype)
 

vllm_ascend/torchair/quantization/torchair_w8a8_dynamic.py

@@ -925,7 +925,7 @@ def apply(
         scoring_func: str = "softmax",
         e_score_correction_bias: Optional[torch.Tensor] = None,
         is_prefill: bool = True,
-        enable_force_load_balance: bool = True,
+        enable_force_load_balance: bool = False,
         log2phy: torch.Tensor = None,
         global_redundant_expert_num: int = 0,
         shared_experts: Optional[Any] = None,