[None][perf] Close Mamba hybrid warmup gap in autotuner warmup by chenfeiz0326 · Pull Request #16177 · NVIDIA/TensorRT-LLM

chenfeiz0326 · 2026-07-09T05:57:19Z

Summary

Fixes 5 unstable perf cases across 4 clusters (GB200 / GB300 / B300) on Nemotron 3 Super 120B and Nemotron Nano 12B v2. For each case, the tables below show run 1 and run 2 within a single Slurm job before and after the fix.

Nemotron 3 Super 120B (and other MambaHybridCacheManager models — including nemotron_nano_12b_v2) trail on their first serve/bench 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_warmup is skipped entirely for Mamba hybrid models — can_run_general_warmup is False when isinstance(kv_cache_manager, MambaHybridCacheManager).
_run_autotuner_warmup issues a single least_requests=True prefill = 1 sequence of curr_max_num_tokens. cu_seqlens_to_chunk_indices_offsets_triton takes its num_seqs == 1 fast path and _cu_seqlens_triton_kernel never launches.
Dummy warmup requests have num_cached_tokens_per_seq = 0, so use_initial_states is False and every HAS_INITSTATES=True kernel 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.

This PR supersedes #15876 (which pre-JIT'd Mamba SSD HAS_INITSTATES=True kernels for nemotron_nano_12b_v2 on B300). The autotuner-warmup approach here fixes both model families and both cluster classes with one code path.

Fix

tensorrt_llm/_torch/pyexecutor/model_engine.py — 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 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).
tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py — add the TLLM_MAMBA_WARMUP_FORCE_INITSTATES hook in Mamba2Metadata.prepare() so warmup can force the HAS_INITSTATES=True code 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 — 5 unstable perf cases across 4 clusters

Each case ran back-to-back reps in a single Slurm job (isolated output dirs / caches per side). "Before" = no-fix baseline; "After" = PR head (8fac3c21c9). Serve-mode tests keep a single trtllm-serve process across reps, so warmup is amortized after run 1 — the "before" side shows the classic run-1 << run-2 gap that the fix closes. Bench-mode tests spawn a fresh worker per rep, so the JIT stall recurs each run and shows up primarily as elevated TTFT + reduced steady-state throughput on both baseline reps; the fix restores both.

Case 1 — GB200 (lyris) `nemotron_3_super_120b_nvfp4-serve reqs:640 con:128 ep:4 tp:4 gpus:4` — serve mode

Metric	Before r1	Before r2	After r1	After r2
Output tok/s	6344.72	8283.56	9725.66	9730.43
Wall (s)	103.29	79.12	67.38	67.35
Mean TPOT (ms)	21.15	16.43	14.13	14.13
P99 TPOT (ms)	53.58	29.78	18.32	18.41
Mean E2EL (ms)	20640	15804	13458	13452
P99 E2EL (ms)	53811	29466	17745	17829

Before-fix within-job: r2 is +31 % faster than r1 (warmup gap). After fix: r2/r1 = +0.05 % (essentially identical — warmup gap closed). After-fix run 1 is also +53 % faster than before-fix run 1 and +17 % faster than before-fix run 2.

Case 2 — B300 (bia) `nemotron_nano_12b_v2-bench-pytorch-streaming-bfloat16 in/out 500/2000 con:250` — bench mode

Metric	Before r1	Before r2	After r1	After r2
Output tok/s	7777.81	7598.08	7893.60	8039.01
TTFT mean (ms)	3590.28	4006.10	1529.06	1195.16
TPOT mean (ms)	27.37	27.92	27.92	27.52
Total latency (ms)	131657	134771	129725	127379
Per-user tps	38.51	37.82	37.81	38.31

Bench mode restarts the worker per rep, so both baseline runs pay the JIT stall equally (r1 ≈ r2 both slow); the fix restores both runs equally too. Mean TTFT drops 3798 → 1362 ms (−64 %), which is where the mamba warmup gap most directly manifests. Aggregate throughput improves modestly (+3 %) because the JIT cost is amortized across a 130 s / 512-request run.

Case 3 — GB300 (lyris) `nemotron_nano_12b_v2-bench-pytorch-bfloat16 in/out 1000/1000 con:250` — bench mode

Metric	Before r1	Before r2	After r1	After r2
Output tok/s	6576.29	6362.10	7450.64	7431.79
Total latency (ms)	77855	80477	68719	68893
Per-user tps	31.03	30.03	34.62	34.53

Bench mode → same-job r1 ≈ r2 on both sides. Baseline avg 6469 → fix avg 7441 tok/s = +15 %; wall time drops ~12 s per rep.

Case 4 — GB200 (lyris) `nemotron_nano_12b_v2-bench-pytorch-bfloat16 in/out 512/512` — bench mode

Metric	Before r1	Before r2	After r1	After r2
Output tok/s	7277.52	7359.45	8587.37	8105.10
Total latency (ms)	36021	35620	30527	32343
Per-user tps	14.88	15.06	17.71	16.66

Bench mode → same-job r1 ≈ r2 on both sides. Baseline avg 7318 → fix avg 8346 tok/s = +14 %.

Case 5 — B300 (bia) `nemotron_3_super_120b_nvfp4-serve reqs:160 con:32 ep:4 tp:4 gpus:4` — serve mode (3 reps per job)

Metric	Before r1	Before r2	Before r3	After r1	After r2	After r3
Output tok/s	1854	2757	2770	3278	3330	3349
Wall (s)	88.4	59.4	59.2	50.0	49.2	48.9
Mean E2EL (ms)	17666	11880	11823	9988	9835	9777
P99 E2EL (ms)	49476	20414	20484	11480	10406	10313
Mean TPOT (ms)	18.25	12.59	12.53	10.74	10.59	10.53

Before-fix within-job: r1 is −33 % slower than the r2/r3 warmed steady state (classic warmup deficit). After fix: r1 ≈ r2 ≈ r3 (max spread 2 %). After-fix run 1 is +77 % over before-fix run 1 and +20 % over the warmed r2/r3 baseline.

Cluster / commit metadata

Lyris (GB200 / GB300): wheel base 4be213c9c (main HEAD at time of test); baseline uses that wheel with no overlay, fix overlays the two patched files.
Bia (B300): wheel base 54cffcd2fb (post-HMAC-revert main HEAD, x86_64 aarch64 respectively).
Warmup sweep shapes logged by the patched _run_autotuner_warmup for these models: [(2048, 0, True, False), (2048, 0, False, False), (2048, 0, False, True)].

Test plan

CI perf regression check for Nemotron 3 Super 120B nvfp4 serve on GB200
CI perf regression check for Nemotron Nano 12B v2 bench on B300 / GB300 / GB200
Autotuner cache size log line unchanged for non-Mamba models when TLLM_AUTOTUNE_SWEEP=0
Non-Mamba models unaffected by TLLM_MAMBA_MULTISEQ_WARMUP / TLLM_MAMBA_WARMUP_FORCE_INITSTATES when unset (default off for the metadata hook when unset)

Summary by CodeRabbit

New Features
- Expanded startup warmup coverage to exercise more request shapes before normal execution, helping the system adapt sooner to common workloads.
- Improved warmup handling for Mamba-based hybrid models so mixed prefill and generation flows are better represented.
Bug Fixes
- Reduced the chance of early autotuning misses during the first requests.
- Improved warmup consistency for multi-sequence runs and parallel execution paths.

coderabbitai · 2026-07-09T06:00:59Z

📝 Walkthrough

Walkthrough

This 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.

Changes

Autotuner Warmup Sweep and Mamba Init-State Override

Layer / File(s)	Summary
Mamba initial-state warmup override `tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py`	Adds `os` import and a conditional override in `Mamba2Metadata.prepare()` forcing `has_initial_states_cpu[:num_contexts]` to `True` when `TLLM_MAMBA_WARMUP_FORCE_INITSTATES="1"`, bypassing the normal computed cache-status flags.
Autotuner warmup shape sweep `tensorrt_llm/_torch/pyexecutor/model_engine.py`	Replaces the single `(curr_max_num_tokens, 0)` warmup call in `_run_autotuner_warmup` with a sweep over multiple `(num_tokens, num_gen_requests, least_requests, force_mamba_initstates)` shapes, gated by `TLLM_AUTOTUNE_SWEEP` and `TLLM_MAMBA_MULTISEQ_WARMUP` (for `MambaHybridCacheManager`), including a local context manager to temporarily set `TLLM_MAMBA_WARMUP_FORCE_INITSTATES`, retained KV-space/TP-rank skip checks, Eagle3 draft state reset per shape, and per-shape PP bookkeeping (`cache_pp_recv`, `cache_pp_send`, `clean_pp_flag`) followed by `torch.cuda.synchronize()`. Docstring expanded accordingly.

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)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise and clearly describes the main change to Mamba hybrid warmup behavior.
Description check	✅ Passed	The description is detailed and covers the problem, fix, benchmarks, and test plan, with only minor template-heading deviations.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

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

📥 Commits

Reviewing files that changed from the base of the PR and between 1bb55a3 and 8fac3c2.

📒 Files selected for processing (2)

tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py
tensorrt_llm/_torch/pyexecutor/model_engine.py

Address CodeRabbit review on PR NVIDIA#16177: Mamba2Metadata.prepare() was reading TLLM_MAMBA_WARMUP_FORCE_INITSTATES from os.environ on every call. A stray shell export, an inherited env in a forked worker, or an exception-skipped context-manager finally would silently force real prefills onto the HAS_INITSTATES=True path with has_initial_states set for cache-miss contexts, corrupting SSD chunk metadata. Replace the env var with a class-scoped context manager on Mamba2Metadata: * Mamba2Metadata gains a private class attribute _warmup_force_initial_states (default False) and a classmethod contextmanager force_initial_states_for_warmup() that flips it for the wrapped scope and restores the previous value in finally. * prepare() reads Mamba2Metadata._warmup_force_initial_states instead of os.environ. * PyTorchModelEngine._run_mamba_hybrid_warmup() drops the local _force_mamba_initstates_env helper and wraps the force-initstates warmup shape with Mamba2Metadata.force_initial_states_for_warmup(). * Drops `import os` in mamba2_metadata.py (no other users). Behavior is identical to the previous env-var handshake for the warmup shapes; only the coupling mechanism changes. Nothing outside those two warmup forward passes can set the flag. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com> Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

chenfeiz0326 · 2026-07-12T15:25:02Z

/bot run --disable-fail-fast --stage-list "*PerfSanity*"

chenfeiz0326 · 2026-07-13T02:11:55Z

/bot run --disable-fail-fast --stage-list "*PerfSanity*"

tensorrt-cicd · 2026-07-13T02:17:26Z

PR_Github #58858 [ run ] triggered by Bot. Commit: ffdcbdf Link to invocation

tensorrt-cicd · 2026-07-13T05:52:27Z

PR_Github #58858 [ run ] completed with state FAILURE. Commit: ffdcbdf
/LLM/main/L0_MergeRequest_PR pipeline #47408 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chenfeiz0326 · 2026-07-13T08:45:18Z

/bot run --disable-fail-fast --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU1-GEN1-NODE1-GPU2-Post-Merge-3,DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-2"

tensorrt-cicd · 2026-07-13T08:51:01Z

PR_Github #58938 [ run ] triggered by Bot. Commit: ffdcbdf Link to invocation

tensorrt-cicd · 2026-07-13T09:25:14Z

PR_Github #58938 [ run ] completed with state SUCCESS. Commit: ffdcbdf
/LLM/main/L0_MergeRequest_PR pipeline #47472 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

Address Wanli Jiang's review on PR NVIDIA#16177: `_run_mamba_hybrid_warmup` runs two extra multi-seq forward passes on top of the standard autotuner warmup, both under the multi-seq / HAS_INITSTATES=True path that allocates more chunk-varlen SSD scratch than the single-seq warmup. Under a tight `--kv_cache_free_gpu_mem_fraction` an OOM here would abort init instead of degrading gracefully. Wrap the per-shape batch-context + forward + sync in a `try / except torch.OutOfMemoryError`, log a warning with the shape tuple, `torch.cuda.empty_cache()`, and continue to the next shape. Matches `_general_warmup_impl` / `_general_warmup_encoder`. No MoE reset needed — Mamba hybrid models covered by this path (Nemotron 3 Super 120B, Nemotron-Nano-12B-v2) are not MoE. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

tensorrt-cicd · 2026-07-19T06:18:31Z

PR_Github #60148 [ run ] triggered by Bot. Commit: b8e22fa Link to invocation

tensorrt-cicd · 2026-07-19T10:56:32Z

PR_Github #60148 [ run ] completed with state FAILURE. Commit: b8e22fa
/LLM/main/L0_MergeRequest_PR pipeline #48526 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chenfeiz0326 · 2026-07-19T11:03:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-19T11:10:06Z

PR_Github #60170 [ run ] triggered by Bot. Commit: 074e7e1 Link to invocation

tensorrt-cicd · 2026-07-19T14:55:25Z

PR_Github #60170 [ run ] completed with state FAILURE. Commit: 074e7e1
/LLM/main/L0_MergeRequest_PR pipeline #48545 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

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>

Address CodeRabbit review on PR NVIDIA#16177: Mamba2Metadata.prepare() was reading TLLM_MAMBA_WARMUP_FORCE_INITSTATES from os.environ on every call. A stray shell export, an inherited env in a forked worker, or an exception-skipped context-manager finally would silently force real prefills onto the HAS_INITSTATES=True path with has_initial_states set for cache-miss contexts, corrupting SSD chunk metadata. Replace the env var with a class-scoped context manager on Mamba2Metadata: * Mamba2Metadata gains a private class attribute _warmup_force_initial_states (default False) and a classmethod contextmanager force_initial_states_for_warmup() that flips it for the wrapped scope and restores the previous value in finally. * prepare() reads Mamba2Metadata._warmup_force_initial_states instead of os.environ. * PyTorchModelEngine._run_mamba_hybrid_warmup() drops the local _force_mamba_initstates_env helper and wraps the force-initstates warmup shape with Mamba2Metadata.force_initial_states_for_warmup(). * Drops `import os` in mamba2_metadata.py (no other users). Behavior is identical to the previous env-var handshake for the warmup shapes; only the coupling mechanism changes. Nothing outside those two warmup forward passes can set the flag. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com> Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Address Wanli Jiang's review on PR NVIDIA#16177: `_run_mamba_hybrid_warmup` runs two extra multi-seq forward passes on top of the standard autotuner warmup, both under the multi-seq / HAS_INITSTATES=True path that allocates more chunk-varlen SSD scratch than the single-seq warmup. Under a tight `--kv_cache_free_gpu_mem_fraction` an OOM here would abort init instead of degrading gracefully. Wrap the per-shape batch-context + forward + sync in a `try / except torch.OutOfMemoryError`, log a warning with the shape tuple, `torch.cuda.empty_cache()`, and continue to the next shape. Matches `_general_warmup_impl` / `_general_warmup_encoder`. No MoE reset needed — Mamba hybrid models covered by this path (Nemotron 3 Super 120B, Nemotron-Nano-12B-v2) are not MoE. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

Mirror _general_warmup_impl's OOM handler: if the mamba hybrid warmup forward OOMs between dispatch() and combine(), the MoE all-to-all state (MoeAlltoAll / NVLinkOneSided) is stuck in ``dispatched`` and the next warmup pass or first real request fails with ``dispatch called twice``. Reset before falling through to the next shape. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

The multi-seq warmup pass added in this PR was crashing DGX_H100 CI tests that combine a MambaHybridCacheManager model (Qwen3.5-4B) with spec decoding on a small KV cache pool (259 blocks total for max_num_tokens=8192). ``_run_autotuner_warmup`` fit because it uses ``least_requests=True`` (1-2 long sequences). ``_run_mamba_hybrid_warmup`` uses ``least_requests=False``, which fans the token budget across batch_size (32) short sequences. When each sequence's length lands on a block boundary AND the target KV cache manager has extra tokens per request (spec decoding paths), ``add_token`` allocates one extra block per sequence that ``_create_warmup_request``'s ``blocks_to_use`` check doesn't account for, tipping allocation over the pool ceiling and crashing with "Can't allocate new blocks for window size N". Two changes, both narrow: 1. Cap the multi-seq token budget at ``min(curr_max_num_tokens, 4096)``. The warmup's job is only to compile ``num_seqs > 1`` + ``HAS_INITSTATES=True`` SSD kernel variants — a modest token count achieves that with far more block headroom. 2. Catch ``RuntimeError`` alongside ``torch.OutOfMemoryError``. C++ KV cache block-allocation failures surface as RuntimeError, not torch OOM. If any shape still doesn't fit, log and skip — the model then JIT-compiles the missing kernel variants lazily on the first real request (pre-fix behavior). Fixes the DGX_H100-PyTorch-6 stage failure on ``accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash``. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

Release-Check flagged the multi-line ``logger.warning(...)`` f-string in ``_run_mamba_hybrid_warmup``'s exception handler. Reflow to the single-anchor continuation style that ``yapf --style '{based_on_style: pep8, column_limit: 80}'`` produces so pre-commit passes. No behavior change. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

chenfeiz0326 · 2026-07-20T02:48:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T02:54:21Z

PR_Github #60227 [ run ] triggered by Bot. Commit: e93b50c Link to invocation

tensorrt-cicd · 2026-07-20T07:00:01Z

PR_Github #60227 [ run ] completed with state FAILURE. Commit: e93b50c
/LLM/main/L0_MergeRequest_PR pipeline #48594 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chenfeiz0326 · 2026-07-20T09:46:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T09:51:43Z

PR_Github #60345 [ run ] triggered by Bot. Commit: e93b50c Link to invocation

tensorrt-cicd · 2026-07-20T13:42:50Z

PR_Github #60345 [ run ] completed with state FAILURE. Commit: e93b50c
/LLM/main/L0_MergeRequest_PR pipeline #48688 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chenfeiz0326 · 2026-07-20T15:00:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T15:06:17Z

PR_Github #60397 [ run ] triggered by Bot. Commit: e93b50c Link to invocation

tensorrt-cicd · 2026-07-20T16:19:15Z

PR_Github #60397 [ run ] completed with state FAILURE. Commit: e93b50c
/LLM/main/L0_MergeRequest_PR pipeline #48736 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chenfeiz0326 · 2026-07-21T02:18:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T02:23:50Z

PR_Github #60533 [ run ] triggered by Bot. Commit: e93b50c Link to invocation

tensorrt-cicd · 2026-07-21T03:45:39Z

PR_Github #60533 [ run ] completed with state SUCCESS. Commit: e93b50c
/LLM/main/L0_MergeRequest_PR pipeline #48854 completed with status: 'SUCCESS'

CI Report

Link to invocation

…p pass _run_mamba_hybrid_warmup (NVIDIA#16177) runs two extra prefill forwards on Mamba hybrid models (Nemotron-H, Qwen3-Next/3.5, Kimi Linear / K3): a batch of many short sequences, then the same batch with every sequence forced to have an initial state (HAS_INITSTATES=True), with chunk alignment splitting them into an aligned and a ragged batch (NVIDIA#18813). The first pass only compiles: multi-sequence variants of the SSD, GDN and KDA kernels and of _cu_seqlens_triton_kernel. MambaSSDProvider, LinearAttnProvider and KdaPrefillProvider compile those variants (and the other chunk alignment) in CPU-only helpers when JIT prefetch is on, so the pass is skipped then. The HAS_INITSTATES pass stays. Its kernels are @triton.autotune'd, and Triton's autotune key includes the dtype of every tensor argument, so initial_states / chunk_indices / chunk_offsets turning from None into tensors make new keys whose configs must be benchmarked on the GPU. A helper can compile the candidates but cannot run that benchmark, and the tuned configs are not seeded or reused, so this pass remains the only way to keep that benchmark out of the first cached-prefix batch. TLLM_WARMUP_JIT_COVERAGE=1 restores the first pass with prefetch on. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>

…serving Prefetch compiled every candidate config of an untuned Triton autotune key but left the benchmark to the first forward that needed the key, so NVIDIA#16177's forced-initial-states Mamba warmup pass was kept only to run that benchmark during warmup. A provider may now tune ahead (tune_batches / run_tune_batch). At the end of each warmup, once the helpers are up, JitPrefetcher.tune() plans each tune batch whose launches reach an untuned key: the variants it needs that are not compiled yet, every candidate config of an untuned key included, go to the helpers at batch priority and are waited for. The batch then runs on the GPU through the module's real launchers on scratch tensors, so each untuned Autotuner.run benchmarks its candidates, as the first forward would, but on helper-built cubins and before serving. An SSD kernel's autotune key depends only on whether some request has a cached prefix (initial states and chunk indices then join the key of _chunk_scan_fwd_kernel, _chunk_state_varlen_kernel and _state_passing_fwd_kernel), not on the batch's size or sequence count. So MambaSSDProvider's two tune batches, the largest context batch the engine admits (capped at 4096 tokens, max_batch_size sequences), without and with cached prefixes, reach every key any batch can produce. Layers on FlashInfer's fused SSD path plan no tune batch. With prefetch tuning ahead, _run_mamba_hybrid_warmup skips the forced-initial-states pass when every recurrent layer of every rank is a Triton-SSD Mamba2 mixer (Nemotron-H); GDN and KDA models keep it, their keys are not tuned ahead. The multi-sequence pass stays: it is the warmup's first multi-request context forward, which loads the attention backend's kernels for such batches (one FlashInfer trtllm-gen context cubin on Nemotron, fetched at the first request otherwise). TLLM_JIT_PREFETCH_TUNE=0 turns tune ahead off. The JIT stats count tune-ahead batches, benchmarks, in-process compiles and waits (ta_* in the summary); a helper that dies stops the wait. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>

chenfeiz0326 requested review from a team as code owners July 9, 2026 05:57

chenfeiz0326 requested review from achartier, shaharmor98 and symphonylyh July 9, 2026 05:57

github-actions Bot assigned chenfeiz0326 Jul 9, 2026

coderabbitai Bot reviewed Jul 9, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py Resolved

Comment thread tensorrt_llm/_torch/pyexecutor/model_engine.py Resolved

chenfeiz0326 requested review from hyukn and longlee0622 July 10, 2026 07:46

chenfeiz0326 mentioned this pull request Jul 12, 2026

[None][perf] Pre-JIT Mamba SSD HAS_INITSTATES=True kernels during warmup #15876

Closed

3 tasks done

chenfeiz0326 force-pushed the chenfeiz/mamba-hybrid-warmup-fix branch from 8fac3c2 to b14af3d Compare July 12, 2026 15:18

chenfeiz0326 force-pushed the chenfeiz/mamba-hybrid-warmup-fix branch from b14af3d to 6d81735 Compare July 12, 2026 15:20

chenfeiz0326 force-pushed the chenfeiz/mamba-hybrid-warmup-fix branch from 6d81735 to ffdcbdf Compare July 12, 2026 15:22

nv-guomingz approved these changes Jul 13, 2026

View reviewed changes

Wanli-Jiang reviewed Jul 13, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/model_engine.py Resolved

Wanli-Jiang approved these changes Jul 13, 2026

View reviewed changes

chenfeiz0326 requested a review from a team as a code owner July 14, 2026 14:45

chenfeiz0326 and others added 6 commits July 20, 2026 10:47

chenfeiz0326 force-pushed the chenfeiz/mamba-hybrid-warmup-fix branch from 074e7e1 to e93b50c Compare July 20, 2026 02:47

achartier approved these changes Jul 20, 2026

View reviewed changes

chenfeiz0326 merged commit 002f099 into NVIDIA:main Jul 21, 2026
7 checks passed

brnguyen2 mentioned this pull request Aug 1, 2026

[TRTLLM-14903][fix] Free partially-allocated warmup dummy KV blocks and count spec extra tokens in warmup block estimates #17162

Merged

2 tasks done

chenfeiz0326 mentioned this pull request Oct 3, 2026

[None][perf] Experimental: compile and tune Triton, DeepGEMM and CuTe DSL JIT variants before serving #19843

Draft

Conversation

chenfeiz0326 commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Fix

Verification — 5 unstable perf cases across 4 clusters

Case 1 — GB200 (lyris) nemotron_3_super_120b_nvfp4-serve reqs:640 con:128 ep:4 tp:4 gpus:4 — serve mode

Case 2 — B300 (bia) nemotron_nano_12b_v2-bench-pytorch-streaming-bfloat16 in/out 500/2000 con:250 — bench mode

Case 3 — GB300 (lyris) nemotron_nano_12b_v2-bench-pytorch-bfloat16 in/out 1000/1000 con:250 — bench mode

Case 4 — GB200 (lyris) nemotron_nano_12b_v2-bench-pytorch-bfloat16 in/out 512/512 — bench mode

Case 5 — B300 (bia) nemotron_3_super_120b_nvfp4-serve reqs:160 con:32 ep:4 tp:4 gpus:4 — serve mode (3 reps per job)

Cluster / commit metadata

Test plan

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

chenfeiz0326 commented Jul 12, 2026

Uh oh!

chenfeiz0326 commented Jul 13, 2026

Uh oh!

tensorrt-cicd commented Jul 13, 2026

Uh oh!

tensorrt-cicd commented Jul 13, 2026

Uh oh!

chenfeiz0326 commented Jul 13, 2026

Uh oh!

tensorrt-cicd commented Jul 13, 2026

Uh oh!

tensorrt-cicd commented Jul 13, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 19, 2026

Uh oh!

tensorrt-cicd commented Jul 19, 2026

Uh oh!

chenfeiz0326 commented Jul 19, 2026

Uh oh!

tensorrt-cicd commented Jul 19, 2026

Uh oh!

tensorrt-cicd commented Jul 19, 2026

Uh oh!

chenfeiz0326 commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

chenfeiz0326 commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

chenfeiz0326 commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

chenfeiz0326 commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

chenfeiz0326 commented Jul 9, 2026 •

edited

Loading

Case 1 — GB200 (lyris) `nemotron_3_super_120b_nvfp4-serve reqs:640 con:128 ep:4 tp:4 gpus:4` — serve mode

Case 2 — B300 (bia) `nemotron_nano_12b_v2-bench-pytorch-streaming-bfloat16 in/out 500/2000 con:250` — bench mode

Case 3 — GB300 (lyris) `nemotron_nano_12b_v2-bench-pytorch-bfloat16 in/out 1000/1000 con:250` — bench mode

Case 4 — GB200 (lyris) `nemotron_nano_12b_v2-bench-pytorch-bfloat16 in/out 512/512` — bench mode

Case 5 — B300 (bia) `nemotron_3_super_120b_nvfp4-serve reqs:160 con:32 ep:4 tp:4 gpus:4` — serve mode (3 reps per job)

coderabbitai Bot commented Jul 9, 2026 •

edited

Loading