From 5ab6d1b3fd407404cd78488bf6f4cbcde6d912b7 Mon Sep 17 00:00:00 2001 From: Rishapveer Singh Date: Mon, 18 May 2026 10:14:36 +0200 Subject: [PATCH] [Model] [Perf] Use flatten for Qwen3.5's GDN output projection (#42311) Signed-off-by: Rishapveer Singh --- vllm/model_executor/layers/mamba/gdn_linear_attn.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index 011a18ec7b..db55e8ba48 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -680,7 +680,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): z = z.reshape(-1, z.shape[-1]) core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) output[:num_tokens], _ = self.out_proj(core_attn_out) def forward_hip( @@ -828,7 +828,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): z = z.reshape(-1, z.shape[-1]) core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) output[:num_tokens], _ = self.out_proj(core_attn_out) def forward_cpu( @@ -878,7 +878,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): z = z.reshape(-1, z.shape[-1]) core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) output[:num_tokens], _ = self.out_proj(core_attn_out) def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None: