From baefcb288f62f69a07c7ace5ee3a2f85522c17e5 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Thu, 2 Jul 2026 00:59:24 -0700 Subject: [PATCH 1/2] [None][perf] Pre-JIT Mamba SSD HAS_INITSTATES=True kernels during warmup 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 --- .../_torch/modules/mamba/mamba2_mixer.py | 76 +++++++++++++++++++ .../_torch/pyexecutor/model_engine.py | 11 +++ 2 files changed, 87 insertions(+) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 816c6bd1e27e..5cef55f7d6f7 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -41,6 +41,7 @@ fused_split_rearrange_after_conv1d) from .layernorm_gated import RMSNorm as RMSNormGated from .layernorm_gated import fused_gated_rmsnorm_quant_shape_ok +from .mamba2_metadata import cu_seqlens_to_chunk_indices_offsets_triton from .replay_selective_state_update import replay_selective_state_update from .selective_state_update import \ selective_state_update as selective_state_update_native @@ -285,6 +286,81 @@ def _try_attach_nvfp4_scale(self): else: self.norm.is_nvfp4 = False + @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). + try: + weight = self.conv1d.weight + device = weight.device + in_dtype = weight.dtype + state_dtype = self._mamba_ssm_cache_dtype or in_dtype + num_prefills = 2 + seq_per = int(self.chunk_size) * 2 + total = num_prefills * seq_per + x_p = torch.zeros((1, total, self.tp_nheads, self.head_dim), + dtype=in_dtype, + device=device) + dt_p = torch.zeros((1, total, self.tp_nheads), + dtype=in_dtype, + device=device) + B_p = torch.zeros((1, total, self.tp_ngroups, self.d_state), + dtype=in_dtype, + device=device) + C_p = torch.zeros((1, total, self.tp_ngroups, self.d_state), + dtype=in_dtype, + device=device) + initial_states = torch.zeros( + (num_prefills, self.tp_nheads, self.head_dim, self.d_state), + dtype=state_dtype, + device=device) + cu_seqlens = torch.arange(0, + total + 1, + seq_per, + dtype=torch.int32, + device=device) + seq_idx = torch.repeat_interleave( + torch.arange(num_prefills, dtype=torch.int32, device=device), + seq_per).unsqueeze(0) + chunk_indices, chunk_offsets = ( + cu_seqlens_to_chunk_indices_offsets_triton( + cu_seqlens=cu_seqlens, + chunk_size=int(self.chunk_size), + total_seqlens=total)) + out_buf = torch.empty((1, total, self.tp_nheads, self.head_dim), + dtype=in_dtype, + device=device) + mamba_chunk_scan_combined( + x_p, + dt_p, + self.A, + B_p, + C_p, + chunk_size=int(self.chunk_size), + D=self.D, + z=None, + dt_bias=self.dt_bias, + initial_states=initial_states, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + dt_softplus=self.delta_softplus, + dt_limit=(0.0, float("inf")), + cu_seqlens=cu_seqlens, + seq_idx=seq_idx, + return_varlen_states=True, + return_final_states=False, + out=out_buf, + state_dtype=state_dtype, + ) + torch.cuda.synchronize() + except Exception as e: + logger.warning_once( + f"Mamba SSD HAS_INITSTATES=True kernel warmup skipped: {e}", + key="mamba_ssd_initstates_warmup_skipped") + def forward( self, hidden_states: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b1a4f1b49cd6..6e5836250aab 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1092,6 +1092,17 @@ def warmup(self, resource_manager: ResourceManager) -> None: self._general_warmup(resource_manager, warmup_requests_configs) log_mem_snapshot("warmup/after_memory_pool_prepop") + # Pre-JIT Mamba SSD HAS_INITSTATES=True Triton kernels so the first + # prefill that carries a cached prefix does not pay the compile cost + # inline. Runs for all cache-manager kinds (incl. MambaHybridCacheManager), + # since that manager is exactly the one that exposes the cold path. + model_root = getattr(self, "model", None) + if model_root is not None: + for module in model_root.modules(): + hook = getattr(module, "warmup_ssd_initstates_kernels", None) + if callable(hook): + hook() + def _general_warmup(self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]): """ From cc6f7ea97792a13eea7fdce06d4b337682fa2dc1 Mon Sep 17 00:00:00 2001 From: Chenfei Zhang Date: Mon, 6 Jul 2026 02:02:51 -0700 Subject: [PATCH 2/2] Address review: narrow warmup exception, dedupe SSD scan call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- .../_torch/modules/mamba/mamba2_mixer.py | 64 +++++++++++-------- .../_torch/pyexecutor/model_engine.py | 10 ++- 2 files changed, 42 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 5cef55f7d6f7..d27d5608054d 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -286,6 +286,34 @@ def _try_attach_nvfp4_scale(self): else: self.norm.is_nvfp4 = False + def _run_ssd_scan(self, x, dt, B, C, cu_seqlens, seq_idx, chunk_indices, + chunk_offsets, initial_states, out, state_dtype): + # Shared entry point for the SSD Triton kernel used by forward() and + # by the warmup pre-JIT. Kept small so forward() semantics are + # unchanged. + return mamba_chunk_scan_combined( + x, + dt, + self.A, + B, + C, + chunk_size=self.chunk_size, + D=self.D, + z=None, + dt_bias=self.dt_bias, + initial_states=initial_states, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + dt_softplus=self.delta_softplus, + dt_limit=(0.0, float("inf")), + cu_seqlens=cu_seqlens, + seq_idx=seq_idx, + return_varlen_states=True, + return_final_states=False, + out=out, + state_dtype=state_dtype, + ) + @torch.inference_mode() def warmup_ssd_initstates_kernels(self) -> None: # Pre-JIT the mamba_chunk_scan_combined + _state_passing_fwd Triton @@ -299,7 +327,7 @@ def warmup_ssd_initstates_kernels(self) -> None: in_dtype = weight.dtype state_dtype = self._mamba_ssm_cache_dtype or in_dtype num_prefills = 2 - seq_per = int(self.chunk_size) * 2 + seq_per = self.chunk_size * 2 total = num_prefills * seq_per x_p = torch.zeros((1, total, self.tp_nheads, self.head_dim), dtype=in_dtype, @@ -328,35 +356,19 @@ def warmup_ssd_initstates_kernels(self) -> None: chunk_indices, chunk_offsets = ( cu_seqlens_to_chunk_indices_offsets_triton( cu_seqlens=cu_seqlens, - chunk_size=int(self.chunk_size), + chunk_size=self.chunk_size, total_seqlens=total)) out_buf = torch.empty((1, total, self.tp_nheads, self.head_dim), dtype=in_dtype, device=device) - mamba_chunk_scan_combined( - x_p, - dt_p, - self.A, - B_p, - C_p, - chunk_size=int(self.chunk_size), - D=self.D, - z=None, - dt_bias=self.dt_bias, - initial_states=initial_states, - chunk_indices=chunk_indices, - chunk_offsets=chunk_offsets, - dt_softplus=self.delta_softplus, - dt_limit=(0.0, float("inf")), - cu_seqlens=cu_seqlens, - seq_idx=seq_idx, - return_varlen_states=True, - return_final_states=False, - out=out_buf, - state_dtype=state_dtype, - ) - torch.cuda.synchronize() - except Exception as e: + self._run_ssd_scan(x_p, dt_p, B_p, C_p, cu_seqlens, seq_idx, + chunk_indices, chunk_offsets, initial_states, + out_buf, state_dtype) + torch.cuda.current_stream().synchronize() + except torch.cuda.OutOfMemoryError: + # OOM at warmup means we won't fit at inference either — surface it. + raise + except Exception as e: # noqa: BLE001 - best-effort JIT prewarm; must not break inference startup logger.warning_once( f"Mamba SSD HAS_INITSTATES=True kernel warmup skipped: {e}", key="mamba_ssd_initstates_warmup_skipped") diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6e5836250aab..df51dc08f500 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1096,12 +1096,10 @@ def warmup(self, resource_manager: ResourceManager) -> None: # prefill that carries a cached prefix does not pay the compile cost # inline. Runs for all cache-manager kinds (incl. MambaHybridCacheManager), # since that manager is exactly the one that exposes the cold path. - model_root = getattr(self, "model", None) - if model_root is not None: - for module in model_root.modules(): - hook = getattr(module, "warmup_ssd_initstates_kernels", None) - if callable(hook): - hook() + for module in self.model.modules(): + hook = getattr(module, "warmup_ssd_initstates_kernels", None) + if callable(hook): + hook() def _general_warmup(self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]):