[None][fix] Stabilize Mamba replay state update

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

@@ -169,6 +169,8 @@ def __init__(
         # cache manager), so precompute both gate values here.
         sr_base = (self._stochastic_rounding_requested
                    and self._mamba_ssm_cache_dtype == torch.float16)
+        # Keep replay SSM-cache writes on the same stochastic-rounding policy
+        # as flashinfer; the replay kernel masks stale slots before using them.
         self._stochastic_rounding_for_replay = sr_base
         self._stochastic_rounding_for_flashinfer = sr_base and self._use_flashinfer
 
@@ -509,8 +511,16 @@ def convert_dt():
 
                 philox_kwargs = {}
                 if use_stochastic_rounding:
-                    philox_kwargs['rand_seed'] = torch.randint(
-                        0, 2**62, (1, ), device=x_d.device, dtype=torch.int64)
+                    if use_replay:
+                        rand_seed = layer_cache.mamba_ssm_rand_seed
+                        rand_seed.add_(1)
+                        philox_kwargs['rand_seed'] = rand_seed
+                    else:
+                        philox_kwargs['rand_seed'] = torch.randint(
+                            0,
+                            2**62, (1, ),
+                            device=x_d.device,
+                            dtype=torch.int64)
                     philox_kwargs['philox_rounds'] = self._philox_rounds
 
                 if use_replay:

tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py

@@ -180,8 +180,16 @@ def _replay_precompute_kernel(
             dt = dt + dt_bias
         if DT_SOFTPLUS:
             dt = softplus(dt)
+        # softplus guarantees dt >= 0 mathematically; sanitize Inf/NaN that
+        # could leak through if upstream activation drifted catastrophically.
+        dt = tl.where((dt == dt) & (tl.abs(dt) < 1e30), dt, 0.0)
 
         A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32)
+        # SSM stability requires A < 0; sanitize Inf/NaN and clamp at 0 so
+        # dA_cumsum = cumsum(A * dt) stays <= 0 by construction, preventing
+        # exp(dA_cumsum) from overflowing to +Inf downstream.
+        A = tl.where((A == A) & (tl.abs(A) < 1e30), A, 0.0)
+        A = tl.minimum(A, 0.0)
         dA_cumsum = tl.cumsum(A * dt, axis=0)
         decay_vec = tl.exp(dA_cumsum)
 
@@ -227,7 +235,11 @@ def _replay_precompute_kernel(
     )
 
     # Compute raw CB once — shared across all heads in this block
-    raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16))
+    raw_CB = tl.dot(
+        C_all.to(tl.float32),
+        tl.trans(B_all).to(tl.float32),
+        input_precision="ieee",
+    )
 
     # Store B to cache (once per group, only if this block covers the first heads)
     if first_head % nheads_ngroups_ratio == 0:
@@ -270,6 +282,8 @@ def _replay_precompute_kernel(
         # Scale raw_CB with per-head decay and dt
         decay_matrix = tl.exp(dA_cumsum[:, None] - dA_cumsum[None, :])
         CB_scaled = tl.where(valid_mask, raw_CB * decay_matrix * dt[None, :], 0.0)
+        # Filter both NaN and Inf for the same reason as `coeff` above.
+        CB_scaled = tl.where((CB_scaled == CB_scaled) & (tl.abs(CB_scaled) < 1e30), CB_scaled, 0.0)
 
         cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head
         tl.store(
@@ -421,6 +435,10 @@ def _replay_state_update_kernel(
     )
     state_mask = m_mask[:, None] & n_mask[None, :]
     state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32)
+    # Defensive sanitize: a prior step's fp16 store may have saturated to Inf
+    # (or rarely a NaN slipped through).  Reset such elements to zero so the
+    # corruption does not persist forever in this slot's SSM cache.
+    state = tl.where((state == state) & (tl.abs(state) < 65504.0), state, 0.0)
     prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx)
 
     # Phase 1: Replay via tl.dot fast-forward (reads from READ buffer)
@@ -436,6 +454,14 @@ def _replay_state_update_kernel(
     old_dt_all = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to(
         tl.float32
     )
+    # Sanitize Inf/NaN loaded from the dt cache.  A prior step could have
+    # stored junk if upstream activations drifted; without this, `coeff =
+    # exp(dA) * dt` propagates Inf into the state update.
+    old_dt_all = tl.where(
+        (old_dt_all == old_dt_all) & (tl.abs(old_dt_all) < 1e30),
+        old_dt_all,
+        0.0,
+    )
 
     old_dA_cumsum_base = (
         old_dA_cumsum_ptr
@@ -446,25 +472,49 @@ def _replay_state_update_kernel(
     old_dA_cumsum_all = tl.load(
         old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0
     ).to(tl.float32)
+    # dA_cumsum = sum of (A * dt) and A is required to be negative for SSM
+    # stability; sanitize Inf/NaN then clamp to <= 0 to keep
+    # `exp(total_dA_cumsum - old_dA_cumsum_all)` in the safe range (0, 1].
+    # An out-of-range positive value here makes coeff = exp(positive) explode.
+    old_dA_cumsum_all = tl.where(
+        (old_dA_cumsum_all == old_dA_cumsum_all) & (tl.abs(old_dA_cumsum_all) < 1e30),
+        old_dA_cumsum_all,
+        0.0,
+    )
+    old_dA_cumsum_all = tl.minimum(old_dA_cumsum_all, 0.0)
 
     # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction).
     # Clamp to [0, T-1] defensively — out-of-contract PNAT > T would read OOB.
     prev_k_idx = tl.minimum(tl.maximum(prev_num_accepted_tokens - 1, 0), T - 1)
     total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to(
         tl.float32
     )
+    # Same sanitize+clamp as old_dA_cumsum_all.  total_decay = exp(total_dA_cumsum);
+    # without the clamp, a positive Inf here yields total_decay = +Inf which
+    # then poisons state via `state *= total_decay`.
+    total_dA_cumsum = tl.where(
+        (total_dA_cumsum == total_dA_cumsum) & (tl.abs(total_dA_cumsum) < 1e30),
+        total_dA_cumsum,
+        0.0,
+    )
+    total_dA_cumsum = tl.minimum(total_dA_cumsum, 0.0)
 
     # Step 0 invariant: PNAT=0 means `state` is already last step's state (not
     # two back).  coeff is all-zero (offs_t < 0), total_decay is 1.0, so the
     # replay leaves `state` unchanged — cache contents don't matter on step 0.
     coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all
-    coeff = tl.where(offs_t < prev_num_accepted_tokens, coeff, 0.0)
+    pnat_t_mask = offs_t < prev_num_accepted_tokens
+    coeff = tl.where(pnat_t_mask, coeff, 0.0)
+    # Filter both NaN and Inf; `coeff == coeff` alone passes Inf since
+    # Inf == Inf is True.  Inf reaching dB_scaled corrupts the SSM state
+    # via tl.dot and propagates to fp16 storage after SR.
+    coeff = tl.where((coeff == coeff) & (tl.abs(coeff) < 1e30), coeff, 0.0)
 
     # Load old_x: (BLOCK_SIZE_T, BLOCK_SIZE_M) — single-buffered
     old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head
     old_x_all = tl.load(
         old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim,
-        mask=t_mask[:, None] & m_mask[None, :],
+        mask=pnat_t_mask[:, None] & m_mask[None, :],
         other=0.0,
     )
 
@@ -477,7 +527,7 @@ def _replay_state_update_kernel(
     )
     old_B_all = tl.load(
         old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate,
-        mask=t_mask[:, None] & n_mask[None, :],
+        mask=pnat_t_mask[:, None] & n_mask[None, :],
         other=0.0,
     ).to(tl.float32)
 
@@ -489,15 +539,27 @@ def _replay_state_update_kernel(
     state *= total_decay
 
     # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate)
-    state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16))
+    state += tl.dot(
+        tl.trans(old_x_all).to(tl.float32),
+        dB_scaled.to(tl.float32),
+        input_precision="ieee",
+    )
+
+    # Sanitize state before storage to prevent Inf/NaN from persisting in the
+    # fp16 SSM cache and corrupting every subsequent decode step for this slot.
+    # NaN gets reset to zero (uncorrupted but useless on its own), while finite
+    # values are clamped to the fp16 normal range so they survive the cast
+    # below instead of saturating to fp16 Inf.  Zeroing legitimate large finite
+    # values would discard information; clamping preserves it.
+    state = tl.where(state == state, tl.minimum(tl.maximum(state, -65504.0), 65504.0), 0.0)
 
     # Write post-replay state
     if USE_RS_ROUNDING:
         # Stochastic rounding for fp16 state using Philox-4x32 PRNG.
         # Each Philox call produces 4 random ints.  We call randint4x on
         # quarter-sized dstate offsets and join+reshape to get the full
         # (M, dstate) random tensor — 4x fewer PRNG rounds.
-        rand_seed = tl.load(rand_seed_ptr)
+        rand_seed = tl.load(rand_seed_ptr + cache_batch_idx)
         base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head
         offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4)
         rand_offsets_q = (
@@ -573,16 +635,36 @@ def _replay_state_update_kernel(
     )
 
     # init_out = C_all @ state^T * decay_vec
-    init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None]
+    init_out = (
+        tl.dot(
+            C_all.to(tl.float32),
+            tl.trans(state).to(tl.float32),
+            input_precision="ieee",
+        )
+        * decay_vec[:, None]
+    )
 
     # cb_out = CB_scaled @ x_all
-    cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16))
+    cb_out = tl.dot(
+        CB_scaled.to(tl.float32),
+        x_all.to(tl.float32),
+        input_precision="ieee",
+    )
 
     out_all = init_out + cb_out
 
     if HAS_D:
         out_all = out_all + x_all * D[None, :]
 
+    # Sanitize the SSM output before it leaves the kernel.  If any element is
+    # NaN or saturates the downstream bf16/fp16 path it cascades into the next
+    # layer's residual stream, in_proj produces garbage dt, and the model
+    # collapses to `<unk>`.  Reset NaN to zero and clamp finite values to fp16
+    # normal range; the post-RMSNorm path then sees a well-conditioned signal
+    # regardless of whether the state itself drifted toward a degenerate
+    # region during long c=32 generations.
+    out_all = tl.where(out_all == out_all, tl.minimum(tl.maximum(out_all, -65504.0), 65504.0), 0.0)
+
     if HAS_Z:
         for t in range(T):
             z_t = tl.load(
@@ -671,8 +753,11 @@ def replay_selective_state_update(
         z: (batch, T, nheads, dim) optional silu gate.
         dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim).
         state_batch_indices: (batch,) optional cache slot mapping.
-        rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed.
-            When provided, state is stochastically rounded to fp16 on store.
+        rand_seed: optional (cache_size,) int64 CUDA tensor of per-cache-slot
+            Philox PRNG seeds.  The caller bumps this tensor in-place for each
+            replay invocation so CUDA graph replay still gets fresh draws; the
+            kernel indexes it by cache_batch_idx.  When provided, state is
+            stochastically rounded to fp16 on store.
             When None, standard deterministic rounding is used.
         philox_rounds: number of Philox PRNG rounds (default 10).
         launch_with_pdl: enable external PDL (conv1d → precompute chain).
@@ -743,6 +828,19 @@ def replay_selective_state_update(
     assert old_dA_cumsum.shape == (cache_size, 2, nheads, T)
     assert cache_buf_idx.shape == (cache_size,)
     assert prev_num_accepted_tokens.shape == (cache_size,)
+    if rand_seed is not None:
+        assert rand_seed.dtype == torch.int64, (
+            f"rand_seed dtype must be int64, got {rand_seed.dtype}"
+        )
+        assert rand_seed.dim() == 1, (
+            f"rand_seed must be a 1D tensor; got shape {tuple(rand_seed.shape)}"
+        )
+        if rand_seed.shape[0] == 1 and cache_size > 1:
+            rand_seed = rand_seed.expand(cache_size).contiguous()
+        assert rand_seed.shape[0] >= cache_size, (
+            f"rand_seed must have length 1 or >= cache_size ({cache_size}); "
+            f"got shape {tuple(rand_seed.shape)}"
+        )
 
     tie_hdim = (
         A.stride(-1) == 0

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -276,8 +276,9 @@ class SpeculativeState(State):
         - Legacy: caches full intermediate SSM states (intermediate_ssm)
         - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum)
         """
-        _SHARED_FIELDS = frozenset(
-            {"prev_num_accepted_tokens", "cache_buf_idx"})
+        _SHARED_FIELDS = frozenset({
+            "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed"
+        })
 
         intermediate_conv_window: torch.Tensor  # always allocated
 
@@ -289,6 +290,10 @@ class SpeculativeState(State):
         # 0 means temporal saved state is actually the last state, not two back.
         prev_num_accepted_tokens: torch.Tensor | None = None  # (cache,) int — shared across layers
         cache_buf_idx: torch.Tensor | None = None  # (cache,) int32 — shared across layers
+        # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so
+        # CUDA graph replay uses fresh SR draws without allocating RNG tensors.
+        # (cache,) int64 - shared across layers
+        mamba_ssm_rand_seed: torch.Tensor | None = None
         old_x: torch.Tensor | None = None  # (layers, cache, T, nheads, dim)
         old_B: torch.Tensor | None = None  # (layers, cache, 2, T, ngroups, dstate)
         # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit.
@@ -406,6 +411,9 @@ def __init__(
                 spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size,
                                                            dtype=torch.int32,
                                                            device=device)
+                spec_kwargs['mamba_ssm_rand_seed'] = torch.empty(
+                    max_batch_size, dtype=torch.int64,
+                    device=device).random_(1, 2**62)
                 spec_kwargs['old_x'] = torch.zeros(num_local_layers,
                                                    max_batch_size,
                                                    T,
@@ -534,6 +542,9 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]):
             if (isinstance(self.mamba_cache, self.SpeculativeState)
                     and self._use_replay_state_update):
                 self.mamba_cache.prev_num_accepted_tokens[block] = 0
+                if self.mamba_cache.mamba_ssm_rand_seed is not None:
+                    self.mamba_cache.mamba_ssm_rand_seed[block] = torch.randint(
+                        1, 2**62, (1, ), dtype=torch.int64).item()
 
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
         context_ids = [
@@ -641,6 +652,7 @@ def _drop(tensor):
                 prev_num_accepted_tokens=_drop(
                     self.mamba_cache.prev_num_accepted_tokens),
                 cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx),
+                mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed),
                 old_x=_drop(self.mamba_cache.old_x),
                 old_B=_drop(self.mamba_cache.old_B),
                 old_dt=_drop(self.mamba_cache.old_dt),
@@ -1320,6 +1332,7 @@ def shutdown(self):
         self.intermediate_state_indices = None
         self.prev_num_accepted_tokens = None
         self.cache_buf_idx = None
+        self.mamba_ssm_rand_seed = None
         self.old_x = None
         self.old_B = None
         self.old_dt = None
@@ -1392,13 +1405,41 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests):
         # block (prefix-cache hit or block recycled across requests) may carry
         # stale prev_num_accepted_tokens / cache_buf_idx values from a prior
         # owner; the replay kernel reads these on the first decode step.
+        # Deterministically zero the per-slot stripes of all replay buffers so
+        # the kernel's unconditional loads (e.g. total_dA_cumsum at prev_k=0)
+        # never see Inf/NaN from a prior owner; this prevented intermittent
+        # SSM state corruption at higher concurrency (c=32) where slot reuse
+        # is frequent.
         if self._use_replay_state_update and self.prev_num_accepted_tokens is not None:
             num_contexts = len(scheduled_batch.context_requests)
             if num_contexts > 0:
                 ctx_slots = self.cuda_state_indices[:num_contexts].long()
                 self.prev_num_accepted_tokens[ctx_slots] = 0
-                # don't care which half of doulbe-buffer is using
-                # self.cache_buf_idx[ctx_slots] = 0
+                self.cache_buf_idx[ctx_slots] = 0
+                if self.old_x is not None:
+                    self.old_x[:, ctx_slots] = 0
+                if self.old_B is not None:
+                    self.old_B[:, ctx_slots] = 0
+                if self.old_dt is not None:
+                    self.old_dt[:, ctx_slots] = 0
+                if self.old_dA_cumsum is not None:
+                    self.old_dA_cumsum[:, ctx_slots] = 0
+                # Zero the SSM/conv state stripes for fresh context slots so
+                # the replay kernel cannot observe stale state on the first
+                # decode step of a recycled slot when the prefill does not
+                # fully overwrite the cache.
+                if self.all_ssm_states is not None:
+                    self.all_ssm_states[:, ctx_slots] = 0
+                if self.all_conv_states is not None:
+                    self.all_conv_states[:, ctx_slots] = 0
+                if self.mamba_ssm_rand_seed is not None:
+                    self.mamba_ssm_rand_seed[ctx_slots] = torch.randint(
+                        1,
+                        2**62,
+                        (num_contexts, ),
+                        dtype=torch.int64,
+                        device=self.mamba_ssm_rand_seed.device,
+                    )
 
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
         super().prepare_resources(scheduled_batch)
@@ -1532,6 +1573,7 @@ def mamba_layer_cache(
                 spec_kwargs['cache_buf_idx'] = self.cache_buf_idx
                 spec_kwargs['prev_num_accepted_tokens'] = (
                     self.prev_num_accepted_tokens)
+                spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed
             else:
                 spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[
                     layer_offset]
@@ -1718,6 +1760,7 @@ def _setup_replay_buffers(self, spec_config) -> None:
             # are allocated in _setup_replay_buffers after _setup_states().
             self.prev_num_accepted_tokens = None
             self.cache_buf_idx = None
+            self.mamba_ssm_rand_seed = None
             self.old_x = None
             self.old_B = None
             self.old_dt = None
@@ -1770,6 +1813,9 @@ def _setup_replay_buffers(self, spec_config) -> None:
                                          T,
                                          dtype=torch.float32,
                                          device=device)
+        self.mamba_ssm_rand_seed = torch.empty(cache_size,
+                                               dtype=torch.int64,
+                                               device=device).random_(1, 2**62)
 
     @property
     def use_replay_state_update(self) -> bool: