allgather after rmsnorm

Author: unknown

vllm_ascend/ascend_forward_context.py

@@ -87,10 +87,14 @@ def set_ascend_forward_context(
         is_deepseek_v3_r1 = hasattr(
             vllm_config.model_config.hf_config, 'n_routed_experts'
         ) and vllm_config.model_config.hf_config.n_routed_experts == 256
+        is_glm4_moe = hasattr(
+            vllm_config.model_config.hf_config, 'n_routed_experts'
+        ) and vllm_config.model_config.hf_config.model_type == 'glm4_moe'
         fused_moe_state = _get_fused_moe_state(ep_size, with_prefill,
                                                is_deepseek_v3_r1)
         forward_context.fused_moe_state = fused_moe_state
         forward_context.in_profile_run = in_profile_run
+        forward_context.is_glm4_moe = is_glm4_moe
 
         from vllm_ascend.ops.moe.token_dispatcher import get_token_dispatcher
         dispatcher_name = _moe_method_to_dispatcher[moe_comm_method]

vllm_ascend/ops/common_fused_moe.py

@@ -327,11 +327,6 @@ 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
-        if flashcomm_v1_enabled:
-            hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(hidden_states, True)
-            if router_logits.shape[0] != hidden_states.shape[0]:
-                router_logits = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(router_logits, True)
 
         forward_context.moe_comm_method = getattr(self, moe_comm_method_name)
 

vllm_ascend/ops/layernorm.py

@@ -19,6 +19,7 @@
 
 import torch
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.forward_context import get_forward_context
 
 
 class AddRMSNormW8A8Quant(RMSNorm):
@@ -54,6 +55,8 @@ def forward(
                 self.layer.aclnn_input_offset,
                 epsilon=self.variance_epsilon)
             torch.ops.vllm.maybe_wait_prefetch_done(x)
+            is_glm4_moe = get_forward_context().is_glm4_moe
+            x = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(x, is_glm4_moe)
             return x, residual
 
         x, residual = torch_npu.npu_rms_norm(x, self.weight,
@@ -84,6 +87,8 @@ def forward_oot(
                 x, _, residual = torch_npu.npu_add_rms_norm(
                     x, residual, self.weight, self.variance_epsilon)
             torch.ops.vllm.maybe_wait_prefetch_done(x)
+            is_glm4_moe = get_forward_context().is_glm4_moe
+            x = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(x, is_glm4_moe)
             return x, residual
 
         x, residual = torch_npu.npu_rms_norm(x, self.weight,

vllm_ascend/ops/linear.py

@@ -403,7 +403,7 @@ def __init__(
         if prefix.find("gate_up_proj") != -1 and mlp_tp_enable():
             comm_group = get_mlp_tp_group()
             self.forward_type = "mlp_tp"
-        elif dense_optim_enable():
+        elif prefix.find("shared_experts") == -1 and dense_optim_enable():
             comm_group = get_tp_group()
             self.forward_type = "dense_optim"
         else: