[None][perf] Close Mamba hybrid warmup gap in autotuner warmup#16177
Open
chenfeiz0326 wants to merge 1 commit into
Open
[None][perf] Close Mamba hybrid warmup gap in autotuner warmup#16177chenfeiz0326 wants to merge 1 commit into
chenfeiz0326 wants to merge 1 commit into
Conversation
Mamba hybrid models (e.g. Nemotron 3 Super 120B) skip _general_warmup
entirely because can_run_general_warmup is False when the KV cache
manager is a MambaHybridCacheManager. _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 iter 3 for
~30 s and inflating P99 E2EL.
Fix:
* Extend _run_autotuner_warmup with a small shape sweep for all
models (pure-ctx-max plus mixed ctx+gen), gated by
TLLM_AUTOTUNE_SWEEP (default on).
* For Mamba hybrid models specifically, add two more warmup passes
that build multi-sequence prefill batches (least_requests=False)
and force HAS_INITSTATES=True via TLLM_MAMBA_WARMUP_FORCE_INITSTATES,
gated by TLLM_MAMBA_MULTISEQ_WARMUP (default on).
* Add the TLLM_MAMBA_WARMUP_FORCE_INITSTATES hook in
Mamba2Metadata.prepare(): when set, override has_initial_states_cpu
to True for context requests so the HAS_INITSTATES=True variants
of _state_passing_fwd_kernel, _chunk_scan_fwd_kernel, and
_chunk_state_varlen_kernel compile during warmup.
Verified on lyris GB200 with
nemotron_3_super_120b_nvfp4-serve-pytorch-float4 (maxbs 512,
maxnt 2048, kv_frac 0.8, in/out 1024, reqs 640, con 128, ep/tp 4):
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%
Run 1 no longer trails Run 2 - the warmup gap is closed.
Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
Contributor
📝 WalkthroughWalkthroughThis PR expands autotuner warmup in the PyTorch model engine to sweep multiple context/generation shape combinations controlled by environment variables (TLLM_AUTOTUNE_SWEEP, TLLM_MAMBA_MULTISEQ_WARMUP), replacing the prior single-shape warmup, and adds a corresponding Mamba2Metadata override (TLLM_MAMBA_WARMUP_FORCE_INITSTATES) to force initial-state flags during warmup passes. ChangesAutotuner Warmup Sweep and Mamba Init-State Override
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ModelEngine as PyTorchModelEngine
participant AutoTuner
participant Mamba2Metadata
ModelEngine->>ModelEngine: build autotune_shapes list
loop each shape
ModelEngine->>ModelEngine: set TLLM_MAMBA_WARMUP_FORCE_INITSTATES if needed
ModelEngine->>ModelEngine: create warmup request
ModelEngine->>Mamba2Metadata: prepare()
Mamba2Metadata->>Mamba2Metadata: override has_initial_states_cpu if env set
ModelEngine->>ModelEngine: run forward pass
ModelEngine->>AutoTuner: cache_pp_recv()
ModelEngine->>AutoTuner: cache_pp_send()
ModelEngine->>AutoTuner: clean_pp_flag()
ModelEngine->>ModelEngine: torch.cuda.synchronize()
end
🚥 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: 2
🤖 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_metadata.py`:
- Around line 423-425: `Mamba2Metadata.prepare()` is reading
`TLLM_MAMBA_WARMUP_FORCE_INITSTATES` globally on every call, which makes the
init-state override affect non-warmup requests too. Update `prepare()` to take
and use an explicit warmup-only flag from metadata, and apply the
`has_initial_states_cpu` override only when that flag is set; keep the change
scoped around the `Mamba2Metadata.prepare` logic that sets
`has_initial_states_cpu[:num_contexts]`.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1272-1275: The mixed warmup shapes added in model_engine.py need
to account for speculative draft tokens before checking the 4/16 thresholds. In
the warmup shape logic that appends the mixed ctx+gen tuples, use the same token
accounting as _create_warmup_request() by subtracting
self.max_total_draft_tokens from the available generation budget or otherwise
ensuring the context budget stays non-negative. Update the guards around the
autotune_shapes.append calls so they only add shapes when the resulting
context/gen split is valid in speculative modes.
🪄 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: ee69f8c8-f247-4525-bedf-dfc3207a8291
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/mamba/mamba2_metadata.pytensorrt_llm/_torch/pyexecutor/model_engine.py
Comment on lines
+423
to
+425
| if os.environ.get( | ||
| "TLLM_MAMBA_WARMUP_FORCE_INITSTATES") == "1": | ||
| self.has_initial_states_cpu[:num_contexts].fill_(True) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep the init-state override scoped to warmup.
Mamba2Metadata.prepare() now honors TLLM_MAMBA_WARMUP_FORCE_INITSTATES on every call. If that variable is set at process scope, real prefill requests with zero cached tokens are treated as having initial states. Prefer passing an explicit warmup-only flag through metadata instead of reading global env here.
🤖 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_metadata.py` around lines 423 - 425,
`Mamba2Metadata.prepare()` is reading `TLLM_MAMBA_WARMUP_FORCE_INITSTATES`
globally on every call, which makes the init-state override affect non-warmup
requests too. Update `prepare()` to take and use an explicit warmup-only flag
from metadata, and apply the `has_initial_states_cpu` override only when that
flag is set; keep the change scoped around the `Mamba2Metadata.prepare` logic
that sets `has_initial_states_cpu[:num_contexts]`.
Comment on lines
+1272
to
+1275
| if curr_max_num_tokens >= 4: | ||
| autotune_shapes.append((curr_max_num_tokens, 2, True, False)) | ||
| if curr_max_num_tokens >= 16: | ||
| autotune_shapes.append((curr_max_num_tokens, 8, True, False)) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Account for speculative tokens before adding mixed warmup shapes.
These guards only check curr_max_num_tokens >= 4/16, but _create_warmup_request() charges each generation request as 1 + self.max_total_draft_tokens. In speculative modes this can make num_gen_tokens > num_tokens_i, so the “mixed ctx+gen” shape has a negative context-token budget and can warm the wrong shape or over-admit generation work.
Suggested fix
+ gen_tokens_per_request = 1 + self.max_total_draft_tokens
if curr_max_num_tokens >= 4:
- autotune_shapes.append((curr_max_num_tokens, 2, True, False))
+ if curr_max_num_tokens > 2 * gen_tokens_per_request:
+ autotune_shapes.append((curr_max_num_tokens, 2, True, False))
if curr_max_num_tokens >= 16:
- autotune_shapes.append((curr_max_num_tokens, 8, True, False))
+ if curr_max_num_tokens > 8 * gen_tokens_per_request:
+ autotune_shapes.append((curr_max_num_tokens, 8, True, False))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if curr_max_num_tokens >= 4: | |
| autotune_shapes.append((curr_max_num_tokens, 2, True, False)) | |
| if curr_max_num_tokens >= 16: | |
| autotune_shapes.append((curr_max_num_tokens, 8, True, False)) | |
| gen_tokens_per_request = 1 + self.max_total_draft_tokens | |
| if curr_max_num_tokens >= 4: | |
| if curr_max_num_tokens > 2 * gen_tokens_per_request: | |
| autotune_shapes.append((curr_max_num_tokens, 2, True, False)) | |
| if curr_max_num_tokens >= 16: | |
| if curr_max_num_tokens > 8 * gen_tokens_per_request: | |
| autotune_shapes.append((curr_max_num_tokens, 8, True, False)) |
🤖 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/pyexecutor/model_engine.py` around lines 1272 - 1275, The
mixed warmup shapes added in model_engine.py need to account for speculative
draft tokens before checking the 4/16 thresholds. In the warmup shape logic that
appends the mixed ctx+gen tuples, use the same token accounting as
_create_warmup_request() by subtracting self.max_total_draft_tokens from the
available generation budget or otherwise ensuring the context budget stays
non-negative. Update the guards around the autotune_shapes.append calls so they
only add shapes when the resulting context/gen split is valid in speculative
modes.
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
1 participant
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
Nemotron 3 Super 120B (and other
MambaHybridCacheManagermodels) trail on their first serve iterations because the warmup path leaves 12 Mamba SSD Triton kernel variants uncompiled. The first mixed chunked-prefill iter with cached tokens (typically iter 3) then stalls ~30 s JIT-compiling them mid-inference, producing a large run-1 throughput deficit and P99 spike.Root cause chain (all three must hold):
_general_warmupis skipped entirely for Mamba hybrid models —can_run_general_warmupisFalsewhenisinstance(kv_cache_manager, MambaHybridCacheManager)._run_autotuner_warmupissues a singleleast_requests=Trueprefill = 1 sequence ofcurr_max_num_tokens.cu_seqlens_to_chunk_indices_offsets_tritontakes itsnum_seqs == 1fast path and_cu_seqlens_triton_kernelnever launches.num_cached_tokens_per_seq = 0, souse_initial_statesisFalseand everyHAS_INITSTATES=Truekernel variant remains uncompiled.The 12 missing kernels:
_chunk_state_varlen_kernel× 5 configs,_state_passing_fwd_kernel× 4 (HAS_INITSTATES×IS_CONT_BATCHED),_chunk_scan_fwd_kernel× 2 (HAS_INITSTATES),_cu_seqlens_triton_kernel× 1.Fix
tensorrt_llm/_torch/pyexecutor/model_engine.py— extend_run_autotuner_warmupwith a small shape sweep for all models (pure-ctx-max plus mixed ctx+gen), gated byTLLM_AUTOTUNE_SWEEP(default on). For Mamba hybrid models specifically, add two more passes that build multi-sequence prefill batches (least_requests=False) and forceHAS_INITSTATES=TrueviaTLLM_MAMBA_WARMUP_FORCE_INITSTATES, gated byTLLM_MAMBA_MULTISEQ_WARMUP(default on).tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py— add theTLLM_MAMBA_WARMUP_FORCE_INITSTATEShook inMamba2Metadata.prepare()so warmup can force theHAS_INITSTATES=Truecode path without changing real-inference behavior.Both env-var gates default on; nothing else changes behavior for non-Mamba models beyond the extra autotune shapes.
Verification
Cold-cache 2-run benchmark on lyris GB200 with
nemotron_3_super_120b_nvfp4-serve-pytorch-float4-maxbs:512-maxnt:2048-kv_frac:0.8-input_output_len:1024,1024-reqs:640-con:128-ep:4-tp:4-gpus:4(job 2319918):Run 1 no longer trails Run 2 — the warmup gap is closed. Run 2 also improves because the same mid-run autotune cache miss was happening sporadically during Run 2 as well, not only at Run 1 iter 3.
Warmup sweep shapes logged by the patched
_run_autotuner_warmupfor this model:[(2048, 0, True, False), (2048, 0, False, False), (2048, 0, False, True)]Test plan
TLLM_AUTOTUNE_SWEEP=0TLLM_MAMBA_MULTISEQ_WARMUP/TLLM_MAMBA_WARMUP_FORCE_INITSTATESwhen unset (default off for the metadata hook when unset)Summary by CodeRabbit
New Features
Bug Fixes