fix mc2 bug

Signed-off-by: yangcheng (AJ) <yangcheng104@huawei.com>

Author: yangcheng (AJ)

vllm_ascend/models/qwen3_moe.py

@@ -16,31 +16,23 @@
 # Adapted from vllm/model_executor/models/qwen3_moe.py
 # This file is a part of the vllm-ascend project.
 
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Optional
 
 import torch
-import torch.distributed as dist
-import torch_npu
-import vllm
-import vllm.envs as envs
 from torch import nn
 from transformers import PretrainedConfig
+from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.distributed.parallel_state import get_ep_group
+from vllm_ascend.ops.fused_moe import AscendFusedMoE
+
+import vllm
 from vllm.attention import AttentionMetadata
-from vllm.distributed import (get_tensor_model_parallel_world_size,
-                              get_tp_group)
+from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group
 from vllm.distributed.parallel_state import get_dp_group
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.linear import ReplicatedLinear
-                                               
 from vllm.model_executor.layers.quantization import QuantizationConfig
-
-from vllm_ascend.ascend_config import get_ascend_config
-from vllm_ascend.distributed.parallel_state import get_ep_group
-from vllm_ascend.ops.fused_moe import AscendFusedMoE
-
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
-from transformers import PretrainedConfig
-from vllm.model_executor.layers.quantization import QuantizationConfig
 
 
 class CustomQwen3MoeForCausalLM(Qwen3MoeForCausalLM):

vllm_ascend/ops/fused_moe.py

@@ -118,9 +118,13 @@ def fused_experts_with_mc2(
     top_k: int,
     expert_map: torch.Tensor = None,
     moe_all_to_all_group_name: Optional[str] = None,
-    shared_experts: Optional[Any] = None
+    shared_experts: Optional[Any] = None,
+    global_batch_size: int = 256,
 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-    global_bs = 0
+
+    ep_group = get_ep_group().device_group
+    all_to_all_group_size = torch.distributed.get_world_size(ep_group)
+    global_bs = global_batch_size * all_to_all_group_size
     moe_expert_num = len(expert_map)
     kwargs_mc2 = {
         "x": hidden_states,
@@ -132,11 +136,8 @@ def fused_experts_with_mc2(
     }
 
     rank = torch.distributed.get_rank()
-
     quant_mode = 0
-    ep_group = get_ep_group().device_group
     local_rank = torch.distributed.get_rank(group=ep_group)
-    all_to_all_group_size = torch.distributed.get_world_size(ep_group)
 
     tp_size = get_etp_group().world_size
     tp_rank = rank % tp_size
@@ -204,7 +205,7 @@ def fused_experts_with_mc2(
         "expert_shard_type": 0,
         "shared_expert_rank_num": 0,
         "moe_expert_num": moe_expert_num,
-        "global_bs": 0,
+        "global_bs": global_bs,
     }
     tp_recv_counts = output[5]
     stage3_kwargs = {
@@ -1037,7 +1038,8 @@ def apply(
                 top_k=top_k,
                 expert_map=expert_map,
                 moe_all_to_all_group_name=self.moe_all_to_all_group_name,
-                shared_experts=shared_experts)
+                shared_experts=shared_experts,
+                global_batch_size=self.global_batch_size)
         elif fused_moe_state == FusedMoEState.AllGather:
             return fused_experts(hidden_states=x,
                                  w1=layer.w13_weight,