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

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,93 @@ 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
+        # 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 = 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=self.chunk_size,
+                    total_seqlens=total))
+            out_buf = torch.empty((1, total, self.tp_nheads, self.head_dim),
+                                  dtype=in_dtype,
+                                  device=device)
+            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")
+
     def forward(
         self,
         hidden_states: torch.Tensor,

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1092,6 +1092,15 @@ 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.
+        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]]):
         """