diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 1474f8f1fe2b..0c75aa58bd77 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 @@ -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( 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..9c510f8f4956 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/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 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 50edbee51db7..b4eebeb1a25c 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -56,6 +56,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. @@ -277,8 +350,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 @@ -290,6 +364,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. @@ -313,12 +391,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 @@ -379,6 +468,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 @@ -393,6 +493,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" @@ -535,6 +639,16 @@ 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_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 = [ @@ -621,6 +735,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 @@ -642,6 +766,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), @@ -714,6 +839,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() @@ -752,6 +878,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: @@ -800,6 +927,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) @@ -893,6 +1027,7 @@ def __init__( model_type: str = "nemotron_hybrid", is_draft: bool = False, use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, # Per-pool configurations forwarded to the C++ KVCacheManager ctor. # Lets a single manager host pools with mixed shapes (e.g. Gemma4 # hybrid attention). See KVCacheManager.__init__. @@ -926,6 +1061,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 @@ -1033,6 +1169,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: @@ -1072,6 +1209,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: @@ -1327,6 +1472,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 @@ -1395,17 +1541,41 @@ 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. - 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() + if self._use_replay_state_update and self.prev_num_accepted_tokens is not None: 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 + # 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) @@ -1528,6 +1698,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 @@ -1720,24 +1896,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.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, @@ -1784,3 +1976,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