[Perf] Warm up hybrid Mamba2 Triton kernels reported by the JIT monitor - #48363
Open
Majid-Taheri wants to merge 1 commit into
Open
Majid-Taheri wants to merge 1 commit into
Majid-Taheri wants to merge 1 commit into
Conversation
On hybrid Mamba2 models (e.g. NemotronH), the JIT monitor reports three Triton kernels compiling during the first inference request, each causing a latency spike on an otherwise fully warmed server: - _causal_conv1d_fwd_kernel: the Mamba2 SSD warmup covers only the SSD chunk kernels, and it runs in the profile pass before the conv cache exists, so the prefill conv kernel cannot be warmed there. - _zero_kv_blocks_kernel: warmed only for Qwen model types (vllm-project#46750). - _compute_slot_mapping_kernel: the generic block-table warmup misses the block_table_stride == 1 specialization that hybrid models hit, because their mamba-aligned attention block size yields one block per request. Add a hybrid_mamba_triton_warmup step, mirroring the Qwen Triton warmup but keyed on the presence of MambaMixer2 layers instead of model_type, so it covers NemotronH, Granite hybrid, Falcon-H1, Codestral Mamba, etc. The conv1d warmup uses the layer's real weights and conv cache with a single dummy token routed to the null block, so the JIT key matches production prefill without touching a real cache line. Signed-off-by: Majid Taheri Andani <tahemaji@amazon.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Kh4L
approved these changes
Jul 15, 2026
LopezCastroRoberto
left a comment
LopezCastroRoberto
left a comment
Contributor
There was a problem hiding this comment.
Hey @Majid-Taheri
We're currently migrating all kernel warmups to a shared warmup contract. See #47451, RFC: #47456.
Would you mind migrating these warmups to conform to that shared contract? It would help keep the warmup infrastructure consistent and make future maintenance easier.
Since you already identified the compile-keys, I think it should be easy for you to do this migration :) Thanks! I add your PR to this list: #49349
43 tasks
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
On hybrid Mamba2 models (NemotronH and other
MambaMixer2models), three Triton kernels compile during the first inference request, not at warmup. Each one causes a first-request latency spike (this is a cold-start / tail-latency fix, not a steady-state throughput change):This is tracked in #43009
Why they are unwarmed today (checked on main @ 429f405):
_causal_conv1d_fwd_kernel:_warmup_ssd_kernelswarms only the SSD chunk kernels, and runs before the conv cache exists._zero_kv_blocks_kernel: warmed only for Qwen, but every hybrid model needs KV-block zeroing._compute_slot_mapping_kernel: generic warmup covers strides (3, 16) but not stride 1, which hybrid models hit.Change
Add
hybrid_mamba_triton_warmup, keyed on the presence ofMambaMixer2layers (covers NemotronH, Granite hybrid, Falcon-H1, Codestral Mamba, Zamba). It reuses the Qwen warmup helpers and adds a conv1d warmup that uses the real conv weights, routes a dummy token to the null block, and sweeps the 8 pointer-alignment cases (a single warmup call covers only one, so mixed decode+prefill batches kept recompiling without the sweep).One new file plus a 1-call hook in
kernel_warmup. No-op for models withoutMambaMixer2layers.Test Result
RTX PRO 6000 Blackwell (sm_120), Nemotron-3-Super-120B-NVFP4,
--jit-monitor-mode=warn. Before: all 3 warnings on the first request (reproduced in 3 cold serves). After: zero warnings across 12 requests (sequential and concurrent mixed batches). Offline replay of production-shaped calls adds no new JIT keys.Precedents and relation to #43642
Precedents: #21955, #22215 (DeepGEMM hot-path warmup), #46750/#47546 (Qwen warmup this extends), #46621 (JIT diagnostics).
#43642 warms the GDN/Qwen bucket. This PR is complementary: it covers the
MambaMixer2bucket that #43642 does not. Happy to coordinate the sharedkernel_warmup.pytouch points with @lesj0610.