iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Model] [Perf] Use flatten for Qwen3.5's GDN output projection (#42311)

Browse Source 
Signed-off-by: Rishapveer Singh <singhrishapveer@gmail.com>
This commit is contained in:
Rishapveer Singh
2026-05-18 16:14:36 +08:00
committed by GitHub
parent 7d5b033782
commit 5ab6d1b3fd
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+3 -3
View File
@@ -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: