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 5, 08:30 UTC
fix(gdn): Align prefill warmup with real prefill path (#39169)
Signed-off-by: Ibrahim Arshad <38925737+ibrahim1023@users.noreply.github.com>
Ibrahim Arshad committed
Apr 10, 2026 at 04:49 UTC
9853a3c15959a9d8f113501f697de6e4e0edeee6
1 file changed
+26
-12
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+26
-12
@@ -702,19 +702,33 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
702
num_v_heads = self.num_v_heads // self.tp_size
703
_, state_dtype = self.get_state_dtype()
704
705
- # All kernels use BT = chunk_size (FLA_CHUNK_SIZE4), so a single pass with
706
- # T = chunk_size is sufficient to populate every autotuner cache.
705
+ # All kernels use BT = chunk_size, so a single pass with T = chunk_size
706
+ # is sufficient to populate every autotuner cache. Mirror the real
707
+ # prefill path here: build q/k/v/g/beta via fused_post_conv_prep and
708
+ # then run chunk_gated_delta_rule with in-kernel L2 norm disabled.
709
T = FLA_CHUNK_SIZE
708
- q = torch.randn(1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype)
709
- k = torch.randn(1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype)
710
- v = torch.randn(1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype)
711
- # NOTE: g and beta must have the same dtypes as during
712
- # inference, so we construct them with the same function
713
- # (fused_gdn_gating). dummy_a and dummy_b are throwaway
714
- # inputs required by that function.
710
+ dummy_mixed_qkv = torch.randn(
711
+ T, mixed_qkv.shape[-1], device=device, dtype=dtype
712
+ )
713
dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)
714
dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)
717
- g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias)
715
+ q, k, v, g, beta = fused_post_conv_prep(
716
+ conv_output=dummy_mixed_qkv,
717
+ a=dummy_a,
718
+ b=dummy_b,
719
+ A_log=self.A_log,
720
+ dt_bias=self.dt_bias,
721
+ num_k_heads=num_k_heads,
722
+ head_k_dim=self.head_k_dim,
723
+ head_v_dim=self.head_v_dim,
724
+ apply_l2norm=True,
725
+ output_g_exp=False,
726
+ )
727
+ q = q.unsqueeze(0)
728
+ k = k.unsqueeze(0)
729
+ v = v.unsqueeze(0)
730
+ g = g.unsqueeze(0)
731
+ beta = beta.unsqueeze(0)
732
state = torch.zeros(
733
1,
734
num_v_heads,
@@ -735,7 +749,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
749
initial_state=state,
750
output_final_state=True,
751
cu_seqlens=cu_seqlens,
738
- use_qk_l2norm_in_kernel=True,
752
+ use_qk_l2norm_in_kernel=False,
753
)
754
except Exception:
755
logger.warning(
@@ -753,7 +767,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
767
self.prefix,
768
)
769
finally:
756
- del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
770
+ del dummy_mixed_qkv, q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
771
772
torch.accelerator.empty_cache()
773