diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 9c323799a77a..5b97e32a6ee3 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import math from typing import Tuple @@ -214,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 @@ -419,6 +437,15 @@ 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. 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_( 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..e03539d66cde 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, @@ -1150,6 +1151,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 +1499,145 @@ 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`` 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``. + + 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 + + # 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 = [ + (capped_num_tokens, 0, False, False), + (capped_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 = (Mamba2Metadata.force_initial_states_for_warmup() + if force_init_i else contextlib.nullcontext()) + 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) + + if autotuner_enabled: + AutoTuner.get().cache_pp_recv() + AutoTuner.get().cache_pp_send() + AutoTuner.get().clean_pp_flag() + + torch.cuda.synchronize() + 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"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. + self._reset_moe_alltoall_state() + torch.cuda.empty_cache() + + 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.