From 5498ea4b2277a5bc65d8dc3e42c34e04d1b15ebf Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Wed, 8 Jul 2026 22:54:20 -0700 Subject: [PATCH 1/6] [None][perf] Close Mamba hybrid warmup gap in PyTorch model engine 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 #15876 (same root cause, narrower scope, function-level hook on Mamba2Mixer). Signed-off-by: Chenfei Zhang --- .../_torch/modules/mamba/mamba2_metadata.py | 9 ++ .../_torch/pyexecutor/model_engine.py | 116 ++++++++++++++++++ 2 files changed, 125 insertions(+) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 9c323799a77a..112c26d180a5 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -14,6 +14,7 @@ # limitations under the License. import math +import os from typing import Tuple import torch @@ -419,6 +420,14 @@ def prepare(self, attn_metadata: AttentionMetadata): device='cpu') self.has_initial_states_cpu[:num_contexts].copy_(initial_states_cpu) + # Warmup-only override: force HAS_INITSTATES=True path so the + # HAS_INITSTATES=True variants of _state_passing_fwd_kernel, + # _chunk_scan_fwd_kernel, and _chunk_state_varlen_kernel compile + # during warmup instead of the first real-request iter that hits + # chunked prefill with cached tokens. + if os.environ.get( + "TLLM_MAMBA_WARMUP_FORCE_INITSTATES") == "1": + self.has_initial_states_cpu[:num_contexts].fill_(True) # Mirror CPU staging flags to the CUDA-side buffer asynchronously. self.has_initial_states[:num_contexts].copy_( self.has_initial_states_cpu[:num_contexts], non_blocking=True) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e970063addd7..58e9bc3825ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1150,6 +1150,13 @@ def warmup(self, resource_manager: ResourceManager) -> None: if not is_enc_dec and not self.mapping.has_cp_helix(): self._run_autotuner_warmup(resource_manager) log_mem_snapshot("warmup/after_autotuner") + # Pre-JIT Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels + # for Mamba hybrid models. Runs regardless of enable_autotuner, + # since MambaHybridCacheManager skips _general_warmup and the + # default autotuner shape is single-seq / no-initstates. Safe + # no-op for non-Mamba models. + self._run_mamba_hybrid_warmup(resource_manager) + log_mem_snapshot("warmup/after_mamba_hybrid") # Release the autotuner's exploration-mode intermediates. The # exploration leftovers are pure waste that hide tens of GiB from # non-torch allocators (cuBLAS handle workspace, UCX/NIXL, @@ -1491,6 +1498,115 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): clear_memory_buffers() torch.cuda.empty_cache() + def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): + """Pre-JIT the Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels. + + Mamba hybrid models (e.g. Nemotron 3 Super 120B, Nemotron-Nano-12B-v2) + skip ``_general_warmup`` 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 + compiles the ``num_seqs == 1`` / ``HAS_INITSTATES=False`` variants of + the SSD kernels. The first real serve iteration with chunked prefill + and multiple context requests then triggers autotune of the missing + variants mid-inference, producing a ~30 s stall / large P99 spike. + + This method runs two extra forward passes to compile those variants + during warmup: + + 1. ``least_requests=False`` — splits ``curr_max_num_tokens`` into many + short sequences, forcing the multi-seq path of + ``cu_seqlens_to_chunk_indices_offsets_triton`` and its + ``_cu_seqlens_triton_kernel``. + 2. ``least_requests=False`` with ``TLLM_MAMBA_WARMUP_FORCE_INITSTATES=1`` + — same as (1) plus the ``HAS_INITSTATES=True`` variants of + ``_state_passing_fwd_kernel``, ``_chunk_scan_fwd_kernel``, and + ``_chunk_state_varlen_kernel``. + + Runs regardless of ``enable_autotuner``. Wraps in ``autotune()`` when + the autotuner is enabled so op-level (M,N,K) caches also get primed + for these shapes. Set ``TLLM_MAMBA_MULTISEQ_WARMUP=0`` to disable. + """ + if os.environ.get("TLLM_MAMBA_MULTISEQ_WARMUP", "1") != "1": + return + kv_cache_manager = resource_manager.get_resource_manager( + self.kv_cache_manager_key) + if kv_cache_manager is None or not isinstance( + kv_cache_manager, MambaHybridCacheManager): + return + + token_num_upper_bound = min(self.max_num_tokens, + self.batch_size * (self.max_seq_len - 1)) + curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( + token_num_upper_bound=token_num_upper_bound, + max_num_draft_tokens=self.original_max_draft_len) + if curr_max_num_tokens < 4: + return + + logger.info( + "Running Mamba hybrid warmup (multi-seq + HAS_INITSTATES=True)...") + + @contextlib.contextmanager + def _force_mamba_initstates_env(): + prev = os.environ.get("TLLM_MAMBA_WARMUP_FORCE_INITSTATES") + os.environ["TLLM_MAMBA_WARMUP_FORCE_INITSTATES"] = "1" + try: + yield + finally: + if prev is None: + os.environ.pop( + "TLLM_MAMBA_WARMUP_FORCE_INITSTATES", None) + else: + os.environ["TLLM_MAMBA_WARMUP_FORCE_INITSTATES"] = prev + + # (num_tokens, num_gen_requests, least_requests, force_initstates) + mamba_warmup_shapes = [ + (curr_max_num_tokens, 0, False, False), + (curr_max_num_tokens, 0, False, True), + ] + + autotuner_enabled = self.llm_args.enable_autotuner + cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) + autotune_ctx = (autotune(cache_path=cache_path) + if autotuner_enabled else contextlib.nullcontext()) + + with self.no_cuda_graph(), autotune_ctx: + for (num_tokens_i, num_gen_requests_i, least_req_i, + force_init_i) in mamba_warmup_shapes: + init_ctx = (_force_mamba_initstates_env() + if force_init_i else contextlib.nullcontext()) + with init_ctx: + warmup_request = self._create_warmup_request( + resource_manager, num_tokens_i, num_gen_requests_i, + least_requests=least_req_i) + with self._release_batch_context( + warmup_request, resource_manager) as batch: + if batch is None and self.mapping.tp_size <= 1: + continue + self._assert_all_tp_ranks_have_warmup_batch( + batch, num_tokens_i) + if batch is None: + continue + spec_resource_manager = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if self.is_draft_model and isinstance( + spec_resource_manager, Eagle3ResourceManager): + spec_resource_manager.is_first_draft = True + + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) + + if autotuner_enabled: + AutoTuner.get().cache_pp_recv() + AutoTuner.get().cache_pp_send() + AutoTuner.get().clean_pp_flag() + + torch.cuda.synchronize() + + clear_memory_buffers() + torch.cuda.empty_cache() + def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: """Compute graph_bs → draft_len mapping for dynamic draft length feature. From 35dc1b13dee18f8927ea58793ce10b2efb1a4c2a Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Sun, 12 Jul 2026 08:03:41 -0700 Subject: [PATCH 2/6] [None][perf] Scope Mamba warmup init-states override to a classmethod Address CodeRabbit review on PR #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 Co-Authored-By: Claude Opus 4.7 --- .../_torch/modules/mamba/mamba2_metadata.py | 26 ++++++++++++--- .../_torch/pyexecutor/model_engine.py | 33 +++++++------------ 2 files changed, 34 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 112c26d180a5..5b97e32a6ee3 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import math -import os from typing import Tuple import torch @@ -215,6 +215,23 @@ def cu_seqlens_to_chunk_indices_offsets( class Mamba2Metadata: + # Warmup-only knob: when set via ``force_initial_states_for_warmup``, + # ``prepare()`` forces ``has_initial_states_cpu[:num_contexts]`` to True so + # the ``HAS_INITSTATES=True`` variants of the SSD Triton kernels compile + # during warmup. Class-scoped (not env-var) so it cannot leak into real + # inference from a stray shell export or a forked worker. + _warmup_force_initial_states: bool = False + + @classmethod + @contextlib.contextmanager + def force_initial_states_for_warmup(cls): + prev = cls._warmup_force_initial_states + cls._warmup_force_initial_states = True + try: + yield + finally: + cls._warmup_force_initial_states = prev + def __init__(self, max_batch_size: int, chunk_size: int): self.max_batch_size = max_batch_size self.chunk_size = chunk_size @@ -424,9 +441,10 @@ def prepare(self, attn_metadata: AttentionMetadata): # HAS_INITSTATES=True variants of _state_passing_fwd_kernel, # _chunk_scan_fwd_kernel, and _chunk_state_varlen_kernel compile # during warmup instead of the first real-request iter that hits - # chunked prefill with cached tokens. - if os.environ.get( - "TLLM_MAMBA_WARMUP_FORCE_INITSTATES") == "1": + # chunked prefill with cached tokens. Gate is a class-scoped + # context manager (see ``force_initial_states_for_warmup``) so it + # cannot silently affect real inference. + if Mamba2Metadata._warmup_force_initial_states: self.has_initial_states_cpu[:num_contexts].fill_(True) # Mirror CPU staging flags to the CUDA-side buffer asynchronously. self.has_initial_states[:num_contexts].copy_( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 58e9bc3825ed..05c0d0675ec3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -59,6 +59,7 @@ from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.fused_moe.moe_load_balancer import (MoeLoadBalancer, MoeLoadBalancerIterContext) +from ..modules.mamba.mamba2_metadata import Mamba2Metadata from ..peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager from ..speculative import (SpecMetadata, get_draft_kv_cache_manager, get_num_extra_kv_tokens, get_spec_metadata, @@ -1518,8 +1519,9 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): short sequences, forcing the multi-seq path of ``cu_seqlens_to_chunk_indices_offsets_triton`` and its ``_cu_seqlens_triton_kernel``. - 2. ``least_requests=False`` with ``TLLM_MAMBA_WARMUP_FORCE_INITSTATES=1`` - — same as (1) plus the ``HAS_INITSTATES=True`` variants of + 2. ``least_requests=False`` inside + ``Mamba2Metadata.force_initial_states_for_warmup()`` — same as (1) + plus the ``HAS_INITSTATES=True`` variants of ``_state_passing_fwd_kernel``, ``_chunk_scan_fwd_kernel``, and ``_chunk_state_varlen_kernel``. @@ -1531,8 +1533,8 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): return kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) - if kv_cache_manager is None or not isinstance( - kv_cache_manager, MambaHybridCacheManager): + if kv_cache_manager is None or not isinstance(kv_cache_manager, + MambaHybridCacheManager): return token_num_upper_bound = min(self.max_num_tokens, @@ -1546,19 +1548,6 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): logger.info( "Running Mamba hybrid warmup (multi-seq + HAS_INITSTATES=True)...") - @contextlib.contextmanager - def _force_mamba_initstates_env(): - prev = os.environ.get("TLLM_MAMBA_WARMUP_FORCE_INITSTATES") - os.environ["TLLM_MAMBA_WARMUP_FORCE_INITSTATES"] = "1" - try: - yield - finally: - if prev is None: - os.environ.pop( - "TLLM_MAMBA_WARMUP_FORCE_INITSTATES", None) - else: - os.environ["TLLM_MAMBA_WARMUP_FORCE_INITSTATES"] = prev - # (num_tokens, num_gen_requests, least_requests, force_initstates) mamba_warmup_shapes = [ (curr_max_num_tokens, 0, False, False), @@ -1573,14 +1562,16 @@ def _force_mamba_initstates_env(): with self.no_cuda_graph(), autotune_ctx: for (num_tokens_i, num_gen_requests_i, least_req_i, force_init_i) in mamba_warmup_shapes: - init_ctx = (_force_mamba_initstates_env() + init_ctx = (Mamba2Metadata.force_initial_states_for_warmup() if force_init_i else contextlib.nullcontext()) with init_ctx: warmup_request = self._create_warmup_request( - resource_manager, num_tokens_i, num_gen_requests_i, + resource_manager, + num_tokens_i, + num_gen_requests_i, least_requests=least_req_i) - with self._release_batch_context( - warmup_request, resource_manager) as batch: + with self._release_batch_context(warmup_request, + resource_manager) as batch: if batch is None and self.mapping.tp_size <= 1: continue self._assert_all_tp_ranks_have_warmup_batch( From 2cf17e943d8901d3269769663a75bd751c3c057d Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Tue, 14 Jul 2026 07:44:34 -0700 Subject: [PATCH 3/6] [None][perf] Catch OOM in Mamba hybrid warmup shape loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Wanli Jiang's review on PR #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 Signed-off-by: Chenfei Zhang --- .../_torch/pyexecutor/model_engine.py | 63 +++++++++++-------- 1 file changed, 36 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 05c0d0675ec3..49cec6c1726c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1564,36 +1564,45 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): force_init_i) in mamba_warmup_shapes: init_ctx = (Mamba2Metadata.force_initial_states_for_warmup() if force_init_i else contextlib.nullcontext()) - with init_ctx: - warmup_request = self._create_warmup_request( - resource_manager, - num_tokens_i, - num_gen_requests_i, - least_requests=least_req_i) - with self._release_batch_context(warmup_request, - resource_manager) as batch: - if batch is None and self.mapping.tp_size <= 1: - continue - self._assert_all_tp_ranks_have_warmup_batch( - batch, num_tokens_i) - if batch is None: - continue - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance( - spec_resource_manager, Eagle3ResourceManager): - spec_resource_manager.is_first_draft = True + try: + with init_ctx: + warmup_request = self._create_warmup_request( + resource_manager, + num_tokens_i, + num_gen_requests_i, + least_requests=least_req_i) + with self._release_batch_context( + warmup_request, resource_manager) as batch: + if batch is None and self.mapping.tp_size <= 1: + continue + self._assert_all_tp_ranks_have_warmup_batch( + batch, num_tokens_i) + if batch is None: + continue + spec_resource_manager = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if self.is_draft_model and isinstance( + spec_resource_manager, + Eagle3ResourceManager): + spec_resource_manager.is_first_draft = True - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) - if autotuner_enabled: - AutoTuner.get().cache_pp_recv() - AutoTuner.get().cache_pp_send() - AutoTuner.get().clean_pp_flag() + if autotuner_enabled: + AutoTuner.get().cache_pp_recv() + AutoTuner.get().cache_pp_send() + AutoTuner.get().clean_pp_flag() - torch.cuda.synchronize() + torch.cuda.synchronize() + except torch.OutOfMemoryError: + logger.warning( + f"OOM during Mamba hybrid warmup with " + f"{num_tokens_i} tokens, {num_gen_requests_i} " + f"generation requests, " + f"force_initstates={force_init_i}. Skipping.") + torch.cuda.empty_cache() clear_memory_buffers() torch.cuda.empty_cache() From cdc9cab1eea1b0e628f3bc7c62cd98375414fe16 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Tue, 14 Jul 2026 22:09:18 -0700 Subject: [PATCH 4/6] [None][perf] Reset MoE A2A state on OOM in Mamba hybrid warmup 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 --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 49cec6c1726c..e18276b2cc1a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1602,6 +1602,10 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): f"{num_tokens_i} tokens, {num_gen_requests_i} " f"generation requests, " f"force_initstates={force_init_i}. Skipping.") + # Mirror _general_warmup_impl: an OOM between dispatch() + # and combine() leaves MoE A2A state in ``dispatched``, + # tripping ``dispatch called twice`` on the next forward. + self._reset_moe_alltoall_state() torch.cuda.empty_cache() clear_memory_buffers() From 7e2f3b4d4ba246d4b2b52ea91cd34599da36f0be Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Sat, 18 Jul 2026 22:46:55 -0700 Subject: [PATCH 5/6] [None][perf] Cap Mamba hybrid warmup tokens + catch RuntimeError MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../_torch/pyexecutor/model_engine.py | 42 +++++++++++++++---- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e18276b2cc1a..450ebb395987 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1545,13 +1545,33 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): if curr_max_num_tokens < 4: return + # Cap the multi-seq warmup token count so we don't fill the KV cache + # to the brim. The autotuner warmup that ran just before this uses + # ``least_requests=True`` (few long sequences) which fits comfortably + # even when ``curr_max_num_tokens`` is close to the block ceiling. + # ``least_requests=False`` instead spreads the token budget across + # ``batch_size`` short sequences; when each sequence's length lands + # exactly on a block boundary AND the KV cache has ``num_extra_kv_tokens`` + # or ``num_extra_decoding_steps`` > 0 (e.g. spec decoding cases), + # ``add_token`` needs to allocate one extra block per sequence, which + # ``_create_warmup_request``'s ``blocks_to_use`` estimate doesn't + # account for. On a small KV pool (e.g. Qwen3.5 hybrid with DFlash spec + # decoding on a single H100: 259 blocks total, ``max_num_tokens=8192`` + # nearly saturates it), that extra per-sequence block overflows the + # pool and crashes with "Can't allocate new blocks for window size N". + # The point of this warmup is only to trigger ``num_seqs > 1`` + + # ``HAS_INITSTATES=True`` kernel variants — a modest token budget + # achieves that with plenty of block headroom. + WARMUP_TOKEN_CAP = 4096 + capped_num_tokens = min(curr_max_num_tokens, WARMUP_TOKEN_CAP) + logger.info( "Running Mamba hybrid warmup (multi-seq + HAS_INITSTATES=True)...") # (num_tokens, num_gen_requests, least_requests, force_initstates) mamba_warmup_shapes = [ - (curr_max_num_tokens, 0, False, False), - (curr_max_num_tokens, 0, False, True), + (capped_num_tokens, 0, False, False), + (capped_num_tokens, 0, False, True), ] autotuner_enabled = self.llm_args.enable_autotuner @@ -1596,12 +1616,20 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): AutoTuner.get().clean_pp_flag() torch.cuda.synchronize() - except torch.OutOfMemoryError: + except (torch.OutOfMemoryError, RuntimeError) as e: + # Catch both OOM and RuntimeError. C++ KV cache block + # allocation ("Can't allocate new blocks for window size + # N") surfaces as RuntimeError, not torch.OutOfMemoryError. + # This warmup is a pure perf optimization: if a shape + # doesn't fit for any reason, log and skip; the model then + # JIT-compiles the missing kernel variants lazily on the + # first real request (i.e. the pre-fix behavior). logger.warning( - f"OOM during Mamba hybrid warmup with " - f"{num_tokens_i} tokens, {num_gen_requests_i} " - f"generation requests, " - f"force_initstates={force_init_i}. Skipping.") + f"Mamba hybrid warmup skipped for shape " + f"num_tokens={num_tokens_i}, " + f"num_gen_requests={num_gen_requests_i}, " + f"force_initstates={force_init_i}: " + f"{type(e).__name__}: {e}") # Mirror _general_warmup_impl: an OOM between dispatch() # and combine() leaves MoE A2A state in ``dispatched``, # tripping ``dispatch called twice`` on the next forward. From e93b50cdbdc0e3a2402563c6bfe875702e9f0dac Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Sun, 19 Jul 2026 04:02:37 -0700 Subject: [PATCH 6/6] [None][chore] yapf-format warning log in Mamba hybrid warmup 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 --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 450ebb395987..e03539d66cde 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1624,12 +1624,11 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): # doesn't fit for any reason, log and skip; the model then # JIT-compiles the missing kernel variants lazily on the # first real request (i.e. the pre-fix behavior). - logger.warning( - f"Mamba hybrid warmup skipped for shape " - f"num_tokens={num_tokens_i}, " - f"num_gen_requests={num_gen_requests_i}, " - f"force_initstates={force_init_i}: " - f"{type(e).__name__}: {e}") + logger.warning(f"Mamba hybrid warmup skipped for shape " + f"num_tokens={num_tokens_i}, " + f"num_gen_requests={num_gen_requests_i}, " + f"force_initstates={force_init_i}: " + f"{type(e).__name__}: {e}") # Mirror _general_warmup_impl: an OOM between dispatch() # and combine() leaves MoE A2A state in ``dispatched``, # tripping ``dispatch called twice`` on the next forward.