use ep/tp size in fusedmoe parallel config

Signed-off-by: MengqingCao <cmq0113@163.com>

Author: MengqingCao

vllm_ascend/ops/common_fused_moe.py

@@ -85,6 +85,7 @@ def forward_oot(
             hidden_states=x,
             w1=layer.w13_weight,
             w2=layer.w2_weight,
+            moe_parallel_config=self.moe.moe_parallel_config,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             top_k=top_k,

vllm_ascend/ops/fused_moe.py

@@ -124,6 +124,7 @@ def fused_experts_with_mc2(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     top_k: int,
+    moe_parallel_config: FusedMoEParallelConfig,
     expert_map: torch.Tensor = None,
     moe_all_to_all_group_name: Optional[str] = None,
     shared_experts: Optional[Any] = None
@@ -142,11 +143,10 @@ def fused_experts_with_mc2(
     rank = torch.distributed.get_rank()
 
     quant_mode = 0
-    ep_group = get_ep_group()
-    ep_rank_id = ep_group.rank_in_group
-    ep_world_size = ep_group.world_size
+    ep_rank_id = moe_parallel_config.ep_rank
+    ep_world_size = moe_parallel_config.ep_size
 
-    tp_world_size = get_tp_group().world_size
+    tp_world_size = moe_parallel_config.tp_size
     tp_rank = rank % tp_world_size
 
     stage1_kwargs = {
@@ -559,6 +559,7 @@ def fused_experts_moge(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
     w2: torch.Tensor,
+    moe_parallel_config: FusedMoEParallelConfig,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     top_k: int,
@@ -580,7 +581,7 @@ def fused_experts_moge(
     Returns:
         hidden_states: Hidden states after routing.
     """
-    ep_size = get_ep_group().world_size
+    ep_size = moe_parallel_config.ep_size
     local_num_experts = global_num_experts // ep_size
     local_num_group = top_k // ep_size
 
@@ -981,7 +982,7 @@ def __init__(self, moe: FusedMoEConfig = None):
         vllm_config = get_current_vllm_config()
 
         self.ep_group = get_ep_group()
-        self.ep_size = self.ep_group.world_size
+        self.ep_size = self.moe.moe_parallel_config.ep_size
         self.global_batch_size = vllm_config.scheduler_config.max_num_seqs
         self.local_batch_size = self.global_batch_size // self.ep_size
         self.max_model_len = vllm_config.model_config.max_model_len
@@ -1073,13 +1074,14 @@ def apply(
         if enable_force_load_balance:
             topk_ids = torch.randint_like(topk_ids, 0, global_num_experts)
 
-        fused_moe_state = get_fused_moe_state(self.ep_group.world_size,
-                                              is_prefill, is_deepseek_v3_r1)
+        fused_moe_state = get_fused_moe_state(self.ep_size, is_prefill,
+                                              is_deepseek_v3_r1)
         if fused_moe_state == FusedMoEState.MC2:
             return fused_experts_with_mc2(
                 hidden_states=x,
                 w1=layer.w13_weight,
                 w2=layer.w2_weight,
+                moe_parallel_config=self.moe.moe_parallel_config,
                 topk_weights=topk_weights,
                 topk_ids=topk_ids,
                 top_k=top_k,