flashinfer.mamba.replayssm_materialize

flashinfer.mamba.replayssm_materialize(state_ptrs: Tensor, state_slot_strides: Tensor, x_cache_ptrs: Tensor, x_cache_slot_strides: Tensor, B_cache_ptrs: Tensor, B_cache_slot_strides: Tensor, dt_cache_ptrs: Tensor, dt_cache_slot_strides: Tensor, A_ptrs: Tensor, state_scale_ptrs: Tensor, state_scale_slot_strides: Tensor, src_slots: Tensor, dst_slots: Tensor, ring_start: Tensor, replay_prefix_len: Tensor, active_request_indices: Tensor, *, state_dtype: dtype, input_dtype: dtype, matrixA_dtype: dtype, dim: int, dstate: int, num_heads: int, heads_per_group: int, max_window: int, ring_buffer_len: int, pad_slot_id: int = -1, rand_seed: Tensor | None = None, philox_rounds: int = 0, dependency_inputs: list[Tensor] | None = None, dependency_outputs: list[Tensor] | None = None) → None

Materialize an SSM state at a selected token, from an older state and a ReplaySSM ring buffer.

A source state plus its replay ring represents every state from the source state through the state advanced by the accepted-token prefix, see flashinfer.mamba.checkpointing_ssu(). This operation starts at ring_start, applies a caller-selected prefix of that ring to the source state, and writes the resulting state to a separate destination slot. It does not consume or otherwise modify replay-ring data.

Parameters:

  • state_ptrs (torch.Tensor) – One-dimensional CUDA int64 pointer table of shape (L,). Each entry is the base address of that layer’s state storage.

  • state_slot_strides (torch.Tensor) – One-dimensional CUDA int64 table of shape (L,) giving the state storage outer-slot stride, in elements, for each layer.

  • x_cache_ptrs (torch.Tensor) – One-dimensional CUDA int64 pointer table of shape (L,) for the per-head replay x cache of each layer.

  • x_cache_slot_strides (torch.Tensor) – One-dimensional CUDA int64 table of shape (L,) giving the x cache outer-slot stride, in elements, for each layer.

  • B_cache_ptrs (torch.Tensor) – One-dimensional CUDA int64 pointer table of shape (L,) for the per-group replay B cache of each layer.

  • B_cache_slot_strides (torch.Tensor) – One-dimensional CUDA int64 table of shape (L,) giving the B cache outer-slot stride, in elements, for each layer.

  • dt_cache_ptrs (torch.Tensor) – One-dimensional CUDA int64 pointer table of shape (L,) for the replay dt cache of each layer.

  • dt_cache_slot_strides (torch.Tensor) – One-dimensional CUDA int64 table of shape (L,) giving the dt cache outer-slot stride, in elements, for each layer.

  • A_ptrs (torch.Tensor) – One-dimensional CUDA int64 pointer table of shape (L,) for each layer’s A vector.

  • state_scale_ptrs (torch.Tensor) – One-dimensional CUDA int64 pointer table of shape (L,) for quantized-state scales. Use zero entries for non-quantized state.

  • state_scale_slot_strides (torch.Tensor) – One-dimensional CUDA int64 table of shape (L,) giving the scale outer-slot stride, in elements, for each layer.

  • src_slots (torch.Tensor) – CUDA int32 tensors of shape (L, B) selecting source and destination slots for each layer and physical batch request. A request/layer is skipped when either slot equals pad_slot_id.

  • dst_slots (torch.Tensor) – CUDA int32 tensors of shape (L, B) selecting source and destination slots for each layer and physical batch request. A request/layer is skipped when either slot equals pad_slot_id.

  • ring_start (torch.Tensor) – CUDA int32 tensors of shape (B,). replay_prefix_len[b] is the number of ring entries, beginning at ring_start[b], to apply to src_slots[:, b] before writing dst_slots[:, b]. Zero performs an exact state/scale copy. For a shared source-state/ring snapshot, it must not exceed checkpointing SSU’s prev_num_accepted_tokens for that request. For every request selected by active_request_indices, ring_start must be in [0, ring_buffer_len) and replay_prefix_len must be in [0, max_window]. This operation does not receive or validate the checkpointing tracker.

  • replay_prefix_len (torch.Tensor) – CUDA int32 tensors of shape (B,). replay_prefix_len[b] is the number of ring entries, beginning at ring_start[b], to apply to src_slots[:, b] before writing dst_slots[:, b]. Zero performs an exact state/scale copy. For a shared source-state/ring snapshot, it must not exceed checkpointing SSU’s prev_num_accepted_tokens for that request. For every request selected by active_request_indices, ring_start must be in [0, ring_buffer_len) and replay_prefix_len must be in [0, max_window]. This operation does not receive or validate the checkpointing tracker.

  • active_request_indices (torch.Tensor) –

    CUDA int32 tensor of shape (B,). Its prefix contains the selected physical batch indices, each exactly once and in any order; remaining entries are -1. The kernel stops at the first -1. Each selected index must be in [0, B) and refer to metadata satisfying the ring_start and replay_prefix_len requirements above.

    The contents of these metadata tensors are GPU-resident and are not checked. Violating their value requirements is invalid input and has undefined behavior.

  • state_dtype (torch.dtype) – JIT state-storage dtype. One-byte state requires dim=64 and dstate=128.

  • input_dtype (torch.dtype) – JIT dtype of x and B cache entries. Must be torch.bfloat16, matching checkpointing SSU’s BF16 replay MMA operands. The pointer tables cannot validate the pointed-to tensor dtypes, so callers must ensure both cache families use BF16 storage.

  • matrixA_dtype (torch.dtype) – JIT dtype of the per-head A values.

  • dim (int) – JIT state head dimension.

  • dstate (int) – JIT SSM state dimension.

  • num_heads (int) – Runtime number of heads in every layer.

  • heads_per_group (int) – JIT number of heads sharing each B cache group; it must divide num_heads.

  • max_window (int) – JIT maximum replay prefix length, in ring entries, in [1, 16].

  • ring_buffer_len (int) – Runtime number of rows in every replay ring.

  • pad_slot_id (int) – Slot sentinel to skip, matching flashinfer.mamba.checkpointing_ssu(). Defaults to -1; zero is a valid slot unless passed here explicitly.

  • rand_seed (Optional[torch.Tensor]) – One-element CUDA int64 seed required when philox_rounds > 0.

  • philox_rounds (int) – JIT number of Philox stochastic-rounding rounds. Zero disables stochastic rounding.

  • dependency_inputs (Optional[list[torch.Tensor]]) – Backing CUDA tensors read indirectly through pointer tables: typically the source state and scale storage selected by src_slots through state_ptrs and state_scale_ptrs, together with storage represented by x_cache_ptrs, B_cache_ptrs, dt_cache_ptrs, and A_ptrs. These tensors are dependency anchors for torch.compile and are not passed to the JIT kernel. Do not repeat a state or scale backing tensor here when it is also an output anchor.

  • dependency_outputs (Optional[list[torch.Tensor]]) – Backing CUDA tensors written to the destination slots selected by dst_slots through state_ptrs and, for quantized state, state_scale_ptrs. These tensors are mutable dependency anchors for torch.compile and are not passed to the JIT kernel. Callers using torch.compile must include every backing allocation whose contents they need PyTorch to order around this call.

Notes

Source state and replay rings are read-only. Caller-owned slot tables and metadata select the tokens to apply; this operation has no acceptance, alignment, or cursor-update semantics.

For example, if src_slots selects the state after token 123 and seven subsequent tokens have been accepted into the ring, a prefix cache that materializes states every 128 tokens passes replay_prefix_len=5 to write the state after token 128.