A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 11:45 UTC
[Model] [Perf] Use flatten for Qwen3.5's GDN output projection (#42311)
Signed-off-by: Rishapveer Singh <singhrishapveer@gmail.com>
Rishapveer Singh committed
May 18, 2026 at 10:14 UTC
5ab6d1b3fd407404cd78488bf6f4cbcde6d912b7
1 file changed
+3
-3
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+3
-3
@@ -680,7 +680,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
680
z = z.reshape(-1, z.shape[-1])
681
core_attn_out = self.norm(core_attn_out, z)
682
core_attn_out = core_attn_out.reshape(z_shape_og)
683
- core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
683
+ core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)
684
output[:num_tokens], _ = self.out_proj(core_attn_out)
685
686
def forward_hip(
@@ -828,7 +828,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
828
z = z.reshape(-1, z.shape[-1])
829
core_attn_out = self.norm(core_attn_out, z)
830
core_attn_out = core_attn_out.reshape(z_shape_og)
831
- core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
831
+ core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)
832
output[:num_tokens], _ = self.out_proj(core_attn_out)
833
834
def forward_cpu(
@@ -878,7 +878,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
878
z = z.reshape(-1, z.shape[-1])
879
core_attn_out = self.norm(core_attn_out, z)
880
core_attn_out = core_attn_out.reshape(z_shape_og)
881
- core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
881
+ core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d)
882
output[:num_tokens], _ = self.out_proj(core_attn_out)
883
884
def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None: