[None][perf] Pre-JIT Mamba SSD HAS_INITSTATES=True kernels during warmup - #15876
Closed
chenfeiz0326 wants to merge 2 commits into
Closed
[None][perf] Pre-JIT Mamba SSD HAS_INITSTATES=True kernels during warmup#15876chenfeiz0326 wants to merge 2 commits into
chenfeiz0326 wants to merge 2 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR adds a ChangesSSD Kernel Warmup
Estimated code review effort: 2 (Simple) | ~12 minutes Sequence Diagram(s)sequenceDiagram
participant ModelEngine
participant ModelModules
participant Mamba2Mixer
ModelEngine->>ModelModules: iterate self.model.modules()
ModelModules->>Mamba2Mixer: locate warmup_ssd_initstates_kernels
ModelEngine->>Mamba2Mixer: call warmup_ssd_initstates_kernels()
Mamba2Mixer->>Mamba2Mixer: build synthetic tensors and cu_seqlens
Mamba2Mixer->>Mamba2Mixer: run mamba_chunk_scan_combined (init-states)
Mamba2Mixer-->>ModelEngine: return, or log warning on exception
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py (1)
286-293: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a docstring instead of a plain comment for this externally-invoked hook.
warmup_ssd_initstates_kernelsis called from outside this file (model_engine.pydiscovers it viagetattr(module, "warmup_ssd_initstates_kernels", None)), making it a cross-file interface. As per coding guidelines, "For interfaces used outside a file, prefer docstrings over comments; reserve comments for code within functions or interfaces local to a file." The rationale currently lives in a#comment block rather than a docstring.♻️ Proposed fix
`@torch.inference_mode`() def warmup_ssd_initstates_kernels(self) -> None: - # Pre-JIT the mamba_chunk_scan_combined + _state_passing_fwd Triton - # kernels with HAS_INITSTATES=True. Without this, the first prefill - # step that carries a non-empty cached prefix pays a one-shot - # kernel-compile cost that shows up as a ~20% latency spike in - # bench iteration ~2000 (Nemotron-Nano-12B-v2, bia B300). + """Pre-JIT the mamba_chunk_scan_combined + _state_passing_fwd Triton kernels. + + Runs a synthetic multi-sequence prefill with ``initial_states`` set so the + ``HAS_INITSTATES=True`` kernel variant is compiled ahead of time. Without + this, the first real prefill that carries a non-empty cached prefix pays a + one-shot kernel-compile cost, observed as a latency spike during benchmarking. + + Any failure here is caught and logged; it must never break inference for + models (or configurations) where this warmup path is not applicable. + """ try:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py` around lines 286 - 293, Move the rationale for warmup_ssd_initstates_kernels into a docstring on that method instead of the current inline # comment block. Since model_engine.py calls it via getattr as an արտաքին hook, update the method-level documentation on warmup_ssd_initstates_kernels to explain the pre-JIT behavior and latency reason, and keep any remaining comments only for local implementation details.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py`:
- Around line 293-346: The warmup guard in mamba2_mixer.py uses a broad except
Exception in the kernel warmup block around mamba_chunk_scan_combined, which
trips BLE001 and looks unintentional. Narrow the exception to the specific
failures you expect from the warmup path, or if this is intentionally
best-effort like _reset_moe_alltoall_state, add the same explicit lint
suppression and a short justification comment near the try/except. Keep the
existing logger.warning_once behavior and use the mamba_chunk_scan_combined /
logger.warning_once block as the location to update.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py`:
- Around line 286-293: Move the rationale for warmup_ssd_initstates_kernels into
a docstring on that method instead of the current inline # comment block. Since
model_engine.py calls it via getattr as an արտաքին hook, update the method-level
documentation on warmup_ssd_initstates_kernels to explain the pre-JIT behavior
and latency reason, and keep any remaining comments only for local
implementation details.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f286de64-7c2b-474e-9205-2eec9f162b43
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/mamba/mamba2_mixer.pytensorrt_llm/_torch/pyexecutor/model_engine.py
chenfeiz0326
force-pushed
the
chenfeiz/mamba-ssd-initstates-warmup
branch
from
July 2, 2026 10:00
061c3c2 to
c288787
Compare
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "*PerfSanity*" |
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "*PerfSanity*" |
Collaborator
|
PR_Github #57308 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57311 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57308 [ run ] completed with state |
Collaborator
|
PR_Github #57311 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "*PerfSanity*" |
Collaborator
|
PR_Github #57419 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57419 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-6,GB300-36_GPUs-9_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU2-GEN1-NODE8-GPU32-Post-Merge-1,DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-3" |
|
|
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-6,GB300-36_GPUs-9_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU2-GEN1-NODE8-GPU32-Post-Merge-1,DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-3" |
Collaborator
|
PR_Github #57615 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57615 [ run ] completed with state
|
chenfeiz0326
force-pushed
the
chenfeiz/mamba-ssd-initstates-warmup
branch
from
July 5, 2026 13:01
c288787 to
6c4591d
Compare
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "*PerfSanity*" |
Collaborator
|
PR_Github #57620 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57620 [ run ] completed with state
|
achartier
reviewed
Jul 5, 2026
chenfeiz0326
force-pushed
the
chenfeiz/mamba-ssd-initstates-warmup
branch
from
July 6, 2026 14:40
b8da7f8 to
4e56c23
Compare
chenfeiz0326
force-pushed
the
chenfeiz/mamba-ssd-initstates-warmup
branch
from
July 6, 2026 14:47
4e56c23 to
31e33ec
Compare
achartier
approved these changes
Jul 6, 2026
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "*PerfSanity*" |
Collaborator
|
PR_Github #58179 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58179 [ run ] completed with state
|
The first prefill iteration that carries a non-empty cached prefix triggers a one-shot Triton JIT compile of the `mamba_chunk_scan_combined` + `_state_passing_fwd` kernels with `HAS_INITSTATES=True`. This shows up as a ~20% mid-run latency spike on Nemotron-Nano-12B-v2 bench-pytorch, producing an unstable 3-rep total_token_throughput (rep1 cold ~7481, rep2/3 warm ~9528) with CV ~14%. Add a `warmup_ssd_initstates_kernels` hook on `Mamba2Mixer` that runs one dummy `mamba_chunk_scan_combined` call with `initial_states != None` during `PyTorchModelEngine._warmup`, so the JIT cost is paid before the measurement window. Verified on bia B300 (Nemotron-Nano-12B-v2, maxbs:512 / maxnt:2048 / isl,osl=500,2000 / con:250, bfloat16): - Before: 3-rep = [7481, 9472, 9612] t/s, CV 13.9%, median 9472 - After: 3-rep = [9580, 9788, 9794] t/s, CV 1.25%, median 9788 (+3.6%) Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
- Let torch.cuda.OutOfMemoryError propagate from warmup_ssd_initstates_kernels (per achartier: OOM must be fatal). Keep best-effort catch for other exceptions, with # noqa: BLE001 + rationale (per CodeRabbit). - Drop redundant int(self.chunk_size) casts (chunk_size is already int). - Extract _run_ssd_scan helper for the mamba_chunk_scan_combined call so the warmup path shares the entry point with forward() (light refactor). - Use current_stream().synchronize() instead of the device-wide synchronize(). - In PyTorchModelEngine.warmup, drop the getattr(self, "model", None) guard — self.model is always set by warmup time. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
chenfeiz0326
force-pushed
the
chenfeiz/mamba-ssd-initstates-warmup
branch
from
July 8, 2026 13:47
31e33ec to
cc6f7ea
Compare
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-3,GB300-36_GPUs-9_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU2-GEN1-NODE8-GPU32-Post-Merge-1,GB300-8_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE1-GPU4-Post-Merge-3" |
Collaborator
|
PR_Github #58240 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58240 [ run ] completed with state
|
4 tasks
Collaborator
Author
chenfeiz0326
added a commit
to chenfeiz0326/TensorRT-LLM
that referenced
this pull request
Jul 12, 2026
Mamba hybrid models (e.g. Nemotron 3 Super 120B, Nemotron-Nano-12B-v2)
skip _general_warmup entirely because can_run_general_warmup is False
when the KV cache manager is a MambaHybridCacheManager. The default
_run_autotuner_warmup then issues a single least_requests=True prefill
= 1 sequence with num_cached_tokens_per_seq=0, which only exercises
the num_seqs==1 / HAS_INITSTATES=False / IS_CONT_BATCHED=False signature
of the Mamba SSD Triton kernels.
The first real serve iteration with chunked prefill and multiple
context requests then triggers autotune of the 12 missing kernel
variants (_chunk_state_varlen_kernel x5 configs,
_state_passing_fwd_kernel x4 across HAS_INITSTATES x IS_CONT_BATCHED,
_chunk_scan_fwd_kernel x2 across HAS_INITSTATES,
_cu_seqlens_triton_kernel x1) mid-inference, stalling for
~30 s and inflating P99 E2EL.
Fix (two changes):
* Add a warmup-only env-var hook TLLM_MAMBA_WARMUP_FORCE_INITSTATES in
Mamba2Metadata.prepare(): when set, override has_initial_states_cpu
to True for context requests so the HAS_INITSTATES=True kernel
variants compile during warmup instead of at first real request.
* Add PyTorchModelEngine._run_mamba_hybrid_warmup(), called from
warmup() after _run_autotuner_warmup. Runs two extra forward passes
for Mamba hybrid models only: (1) least_requests=False for multi-seq
path (compiles _cu_seqlens_triton_kernel and the multi-seq varlen
SSD kernels); (2) same as (1) plus TLLM_MAMBA_WARMUP_FORCE_INITSTATES
set (compiles HAS_INITSTATES=True variants). Fires regardless of
enable_autotuner. Wraps in autotune() when the autotuner is enabled
so op-level (M,N,K) caches also get primed. Set
TLLM_MAMBA_MULTISEQ_WARMUP=0 to disable. Non-Mamba models get a
free early return.
Verified:
Nemotron-3-Super-120B on GB200
(nvfp4-serve-pytorch, maxbs 512, maxnt 2048, kv_frac 0.8, in/out 1024,
reqs 640, con 128, ep/tp 4), cold 2-run:
before after delta
Run 1 tok/s 8142 9604 +18%
Run 1 P99 E2EL 29980 ms 18704 ms -38%
Run 1 P99 TPOT 30.28 ms 19.26 ms -36%
Run 2 tok/s 8181 9735 +19%
Run 2 P99 E2EL 29650 ms 17788 ms -40%
Nemotron-Nano-12B-v2 on B300
(bench-pytorch-streaming-bfloat16, maxbs 512, maxnt 2048,
in/out 500,2000, con 250), 3-rep:
before after
rep1 tok/s 7481 (to be filled after re-verification)
rep2 tok/s 9472
rep3 tok/s 9612
CV 13.9%
Supersedes PR NVIDIA#15876 (same root cause, narrower scope, function-level
hook on Mamba2Mixer).
Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
chenfeiz0326
added a commit
to chenfeiz0326/TensorRT-LLM
that referenced
this pull request
Jul 18, 2026
Mamba hybrid models (e.g. Nemotron 3 Super 120B, Nemotron-Nano-12B-v2)
skip _general_warmup entirely because can_run_general_warmup is False
when the KV cache manager is a MambaHybridCacheManager. The default
_run_autotuner_warmup then issues a single least_requests=True prefill
= 1 sequence with num_cached_tokens_per_seq=0, which only exercises
the num_seqs==1 / HAS_INITSTATES=False / IS_CONT_BATCHED=False signature
of the Mamba SSD Triton kernels.
The first real serve iteration with chunked prefill and multiple
context requests then triggers autotune of the 12 missing kernel
variants (_chunk_state_varlen_kernel x5 configs,
_state_passing_fwd_kernel x4 across HAS_INITSTATES x IS_CONT_BATCHED,
_chunk_scan_fwd_kernel x2 across HAS_INITSTATES,
_cu_seqlens_triton_kernel x1) mid-inference, stalling for
~30 s and inflating P99 E2EL.
Fix (two changes):
* Add a warmup-only env-var hook TLLM_MAMBA_WARMUP_FORCE_INITSTATES in
Mamba2Metadata.prepare(): when set, override has_initial_states_cpu
to True for context requests so the HAS_INITSTATES=True kernel
variants compile during warmup instead of at first real request.
* Add PyTorchModelEngine._run_mamba_hybrid_warmup(), called from
warmup() after _run_autotuner_warmup. Runs two extra forward passes
for Mamba hybrid models only: (1) least_requests=False for multi-seq
path (compiles _cu_seqlens_triton_kernel and the multi-seq varlen
SSD kernels); (2) same as (1) plus TLLM_MAMBA_WARMUP_FORCE_INITSTATES
set (compiles HAS_INITSTATES=True variants). Fires regardless of
enable_autotuner. Wraps in autotune() when the autotuner is enabled
so op-level (M,N,K) caches also get primed. Set
TLLM_MAMBA_MULTISEQ_WARMUP=0 to disable. Non-Mamba models get a
free early return.
Verified:
Nemotron-3-Super-120B on GB200
(nvfp4-serve-pytorch, maxbs 512, maxnt 2048, kv_frac 0.8, in/out 1024,
reqs 640, con 128, ep/tp 4), cold 2-run:
before after delta
Run 1 tok/s 8142 9604 +18%
Run 1 P99 E2EL 29980 ms 18704 ms -38%
Run 1 P99 TPOT 30.28 ms 19.26 ms -36%
Run 2 tok/s 8181 9735 +19%
Run 2 P99 E2EL 29650 ms 17788 ms -40%
Nemotron-Nano-12B-v2 on B300
(bench-pytorch-streaming-bfloat16, maxbs 512, maxnt 2048,
in/out 500,2000, con 250), 3-rep:
before after
rep1 tok/s 7481 (to be filled after re-verification)
rep2 tok/s 9472
rep3 tok/s 9612
CV 13.9%
Supersedes PR NVIDIA#15876 (same root cause, narrower scope, function-level
hook on Mamba2Mixer).
Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
chenfeiz0326
added a commit
to chenfeiz0326/TensorRT-LLM
that referenced
this pull request
Jul 20, 2026
Mamba hybrid models (e.g. Nemotron 3 Super 120B, Nemotron-Nano-12B-v2)
skip _general_warmup entirely because can_run_general_warmup is False
when the KV cache manager is a MambaHybridCacheManager. The default
_run_autotuner_warmup then issues a single least_requests=True prefill
= 1 sequence with num_cached_tokens_per_seq=0, which only exercises
the num_seqs==1 / HAS_INITSTATES=False / IS_CONT_BATCHED=False signature
of the Mamba SSD Triton kernels.
The first real serve iteration with chunked prefill and multiple
context requests then triggers autotune of the 12 missing kernel
variants (_chunk_state_varlen_kernel x5 configs,
_state_passing_fwd_kernel x4 across HAS_INITSTATES x IS_CONT_BATCHED,
_chunk_scan_fwd_kernel x2 across HAS_INITSTATES,
_cu_seqlens_triton_kernel x1) mid-inference, stalling for
~30 s and inflating P99 E2EL.
Fix (two changes):
* Add a warmup-only env-var hook TLLM_MAMBA_WARMUP_FORCE_INITSTATES in
Mamba2Metadata.prepare(): when set, override has_initial_states_cpu
to True for context requests so the HAS_INITSTATES=True kernel
variants compile during warmup instead of at first real request.
* Add PyTorchModelEngine._run_mamba_hybrid_warmup(), called from
warmup() after _run_autotuner_warmup. Runs two extra forward passes
for Mamba hybrid models only: (1) least_requests=False for multi-seq
path (compiles _cu_seqlens_triton_kernel and the multi-seq varlen
SSD kernels); (2) same as (1) plus TLLM_MAMBA_WARMUP_FORCE_INITSTATES
set (compiles HAS_INITSTATES=True variants). Fires regardless of
enable_autotuner. Wraps in autotune() when the autotuner is enabled
so op-level (M,N,K) caches also get primed. Set
TLLM_MAMBA_MULTISEQ_WARMUP=0 to disable. Non-Mamba models get a
free early return.
Verified:
Nemotron-3-Super-120B on GB200
(nvfp4-serve-pytorch, maxbs 512, maxnt 2048, kv_frac 0.8, in/out 1024,
reqs 640, con 128, ep/tp 4), cold 2-run:
before after delta
Run 1 tok/s 8142 9604 +18%
Run 1 P99 E2EL 29980 ms 18704 ms -38%
Run 1 P99 TPOT 30.28 ms 19.26 ms -36%
Run 2 tok/s 8181 9735 +19%
Run 2 P99 E2EL 29650 ms 17788 ms -40%
Nemotron-Nano-12B-v2 on B300
(bench-pytorch-streaming-bfloat16, maxbs 512, maxnt 2048,
in/out 500,2000, con 250), 3-rep:
before after
rep1 tok/s 7481 (to be filled after re-verification)
rep2 tok/s 9472
rep3 tok/s 9612
CV 13.9%
Supersedes PR NVIDIA#15876 (same root cause, narrower scope, function-level
hook on Mamba2Mixer).
Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
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.
Summary
mamba_chunk_scan_combined+_state_passing_fwdTriton kernels withHAS_INITSTATES=TrueduringPyTorchModelEngine._warmup, so the first prefill that carries a cached prefix does not pay the compile cost inline.The warmup hook is opt-in per module (
hasattr(module, "warmup_ssd_initstates_kernels")) and wrapped in a try/except that only logs — non-Mamba models are unaffected.Verification
On bia B300,
nemotron_nano_12b_v2-bench-pytorch-streaming-bfloat16-maxbs:512-maxnt:2048-input_output_len:500,2000-con:250:_state_passing_fwd_kernelat 0.153 ms across all 28 prefill invocations (no first-call compile spike).Test plan
hasattron eachnn.Module).Summary by CodeRabbit
New Features
Bug Fixes