From c7f32495348ec93ecd7bbf0462469e26ff760172 Mon Sep 17 00:00:00 2001 From: qgai Date: Sun, 24 May 2026 19:06:17 -0700 Subject: [PATCH 1/5] [None][fix] Stabilize Mamba replay state update Signed-off-by: qgai --- .../_torch/modules/mamba/mamba2_mixer.py | 14 ++- .../mamba/replay_selective_state_update.py | 118 ++++++++++++++++-- .../_torch/pyexecutor/mamba_cache_manager.py | 54 +++++++- 3 files changed, 170 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 1474f8f1fe2b..168799be5ee9 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/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: diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index c8c6a655929b..714be639cde4 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/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,6 +472,16 @@ 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. @@ -453,18 +489,32 @@ def _replay_state_update_kernel( 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,7 +539,19 @@ 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: @@ -497,7 +559,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 = ( @@ -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 ``. 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 diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 8fbf40599d07..372137c2ca25 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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: From 9663fb0bf907f09e4ecb44422f9894b6084f07f0 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 25 May 2026 21:53:45 -0700 Subject: [PATCH 2/5] [None][fix] Stabilize Mamba SSM SR seeds Signed-off-by: qgai --- .../_torch/modules/mamba/mamba2_mixer.py | 37 ++- tensorrt_llm/_torch/pyexecutor/_util.py | 9 + .../_torch/pyexecutor/mamba_cache_manager.py | 263 +++++++++++++++--- .../modules/mamba/test_mamba_ssm_rand_seed.py | 244 ++++++++++++++++ 4 files changed, 500 insertions(+), 53 deletions(-) create mode 100644 tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 168799be5ee9..2ea58ad3c725 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -511,16 +511,22 @@ def convert_dt(): philox_kwargs = {} if use_stochastic_rounding: + # 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: - 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['rand_seed'] = rand_seed[:1] philox_kwargs['philox_rounds'] = self._philox_rounds if use_replay: @@ -612,10 +618,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( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 82f6c137f3e4..5fdf00c1492d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 372137c2ca25..3087749d0797 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -55,6 +55,79 @@ def get_tensor_size_bytes(tensor): return 0 +# Mamba SSM stochastic-rounding Philox seed plumbing. +# +# Both the replay kernel and flashinfer's selective_state_update consume a +# `rand_seed` int64 tensor. Historically the non-replay paths created this +# tensor via `torch.randint(..., (1,))` on every decode step, which (a) is +# non-deterministic across runs, (b) allocates a fresh CUDA tensor per call +# and is therefore unfriendly to CUDA-graph capture, and (c) has no notion +# of cache-slot identity (so slot reuse cannot rotate the stream). +# +# The functions below produce per-slot int64 seeds via SplitMix64 finalization +# of (counter, slot, rank). Adjacent inputs yield uncorrelated outputs so +# consecutive cache slots and consecutive request counters do not leave +# structural fingerprints in the Philox input stream. All outputs live in +# (0, 2**62) so they avoid Philox's degenerate seed=0 case while staying +# within int64. +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(x: int) -> int: + """SplitMix64 finalizer; pure function, no torch.""" + x = (x + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 30) + x = (x * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 27) + x = (x * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 31) + return x & _MAMBA_SSM_UINT64_MASK + + +def _compute_deterministic_mamba_seed(counter: int, slot: int, + rank_offset: int) -> int: + """Deterministic int64 seed in (0, 2**62) from (counter, slot, rank). + + Pure function (no RNG, no torch.randint). Identical inputs across + process invocations produce identical outputs, which is what the + acceptance criteria require for cross-run reproducibility. + """ + folded = (_MAMBA_SSM_SEED_BASE + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK) + folded &= _MAMBA_SSM_UINT64_MASK + value = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + if value == 0: + value = 1 + return value + + +def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, + device: torch.device) -> torch.Tensor: + """Allocate a (cache_size,) int64 CUDA buffer of deterministic seeds. + + counter=0 at allocation time; per-slot reset on fresh request assignment + bumps a host-side counter and rewrites only the freshly-assigned slot. + """ + slot_seeds = [ + _compute_deterministic_mamba_seed(0, i, rank_offset) + for i in range(cache_size) + ] + return torch.tensor(slot_seeds, dtype=torch.int64, device=device) + + +def _mamba_rank_offset(mapping: Mapping) -> int: + """Distinct seed offset per (tp_rank, pp_rank) so independent ranks + don't draw identical streams when not coordinated.""" + return (mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + + mapping.rank * 1_009) + + def use_cpp_mamba_cache_manager() -> bool: """Check if C++ MambaCacheManager should be used. @@ -317,12 +390,23 @@ def __init__( speculative_num_draft_tokens: Optional[int] = None, model_type: str = "nemotron_hybrid", use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, ) -> None: self.mamba_ssm_cache_dtype = ssm_cache_dtype self.speculative_num_draft_tokens = speculative_num_draft_tokens self.spec_state_size = spec_state_size self._use_replay_state_update = use_replay_state_update + # When True, allocate the per-slot Philox seed buffer even outside + # the replay path so the non-replay flashinfer SR kernel reads a + # persistent deterministic seed instead of a per-call torch.randint. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh cache-slot assignment. Combined + # with slot index and rank offset to produce reproducible per-slot + # seed values. Starts at 0 so the post-init "reset" stream is + # disjoint from the counter=0 stream used at allocation time. + self._seed_request_counter = 0 # get tp size tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size @@ -383,6 +467,17 @@ def __init__( device=device, ) + # Per-slot Philox seeds. Allocated whenever stochastic rounding can + # fire, even outside the replay path and even when MTP/spec is off, + # so all consumers (replay kernel, MTP non-replay flashinfer, + # non-MTP flashinfer) read a persistent deterministic seed from the + # cache manager instead of calling torch.randint per forward. + self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None + if (self._use_replay_state_update + or self._mamba_ssm_stochastic_rounding): + self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + max_batch_size, self._seed_rank_offset, device) + # create state container if speculative_num_draft_tokens is not None: T = speculative_num_draft_tokens + 1 @@ -397,6 +492,10 @@ def __init__( # SSM speculative cache — path-specific tensors spec_kwargs = {} + # Share the manager-level seed buffer through SpeculativeState + # so the MTP path can still read it via layer_cache. + if self._mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed if self._use_replay_state_update: assert n_groups % tp_size == 0, \ "replay state update requires n_groups divisible by tp_size" @@ -411,9 +510,6 @@ 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, @@ -539,12 +635,27 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() self.mamba_cache_index[r] = block + # Defensive zero of the per-slot SSM/conv state stripes for fresh + # context blocks. Without this, stale state from a prior occupant + # can leak into the first decode step of a recycled slot at high + # concurrency, which the non-replay flashinfer SR path reads as + # garbage and causes token-id-0 / `` collapse on greedy + # (temp=0.0) decoding. Harmless when the prefill fully overwrites. + self.mamba_cache.conv[:, block] = 0 + self.mamba_cache.temporal[:, block] = 0 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() + if self._mamba_ssm_rand_seed is not None: + # Deterministic per-slot rotation on fresh assignment. + # `block` is pulled from mamba_cache_free_blocks, which + # excludes _padding_slot by construction (see __init__), + # so padding sentinels never reach this branch. + self._seed_request_counter += 1 + self._mamba_ssm_rand_seed[block] = ( + _compute_deterministic_mamba_seed( + self._seed_request_counter, block, + self._seed_rank_offset)) def prepare_resources(self, scheduled_batch: ScheduledRequests): context_ids = [ @@ -631,6 +742,16 @@ def mamba_layer_cache(self, def get_mamba_ssm_cache_dtype(self) -> torch.dtype: return self.mamba_ssm_cache_dtype + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Return the persistent (cache_size,) int64 Philox seed buffer or + None when stochastic rounding is not active for this manager. + + Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. + Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a + view that matches the consuming kernel's expected shape. + """ + return self._mamba_ssm_rand_seed + @property def use_replay_state_update(self) -> bool: return self._use_replay_state_update @@ -725,6 +846,7 @@ def __init__( speculative_num_draft_tokens: Optional[int] = None, model_type: str = "nemotron_hybrid", use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, ) -> None: max_num_sequences = max_batch_size * mapping.pp_size self._use_cpp = use_cpp_mamba_cache_manager() @@ -763,6 +885,7 @@ def __init__( speculative_num_draft_tokens=speculative_num_draft_tokens, model_type=model_type, use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, ) def get_max_resource_count(self) -> int: @@ -811,6 +934,13 @@ def get_ssm_states(self, layer_idx: int) -> torch.Tensor: def get_mamba_ssm_cache_dtype(self) -> torch.dtype: return self._impl.get_mamba_ssm_cache_dtype() + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Delegate to the underlying Python manager. The C++ manager does + not allocate this buffer because it does not support speculative + decoding (and the SR-on-non-replay bug only fires under MTP).""" + getter = getattr(self._impl, 'get_mamba_ssm_rand_seed', None) + return getter() if getter is not None else None + @property def use_replay_state_update(self) -> bool: return getattr(self._impl, 'use_replay_state_update', False) @@ -904,6 +1034,7 @@ def __init__( model_type: str = "nemotron_hybrid", is_draft: bool = False, use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, ) -> None: # mamba hybrid cache requires block reuse to be disabled in KV cache config @@ -933,6 +1064,7 @@ def __init__( if spec_config is not None else None), model_type=model_type, use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, ) # initialize kv cache manager @@ -1039,6 +1171,7 @@ def __init__( is_estimating_kv_cache: bool = False, is_draft: bool = False, use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, model_type: str = "nemotron_hybrid", **kwargs, ) -> None: @@ -1078,6 +1211,14 @@ def __init__( # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work # on ranks with no local mamba layers. self._use_replay_state_update = use_replay_state_update + # Same allocation gate as PythonMambaCacheManager: the rand_seed + # buffer must exist whenever SR can fire, not only on the replay path. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh context-slot assignment; combined + # with the slot index and rank offset to produce reproducible per-slot + # seed values without any torch.randint. + self._seed_request_counter = 0 self.ssm_state_dtype = mamba_ssm_cache_dtype if self.local_num_mamba_layers == 0: @@ -1410,10 +1551,21 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): # 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() + num_contexts = len(scheduled_batch.context_requests) + if num_contexts > 0: + ctx_slots = self.cuda_state_indices[:num_contexts].long() + # Zero the SSM/conv state stripes for fresh context slots so + # neither the replay kernel nor the non-replay flashinfer SR + # kernel observes stale state on the first decode step of a + # recycled slot when the prefill does not fully overwrite the + # cache. Must fire on BOTH paths — gating this behind + # _use_replay_state_update caused token-id-0 / `` collapse + # on greedy temp=0 replay-off at c=32 (cf. iter4 root cause). + 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._use_replay_state_update and self.prev_num_accepted_tokens is not None: self.prev_num_accepted_tokens[ctx_slots] = 0 self.cache_buf_idx[ctx_slots] = 0 if self.old_x is not None: @@ -1424,22 +1576,27 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): 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, - ) + # Deterministic per-context-slot seed rotation. Runs whenever + # the seed buffer exists, including the non-replay SR path. + # Bump the host counter once per batch and write one new seed + # per fresh context slot from a pure function of + # (counter, slot, rank). No torch.randint involved. + if self.mamba_ssm_rand_seed is not None: + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = ctx_slots.cpu().tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, + rank_offset) + for slot in host_slots + ] + seed_tensor = torch.tensor( + new_seeds, + dtype=torch.int64, + device=self.mamba_ssm_rand_seed.device, + ) + self.mamba_ssm_rand_seed[ctx_slots] = seed_tensor def prepare_resources(self, scheduled_batch: ScheduledRequests): super().prepare_resources(scheduled_batch) @@ -1562,6 +1719,12 @@ def mamba_layer_cache( if self.spec_config is not None: layer_offset = self.mamba_layer_offsets[layer_idx] spec_kwargs = {} + # Per-cache-slot Philox seed buffer is shared across replay and + # non-replay MTP paths. The mixer asserts non-None on both + # branches when SR is enabled, so pass it through whenever it + # exists — not just on the replay branch. + if self.mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed if self._use_replay_state_update: # Per-layer slices for the replay kernel; shared 1D tensors # (cache_buf_idx, prev_num_accepted_tokens) are passed @@ -1573,7 +1736,6 @@ 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] @@ -1755,25 +1917,40 @@ def _setup_replay_buffers(self, spec_config) -> None: ``prev_num_accepted_tokens`` by these block indices, so the buffers must match the pool extent rather than ``max_batch_size``. """ + # Replay tensors require spec_config + replay path enabled. The + # rand_seed buffer is separable from replay and must also be + # allocated for non-replay SR so the flashinfer path has a + # persistent deterministic seed source. + 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 + self.old_dA_cumsum = None + + if (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding): + return + + cache_size = self.all_ssm_states.shape[1] + device = self.all_ssm_states.device + # Always-available deterministic seed buffer when SR (or replay) + # is on. Works for non-MTP runs because we don't depend on + # spec_config to allocate it. + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device) + if spec_config is None or not self._use_replay_state_update: - # Replay tensors are sized by the recurrent-state pool block count and - # 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 - self.old_dA_cumsum = None + # Without spec_config or replay we still keep the seed buffer + # (above) so the non-MTP flashinfer SR path has a persistent + # rand_seed source. return T = spec_config.max_draft_len + 1 num_local_mamba_layers = self.local_num_mamba_layers - # all_ssm_states: [num_local_mamba_layers, num_blocks_in_pool, ...] - cache_size = self.all_ssm_states.shape[1] nheads, head_dim, d_state = self.ssm_state_shape n_groups_per_rank = self._n_groups_per_rank - device = self.all_ssm_states.device # Shared across layers (consumed by the replay kernel via slot index). self.prev_num_accepted_tokens = torch.zeros(cache_size, @@ -1813,9 +1990,6 @@ 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: @@ -1823,3 +1997,8 @@ def use_replay_state_update(self) -> bool: def get_mamba_ssm_cache_dtype(self) -> torch.dtype: return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Return the persistent (cache_size,) int64 Philox seed buffer or + None when stochastic rounding is not active for this manager.""" + return getattr(self, 'mamba_ssm_rand_seed', None) diff --git a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py new file mode 100644 index 000000000000..f94a5f7a28cf --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the Mamba SSM stochastic-rounding Philox seed plumbing. + +The Mamba SSM SR path previously generated `rand_seed` tensors via +`torch.randint(..., (1,))` on every decode forward. The cache manager now +owns a persistent per-cache-slot int64 buffer that is deterministically +initialized and rewritten on fresh request assignment. These tests pin the +contract: pure-function seed generation, deterministic allocation, and +per-slot reset without `torch.randint`. +""" + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + PythonMambaCacheManager, + _allocate_mamba_seed_buffer, + _compute_deterministic_mamba_seed, + _mamba_rank_offset, +) +from tensorrt_llm.mapping import Mapping + + +def test_deterministic_seed_is_pure_function(): + s1 = _compute_deterministic_mamba_seed(counter=7, slot=3, rank_offset=42) + s2 = _compute_deterministic_mamba_seed(counter=7, slot=3, rank_offset=42) + assert s1 == s2 + assert 0 < s1 < (1 << 62) + + +def test_deterministic_seed_distinct_inputs_distinct_outputs(): + s_a = _compute_deterministic_mamba_seed(1, 0, 0) + s_b = _compute_deterministic_mamba_seed(1, 1, 0) + s_c = _compute_deterministic_mamba_seed(1, 0, 1) + s_d = _compute_deterministic_mamba_seed(2, 0, 0) + # All four (counter, slot, rank_offset) keys yield distinct seeds. + assert len({s_a, s_b, s_c, s_d}) == 4 + + +def test_rank_offset_distinct_per_rank(): + o0 = _mamba_rank_offset(Mapping(world_size=4, tp_size=4, pp_size=1, rank=0)) + o1 = _mamba_rank_offset(Mapping(world_size=4, tp_size=4, pp_size=1, rank=1)) + o2 = _mamba_rank_offset(Mapping(world_size=4, tp_size=4, pp_size=1, rank=2)) + o3 = _mamba_rank_offset(Mapping(world_size=4, tp_size=4, pp_size=1, rank=3)) + assert len({o0, o1, o2, o3}) == 4 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_allocated_buffer_is_reproducible_and_nonzero(): + device = torch.device("cuda") + buf1 = _allocate_mamba_seed_buffer(8, rank_offset=5, device=device) + buf2 = _allocate_mamba_seed_buffer(8, rank_offset=5, device=device) + assert buf1.dtype == torch.int64 + assert buf1.shape == (8,) + assert torch.equal(buf1, buf2) + assert (buf1 > 0).all().item() + # All eight slots should differ; collisions would defeat per-slot + # variance in the replay kernel. + assert buf1.unique().numel() == 8 + + +def _make_python_manager(*, sr: bool, replay: bool, max_batch_size: int = 4): + """Build a tiny PythonMambaCacheManager with MTP/spec enabled.""" + mapping = Mapping(world_size=1, tp_size=1, pp_size=1, rank=0) + return PythonMambaCacheManager( + d_state=4, + d_conv=2, + num_heads=2, + n_groups=1, + head_dim=2, + num_layers=1, + max_batch_size=max_batch_size, + spec_state_size=max_batch_size, + mapping=mapping, + dtype=torch.float16, + ssm_cache_dtype=torch.float16, + layer_mask=[True], + speculative_num_draft_tokens=1, + model_type="nemotron_hybrid", + use_replay_state_update=replay, + mamba_ssm_stochastic_rounding=sr, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_buffer_allocated_when_sr_only_no_replay(): + # The regression: SR enabled but replay off must still produce a + # persistent seed buffer so mamba2_mixer's non-replay branch can read it. + mgr = _make_python_manager(sr=True, replay=False) + seed_buf = mgr.get_mamba_ssm_rand_seed() + assert seed_buf is not None + assert seed_buf.dtype == torch.int64 + assert seed_buf.device.type == "cuda" + # Replay-only buffers are not allocated. + assert mgr.mamba_cache.old_x is None + assert mgr.mamba_cache.intermediate_ssm is not None # legacy SSM cache + # SpeculativeState exposes the same buffer (same Python identity). + assert mgr.mamba_cache.mamba_ssm_rand_seed is seed_buf + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_buffer_absent_when_neither_sr_nor_replay(): + mgr = _make_python_manager(sr=False, replay=False) + assert mgr.get_mamba_ssm_rand_seed() is None + assert mgr.mamba_cache.mamba_ssm_rand_seed is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_per_slot_reset_is_deterministic_and_uses_no_randint(): + mgr = _make_python_manager(sr=True, replay=False, max_batch_size=8) + seed_buf = mgr.get_mamba_ssm_rand_seed() + assert seed_buf is not None + pre = seed_buf.clone() + mgr._prepare_mamba_cache_blocks([1001]) + block = mgr.mamba_cache_index[1001] + post = seed_buf.clone() + # Exactly the freshly-assigned slot is rewritten. + diff_mask = pre != post + assert diff_mask.sum().item() == 1 + assert diff_mask[block].item() + # Rewrite is a deterministic function of (counter, slot, rank_offset). + expected = _compute_deterministic_mamba_seed( + mgr._seed_request_counter, block, mgr._seed_rank_offset + ) + assert post[block].item() == expected + assert expected > 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_padding_sentinel_does_not_churn_seeds(): + # CUDA-graph padding sentinels alias to the shared _padding_slot and + # must not rotate the seed entry of live real requests. + mgr = _make_python_manager(sr=True, replay=False, max_batch_size=8) + seed_buf = mgr.get_mamba_ssm_rand_seed() + assert seed_buf is not None + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + + sentinel_id = CUDA_GRAPH_DUMMY_REQUEST_ID + pre = seed_buf.clone() + counter_pre = mgr._seed_request_counter + mgr.add_dummy_requests([sentinel_id]) + post = seed_buf.clone() + # Sentinel assignment must leave the buffer (and the host counter) + # untouched; otherwise the seed buffer would churn every graph capture. + assert torch.equal(pre, post) + assert mgr._seed_request_counter == counter_pre + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_replay_path_still_allocates_seed_buffer(): + # Backward-compatibility: the replay path used to allocate the seed + # buffer; the new wiring must not regress that. + mgr = _make_python_manager(sr=False, replay=True) + seed_buf = mgr.get_mamba_ssm_rand_seed() + assert seed_buf is not None + assert mgr.mamba_cache.mamba_ssm_rand_seed is seed_buf + assert mgr.mamba_cache.old_x is not None # replay buffers also exist + + +def _build_cpp_hybrid(*, spec_config, use_replay: bool, sr: bool, max_batch_size: int = 4): + """Construct a CppMambaHybridCacheManager with one mamba + one attention + layer. Mirrors test_mamba_cache_manager._build_hybrid_with_mamba_layer + but parameterizes the replay / SR flags so we can exercise the + non-replay MTP SR layer-cache hand-off path.""" + from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import CppMambaHybridCacheManager + from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + mamba_mask = [True, False] + attn_mask = [False, True] + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + kv_cache_config = KvCacheConfig(max_tokens=512, enable_block_reuse=False) + return CppMambaHybridCacheManager( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=1, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=mapping, + spec_config=spec_config, + layer_mask=attn_mask, + use_replay_state_update=use_replay, + mamba_ssm_stochastic_rounding=sr, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_cpp_hybrid_non_replay_mtp_layer_cache_carries_rand_seed(): + """Regression: CppMambaHybridCacheManager.mamba_layer_cache() with + spec_config != None AND _use_replay_state_update == False AND + mamba_ssm_stochastic_rounding == True must still surface + mamba_ssm_rand_seed on the returned SpeculativeState. + + The mixer's non-replay MTP SR branch (mamba2_mixer.py) reads + `layer_cache.mamba_ssm_rand_seed` and asserts non-None. Iter5 review + caught the regression where the seed was only forwarded inside the + replay branch of mamba_layer_cache; this test pins both paths.""" + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + spec_config = MTPDecodingConfig(max_draft_len=2) + mgr = _build_cpp_hybrid(spec_config=spec_config, use_replay=False, sr=True) + # Manager-level buffer must exist (SR is on). + seed_buf = mgr.get_mamba_ssm_rand_seed() + assert seed_buf is not None + assert seed_buf.dtype == torch.int64 + # Layer cache exposes the same buffer (Python identity equality). + layer_cache = mgr.mamba_layer_cache(0) + assert layer_cache is not None + assert layer_cache.mamba_ssm_rand_seed is mgr.mamba_ssm_rand_seed + # And the SpeculativeState is on the non-replay legacy branch. + assert layer_cache.intermediate_ssm is not None + assert layer_cache.old_x is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_cpp_hybrid_replay_mtp_layer_cache_still_carries_rand_seed(): + """Backward-compat: the replay branch must keep forwarding the seed + buffer through mamba_layer_cache.""" + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + spec_config = MTPDecodingConfig(max_draft_len=2) + mgr = _build_cpp_hybrid(spec_config=spec_config, use_replay=True, sr=True) + seed_buf = mgr.get_mamba_ssm_rand_seed() + assert seed_buf is not None + layer_cache = mgr.mamba_layer_cache(0) + assert layer_cache is not None + assert layer_cache.mamba_ssm_rand_seed is mgr.mamba_ssm_rand_seed + # Replay-specific compact buffers are populated. + assert layer_cache.old_x is not None + assert layer_cache.cache_buf_idx is not None + assert layer_cache.prev_num_accepted_tokens is not None From 90ae545fee92e77747bb17090ce1d95fa95d67e4 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 26 May 2026 22:01:37 -0700 Subject: [PATCH 3/5] [None][fix] Capture Mamba fresh state reset in forward Signed-off-by: qgai --- .../_torch/modules/mamba/mamba2_mixer.py | 2 ++ .../_torch/pyexecutor/mamba_cache_manager.py | 20 ------------------- 2 files changed, 2 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 2ea58ad3c725..0c75aa58bd77 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -348,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, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 3087749d0797..c5ca7a222b11 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1542,29 +1542,9 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): if self._pending_state_transfers: logger.info(f"Need to transfer mamba state blocks") self._setup_state_indices() - # Reset replay double-buffer state for fresh context blocks. A reused - # 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. num_contexts = len(scheduled_batch.context_requests) if num_contexts > 0: ctx_slots = self.cuda_state_indices[:num_contexts].long() - # Zero the SSM/conv state stripes for fresh context slots so - # neither the replay kernel nor the non-replay flashinfer SR - # kernel observes stale state on the first decode step of a - # recycled slot when the prefill does not fully overwrite the - # cache. Must fire on BOTH paths — gating this behind - # _use_replay_state_update caused token-id-0 / `` collapse - # on greedy temp=0 replay-off at c=32 (cf. iter4 root cause). - 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._use_replay_state_update and self.prev_num_accepted_tokens is not None: self.prev_num_accepted_tokens[ctx_slots] = 0 self.cache_buf_idx[ctx_slots] = 0 From 1105153dfa0eb2acd15ae60e51d55c337206c4f3 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 26 May 2026 22:05:47 -0700 Subject: [PATCH 4/5] [None][fix] Keep replay state update scoped to seed Signed-off-by: qgai --- .../mamba/replay_selective_state_update.py | 96 ++----------------- 1 file changed, 7 insertions(+), 89 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 714be639cde4..9c510f8f4956 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -180,16 +180,8 @@ 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) @@ -235,11 +227,7 @@ def _replay_precompute_kernel( ) # Compute raw CB once — shared across all heads in this block - raw_CB = tl.dot( - C_all.to(tl.float32), - tl.trans(B_all).to(tl.float32), - input_precision="ieee", - ) + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) # Store B to cache (once per group, only if this block covers the first heads) if first_head % nheads_ngroups_ratio == 0: @@ -282,8 +270,6 @@ 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( @@ -435,10 +421,6 @@ 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) @@ -454,14 +436,6 @@ 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 @@ -472,16 +446,6 @@ 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. @@ -489,32 +453,18 @@ def _replay_state_update_kernel( 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 - 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) + coeff = tl.where(offs_t < prev_num_accepted_tokens, 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=pnat_t_mask[:, None] & m_mask[None, :], + mask=t_mask[:, None] & m_mask[None, :], other=0.0, ) @@ -527,7 +477,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=pnat_t_mask[:, None] & n_mask[None, :], + mask=t_mask[:, None] & n_mask[None, :], other=0.0, ).to(tl.float32) @@ -539,19 +489,7 @@ 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.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) + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) # Write post-replay state if USE_RS_ROUNDING: @@ -635,36 +573,16 @@ def _replay_state_update_kernel( ) # init_out = C_all @ state^T * decay_vec - init_out = ( - tl.dot( - C_all.to(tl.float32), - tl.trans(state).to(tl.float32), - input_precision="ieee", - ) - * decay_vec[:, None] - ) + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] # cb_out = CB_scaled @ x_all - cb_out = tl.dot( - CB_scaled.to(tl.float32), - x_all.to(tl.float32), - input_precision="ieee", - ) + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) 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 ``. 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( From 11e8ae210091a0d96bcd01fcd8fac6748e2ae278 Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 27 May 2026 04:22:42 -0700 Subject: [PATCH 5/5] [None][fix] Remove fresh Mamba cache zeroing Signed-off-by: qgai --- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index c5ca7a222b11..c6f4125b6c66 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -635,14 +635,6 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() self.mamba_cache_index[r] = block - # Defensive zero of the per-slot SSM/conv state stripes for fresh - # context blocks. Without this, stale state from a prior occupant - # can leak into the first decode step of a recycled slot at high - # concurrency, which the non-replay flashinfer SR path reads as - # garbage and causes token-id-0 / `` collapse on greedy - # (temp=0.0) decoding. Harmless when the prefill fully overwrites. - self.mamba_cache.conv[:, block] = 0 - self.mamba_cache.temporal[:, block] = 0 if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0