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 6, 20:30 UTC
[Bugfix] Warm up Triton autotuner for GDN layers during V1 profiling (#36599)
Signed-off-by: AuYang <459461160@qq.com>
Xu Jinyang committed
Mar 12, 2026 at 15:51 UTC
3e64fe4a183aae43c039c9467fe2be49c68389fa
1 file changed
+98
-1
vllm/model_executor/models/qwen3_next.py
+98
-1
@@ -645,6 +645,101 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
645
core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
646
output[:num_tokens], _ = self.out_proj(core_attn_out)
647
648
+ def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
649
+ """Warm up GDN prefill kernels during V1 profiling.
650
+
651
+ During V1 profile runs, ``_forward_core`` returns early because
652
+ ``attn_metadata`` is ``None``, so the autotuned kernels used by
653
+ ``chunk_gated_delta_rule`` (e.g. ``solve_tril``,
654
+ ``chunk_scaled_dot_kkt``) are never invoked. After profiling,
655
+ vLLM allocates KV cache using most of the remaining GPU memory.
656
+ When the first real inference triggers the autotuner it OOMs
657
+ because there is not enough memory left for benchmarking.
658
+
659
+ This method runs minimal forward passes through
660
+ ``chunk_gated_delta_rule`` with small dummy tensors to force
661
+ autotuning while GPU memory is still plentiful. The autotuner
662
+ results are cached globally, so only the first layer incurs
663
+ actual benchmarking cost.
664
+
665
+ Most kernels use a fixed ``BT = chunk_size`` (64), but
666
+ ``chunk_fwd_kernel_o`` recomputes ``BT`` from the sequence
667
+ length: ``min(64, max(16, next_power_of_2(T)))``. Since ``BT``
668
+ is part of its autotune key, we run warmup passes with T = 16,
669
+ 32, and 64 to cover all possible ``BT`` values.
670
+
671
+ The decode path uses ``fused_sigmoid_gating_delta_rule_update``
672
+ which has fixed kernel parameters (no autotuning), so only the
673
+ prefill (chunked) path needs warming up.
674
+ """
675
+ if hasattr(self, "_prefill_kernels_warmed_up"):
676
+ return
677
+ self._prefill_kernels_warmed_up = True
678
+
679
+ device = mixed_qkv.device
680
+ dtype = mixed_qkv.dtype
681
+ num_k_heads = self.num_k_heads // self.tp_size
682
+ num_v_heads = self.num_v_heads // self.tp_size
683
+ _, state_dtype = self.get_state_dtype()
684
+
685
+ # Run warmup for each possible BT value of chunk_fwd_kernel_o:
686
+ # T=16 → BT=16, T=32 → BT=32, T=64 → BT=64.
687
+ # Other kernels always use BT=chunk_size(64), so their autotune
688
+ # cache is populated on the first pass and reused thereafter.
689
+ for T in (16, 32, 64):
690
+ q = torch.randn(
691
+ 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
692
+ )
693
+ k = torch.randn(
694
+ 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
695
+ )
696
+ v = torch.randn(
697
+ 1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype
698
+ )
699
+ g = torch.randn(1, T, num_v_heads, device=device, dtype=dtype)
700
+ beta = torch.randn(1, T, num_v_heads, device=device, dtype=dtype)
701
+ state = torch.zeros(
702
+ 1,
703
+ num_v_heads,
704
+ self.head_v_dim,
705
+ self.head_k_dim,
706
+ device=device,
707
+ dtype=state_dtype,
708
+ )
709
+ cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.long)
710
+
711
+ try:
712
+ self.chunk_gated_delta_rule(
713
+ q=q,
714
+ k=k,
715
+ v=v,
716
+ g=g,
717
+ beta=beta,
718
+ initial_state=state,
719
+ output_final_state=False,
720
+ cu_seqlens=cu_seqlens,
721
+ use_qk_l2norm_in_kernel=True,
722
+ )
723
+ except Exception:
724
+ logger.warning(
725
+ "GDN prefill kernel warmup (T=%d) failed for "
726
+ "layer %s. First inference may OOM due to "
727
+ "autotuner.",
728
+ T,
729
+ self.prefix,
730
+ exc_info=True,
731
+ )
732
+ else:
733
+ logger.debug(
734
+ "GDN prefill kernel warmup (T=%d) completed for layer %s",
735
+ T,
736
+ self.prefix,
737
+ )
738
+ finally:
739
+ del q, k, v, g, beta, state, cu_seqlens
740
+
741
+ torch.accelerator.empty_cache()
742
+
743
def _forward_core(
744
self,
745
mixed_qkv: torch.Tensor,
@@ -659,7 +754,9 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
754
attn_metadata: AttentionMetadata = forward_context.attn_metadata
755
756
if attn_metadata is None:
662
- # V1 profile run
757
+ # V1 profile run — warm up prefill kernels so that
758
+ # autotuning completes before KV cache allocation.
759
+ self._warmup_prefill_kernels(mixed_qkv)
760
return
761
762
assert isinstance(attn_metadata, dict)