[Perf] fuse more rmsnorm and all-reduce in qwen3.5 (#46998)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Jul 10, 2026 at 15:37 UTC
300e33797fac8a949a0ed89fab2633de96e65fa8
2 files changed
+21
-31
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+16
-20
@@ -844,17 +844,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
844
def forward(
845
self,
846
hidden_states: torch.Tensor,
847
- output: torch.Tensor,
848
- ):
849
- self._forward_method(hidden_states, output)
847
+ ) -> torch.Tensor:
848
+ return self._forward_method(hidden_states)
849
850
def _output_projection(
851
self,
852
core_attn_out: torch.Tensor,
853
z: torch.Tensor,
855
- output: torch.Tensor,
856
- num_tokens: int,
857
- ):
854
+ ) -> torch.Tensor:
855
"""Part 3: RMSNormGated + output linear projection.
856
857
The RMSNormGated + quant sequence is eligible for fusion
@@ -866,13 +863,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
863
core_attn_out = self.norm(core_attn_out, z)
864
core_attn_out = core_attn_out.reshape(z_shape_og)
865
core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)
869
- output[:num_tokens], _ = self.out_proj(core_attn_out)
866
+ output, _ = self.out_proj(core_attn_out)
867
+ return output
868
869
def forward_hip(
870
self,
871
hidden_states: torch.Tensor,
874
- output: torch.Tensor,
875
- ):
872
+ ) -> torch.Tensor:
873
"""ROCm forward using AITER Triton fused projection+attention when
874
available, otherwise falling back to the generic CUDA path."""
875
if GDN_AITER_TRITON_AVAILABLE:
@@ -901,15 +898,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
898
use_aiter=True,
899
)
900
904
- self._output_projection(core_attn_out, z, output, num_tokens)
901
+ return self._output_projection(core_attn_out, z)
902
else:
906
- self.forward_cuda(hidden_states, output)
903
+ return self.forward_cuda(hidden_states)
904
905
def forward_cuda(
906
self,
907
hidden_states: torch.Tensor,
911
- output: torch.Tensor,
912
- ):
908
+ ) -> torch.Tensor:
909
"""
910
Forward pass with three parts:
911
1. Input projection
@@ -964,13 +960,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
960
# ============================================================
961
# Part 3: Output Projection
962
# ============================================================
967
- self._output_projection(core_attn_out, z, output, num_tokens)
963
+ return self._output_projection(core_attn_out, z)
964
965
def forward_xpu(
966
self,
967
hidden_states: torch.Tensor,
972
- output: torch.Tensor,
973
- ):
968
+ ) -> torch.Tensor:
969
"""
970
Forward pass with three parts:
971
1. Input projection
@@ -1013,13 +1008,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1008
core_attn_out = self.norm(core_attn_out, z)
1009
core_attn_out = core_attn_out.reshape(z_shape_og)
1010
core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)
1016
- output[:num_tokens], _ = self.out_proj(core_attn_out)
1011
+ out, _ = self.out_proj(core_attn_out)
1012
+ return out
1013
1014
def forward_cpu(
1015
self,
1016
hidden_states: torch.Tensor,
1021
- output: torch.Tensor,
1022
- ):
1017
+ ) -> torch.Tensor:
1018
assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU."
1019
1020
mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
@@ -1063,7 +1058,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1058
core_attn_out = self.norm(core_attn_out, z)
1059
core_attn_out = core_attn_out.reshape(z_shape_og)
1060
core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)
1066
- output[:num_tokens], _ = self.out_proj(core_attn_out)
1061
+ out, _ = self.out_proj(core_attn_out)
1062
+ return out
1063
1064
def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None:
1065
"""Warm up GDN prefill kernels during V1 profiling.
vllm/model_executor/models/qwen3_next.py
+5
-11
@@ -381,15 +381,15 @@ class Qwen3NextAttention(nn.Module):
381
def forward(
382
self,
383
positions: torch.Tensor,
384
- output: torch.Tensor,
384
hidden_states: torch.Tensor,
386
- ):
385
+ ) -> torch.Tensor:
386
qkv, _ = self.qkv_proj(hidden_states)
387
q, k, v, gate = self._project_qkv_gate(qkv, positions)
388
attn_output = self.attn(q, k, v)
389
if gate is not None:
390
attn_output = attn_output * torch.sigmoid(gate)
392
- output[:], _ = self.o_proj(attn_output)
391
+ output, _ = self.o_proj(attn_output)
392
+ return output
393
394
395
class Qwen3NextDecoderLayer(nn.Module):
@@ -484,21 +484,15 @@ class Qwen3NextDecoderLayer(nn.Module):
484
else:
485
hidden_states, residual = self.input_layernorm(hidden_states, residual)
486
487
- self_attention_output = torch.empty_like(hidden_states)
487
if self.layer_type == "linear_attention":
489
- self.linear_attn(
490
- hidden_states=hidden_states,
491
- output=self_attention_output,
492
- )
488
+ hidden_states = self.linear_attn(hidden_states=hidden_states)
489
elif self.layer_type == "full_attention":
494
- self.self_attn(
490
+ hidden_states = self.self_attn(
491
hidden_states=hidden_states,
496
- output=self_attention_output,
492
positions=positions,
493
)
494
else:
495
raise ValueError("Invalid layer_type")
501
- hidden_states = self_attention_output
496
497
if self.layer_scale:
498
if len(hidden_states.shape) == 2: