[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
 
@@ -346,6 +348,8 @@ def forward(
             has_initial_states = mamba_metadata.has_initial_states[:
                                                                    num_prefills]
 
+            has_initial_states_p = has_initial_states[:num_prefills]
+            conv_states[state_indices_p[~has_initial_states_p]].zero_()
             # Fused kernel to avoid expensive .contiguous() call in causal_conv1d_fn.
             xbc_p_t = extract_transpose_xbc_prefill(zxbcdt, num_prefill_tokens,
                                                     self.tp_d_inner,
@@ -509,8 +513,22 @@ 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)
+                    # Both replay and flashinfer read from the cache manager's
+                    # persistent per-slot Philox seed buffer; replay indexes by
+                    # cache_batch_idx, flashinfer reads slot 0 from a (1,)
+                    # view.  In-place add_(1) keeps CUDA-graph replay fresh
+                    # without allocating any new CUDA tensors per forward.
+                    rand_seed = layer_cache.mamba_ssm_rand_seed
+                    assert rand_seed is not None, (
+                        "Mamba SSM stochastic rounding is enabled but the "
+                        "rand_seed buffer was not allocated; check that "
+                        "_util.py passes mamba_ssm_stochastic_rounding=True "
+                        "to the cache manager.")
+                    rand_seed.add_(1)
+                    if use_replay:
+                        philox_kwargs['rand_seed'] = rand_seed
+                    else:
+                        philox_kwargs['rand_seed'] = rand_seed[:1]
                     philox_kwargs['philox_rounds'] = self._philox_rounds
 
                 if use_replay:
@@ -602,10 +620,19 @@ def convert_dt():
                 # Non-MTP decode only runs through flashinfer, no replay path.
                 use_stochastic_rounding = self._stochastic_rounding_for_flashinfer
                 if use_stochastic_rounding:
-                    ssu_kwargs['rand_seed'] = torch.randint(0,
-                                                            2**62, (1, ),
-                                                            device=x_d.device,
-                                                            dtype=torch.int64)
+                    # Fetch the persistent (cache_size,) Philox seed buffer
+                    # from the cache manager and pass slot 0 as a (1,) view to
+                    # flashinfer.  No per-call CUDA tensor allocation; the
+                    # in-place add_(1) is CUDA-graph-friendly.
+                    rand_seed = (attn_metadata.kv_cache_manager.
+                                 get_mamba_ssm_rand_seed())
+                    assert rand_seed is not None, (
+                        "Mamba SSM stochastic rounding is enabled but the "
+                        "rand_seed buffer was not allocated; check that "
+                        "_util.py passes mamba_ssm_stochastic_rounding=True "
+                        "to the cache manager.")
+                    rand_seed.add_(1)
+                    ssu_kwargs['rand_seed'] = rand_seed[:1]
                     ssu_kwargs['philox_rounds'] = self._philox_rounds
 
                 self.selective_state_update_func(

tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py

@@ -497,7 +497,7 @@ def _replay_state_update_kernel(
         # 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 = (
@@ -671,8 +671,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 +746,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/_util.py

@@ -1324,6 +1324,14 @@ def _create_kv_cache_manager(
             logger.info(
                 "Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1")
 
+        # Stochastic-rounding seeds must live on the cache manager (not be
+        # re-created with torch.randint per forward) whenever SR can fire
+        # on the fp16 SSM cache.  This mirrors the predicate the mixer uses
+        # internally (`_stochastic_rounding_for_flashinfer` /
+        # `_stochastic_rounding_for_replay`) so allocation matches consumption.
+        mamba_ssm_stochastic_rounding = (stochastic_rounding
+                                         and mamba_params.mamba_ssm_cache_dtype
+                                         == torch.float16)
         kv_cache_manager = kv_cache_manager_cls(
             # mamba cache parameters
             mamba_params.state_size,
@@ -1353,6 +1361,7 @@ def _create_kv_cache_manager(
             execution_stream=execution_stream,
             model_type="nemotron_hybrid",
             use_replay_state_update=use_replay,
+            mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding,
         )
     elif is_qwen3_hybrid(config):
         if max_beam_width > 1: