From 300e33797fac8a949a0ed89fab2633de96e65fa8 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Fri, 10 Jul 2026 15:37:51 +0800 Subject: [PATCH] [Perf] fuse more rmsnorm and all-reduce in qwen3.5 (#46998) Signed-off-by: zjy0516 --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 36 +++++++++---------- vllm/model_executor/models/qwen3_next.py | 16 +++------ 2 files changed, 21 insertions(+), 31 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5de2..6e10b9a9932 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -844,17 +844,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): def forward( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): - self._forward_method(hidden_states, output) + ) -> torch.Tensor: + return self._forward_method(hidden_states) def _output_projection( self, core_attn_out: torch.Tensor, z: torch.Tensor, - output: torch.Tensor, - num_tokens: int, - ): + ) -> torch.Tensor: """Part 3: RMSNormGated + output linear projection. The RMSNormGated + quant sequence is eligible for fusion @@ -866,13 +863,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + output, _ = self.out_proj(core_attn_out) + return output def forward_hip( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ROCm forward using AITER Triton fused projection+attention when available, otherwise falling back to the generic CUDA path.""" if GDN_AITER_TRITON_AVAILABLE: @@ -901,15 +898,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): use_aiter=True, ) - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z) else: - self.forward_cuda(hidden_states, output) + return self.forward_cuda(hidden_states) def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ Forward pass with three parts: 1. Input projection @@ -964,13 +960,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # ============================================================ # Part 3: Output Projection # ============================================================ - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z) def forward_xpu( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ Forward pass with three parts: 1. Input projection @@ -1013,13 +1008,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + out, _ = self.out_proj(core_attn_out) + return out def forward_cpu( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU." mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) @@ -1063,7 +1058,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + out, _ = self.out_proj(core_attn_out) + return out def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None: """Warm up GDN prefill kernels during V1 profiling. diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ef320dd526b..9a18cd4aad7 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -381,15 +381,15 @@ class Qwen3NextAttention(nn.Module): def forward( self, positions: torch.Tensor, - output: torch.Tensor, hidden_states: torch.Tensor, - ): + ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v, gate = self._project_qkv_gate(qkv, positions) attn_output = self.attn(q, k, v) if gate is not None: attn_output = attn_output * torch.sigmoid(gate) - output[:], _ = self.o_proj(attn_output) + output, _ = self.o_proj(attn_output) + return output class Qwen3NextDecoderLayer(nn.Module): @@ -484,21 +484,15 @@ class Qwen3NextDecoderLayer(nn.Module): else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - self_attention_output = torch.empty_like(hidden_states) if self.layer_type == "linear_attention": - self.linear_attn( - hidden_states=hidden_states, - output=self_attention_output, - ) + hidden_states = self.linear_attn(hidden_states=hidden_states) elif self.layer_type == "full_attention": - self.self_attn( + hidden_states = self.self_attn( hidden_states=hidden_states, - output=self_attention_output, positions=positions, ) else: raise ValueError("Invalid layer_type") - hidden_states = self_attention_output if self.layer_scale: if len(hidden_states.shape) == 2: