fix sharedfusedmoe decision method

Co-authored-by: realliujiaxu <realliujiaxu@163.com>

Signed-off-by: zhaozx-cn <zhaozx2116@163.com>

Author: zhaozx-cn

vllm_ascend/ascend_forward_context.py

@@ -144,12 +144,7 @@ def set_ascend_forward_context(
             forward_context.num_hidden_layers = vllm_config.model_config.hf_config.num_hidden_layers
             forward_context.fusion_linear = "gate_up_dense" if forward_context.layer_idx == 0 else "qkv_dense"
         forward_context.addrmsnorm_quant_fusion_enabled = addrmsnorm_quant_fusion_enabled
-        is_shared_fused_moe = hasattr(vllm_config.model_config.hf_config,
-                                      'n_shared_experts')
-        if is_shared_fused_moe:
-            forward_context.n_shared_experts = vllm_config.model_config.hf_config.n_shared_experts
-        else:
-            forward_context.n_shared_experts = 0
+        
         if num_tokens is None and attn_metadata is not None:
             num_tokens = attn_metadata.num_actual_tokens
 

vllm_ascend/ops/common_fused_moe.py

@@ -311,9 +311,8 @@ def maybe_all_reduce_tensor_model_parallel(
         """
         forward_context = get_forward_context()
         moe_comm_method_name = forward_context.moe_comm_method_name
-        flashcomm_v1_enabled = forward_context.flashcomm_v1_enabled
         if moe_comm_method_name in {"alltoallcommimpl", "mc2commimpl"}:
-            if flashcomm_v1_enabled:
+            if forward_context.flashcomm_v1_enabled:
                 pad_size = forward_context.pad_size
                 if pad_size > 0:
                     final_hidden_states = F.pad(final_hidden_states,
@@ -333,9 +332,8 @@ def forward_impl(self, hidden_states: torch.Tensor,
 
         forward_context = get_forward_context()
         moe_comm_method_name = forward_context.moe_comm_method_name
-        flashcomm_v1_enabled = forward_context.flashcomm_v1_enabled
-        n_shared_experts = forward_context.n_shared_experts
-        if n_shared_experts == 0 and flashcomm_v1_enabled:
+
+        if self._shared_experts is None and forward_context.flashcomm_v1_enabled:
             hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
                 hidden_states, True)
             router_logits = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
@@ -445,8 +443,8 @@ def __init__(
         use_overlapped: bool = True,
         **kwargs,
     ):
-        AscendFusedMoE.__init__(self, **kwargs)
         self._shared_experts = shared_experts
+        AscendFusedMoE.__init__(self, **kwargs)
         self.use_overlapped = use_overlapped
         self.shared_expert_stream = None
         ascend_config = get_ascend_config()
@@ -460,8 +458,7 @@ def forward(
         router_logits: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         forward_context = get_forward_context()
-        flashcomm_v1_enabled = forward_context.flashcomm_v1_enabled
-        if flashcomm_v1_enabled:
+        if forward_context.flashcomm_v1_enabled:
             hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
                 hidden_states, True)
             router_logits = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(