`core.ssm.ops.batch_invariant_decode`#

Batch-invariant Mamba decode using buffered chunk replay.

Module Contents#

Classes#

`BatchInvariantDecodeBuffers`	Per-slot persistent state for the buffered decode scan.
`MambaBatchInvariantDecode`	Adapter between a MambaMixer and the buffered decode.

Functions#

`_masked_update_rows_kernel`	Copy contiguous rows, skipping entries whose destination index is -1.
`_masked_update_rows`	Copy rows into persistent BIK buffers without touching inactive graph lanes.
`_mamba_chunk_scan_decode_rows`	Run the training scan pipeline over buffered decode chunks.
`batch_invariant_decode_buffered_scan`	Run one decode token with full chunk-scan arithmetic.

API#

core.ssm.ops.batch_invariant_decode._masked_update_rows_kernel( states_ptr, indices_ptr, values_ptr, state_row_stride, value_row_stride, ROW_SIZE: triton.language.constexpr, BLOCK_SIZE: triton.language.constexpr, )#: Copy contiguous rows, skipping entries whose destination index is -1.

core.ssm.ops.batch_invariant_decode._masked_update_rows( states: torch.Tensor, indices: torch.Tensor, values: torch.Tensor, ) → None#: Copy rows into persistent BIK buffers without touching inactive graph lanes.

core.ssm.ops.batch_invariant_decode._mamba_chunk_scan_decode_rows( x, z, dt, A, B, C, chunk_size, chunk_starts, slots, target_rows, chunk_flags, initial_states, out, D=None, dt_bias=None, dt_softplus=False, dt_limit=(0.0, float('inf')), )#

Run the training scan pipeline over buffered decode chunks.

Each kernel computes only the row or boundary consumed by this decode step, while preserving the training kernel’s arithmetic for that result.

class core.ssm.ops.batch_invariant_decode.BatchInvariantDecodeBuffers#

Per-slot persistent state for the buffered decode scan.

x: torch.Tensor#: None

z: torch.Tensor#: None

dt: torch.Tensor#: None

B: torch.Tensor#: None

C: torch.Tensor#: None

num_buffered: torch.Tensor#: None

out: torch.Tensor#: None

target_rows: torch.Tensor#: None

chunk_flags: torch.Tensor#: None

classmethod allocate( max_requests: int, chunk_size: int, nheads: int, headdim: int, ngroups: int, dstate: int, device: torch.device, dtype: torch.dtype, ) → core.ssm.ops.batch_invariant_decode.BatchInvariantDecodeBuffers#: Allocate the per-slot decode buffers.

seed( x: torch.Tensor, z: torch.Tensor, dt: torch.Tensor, B: torch.Tensor, C: torch.Tensor, cu_seqlens: torch.Tensor, batch_indices: torch.Tensor, ) → None#: Store each prefill’s unfinished chunk for decode replay.

core.ssm.ops.batch_invariant_decode.batch_invariant_decode_buffered_scan( buffers: core.ssm.ops.batch_invariant_decode.BatchInvariantDecodeBuffers, x: torch.Tensor, z: torch.Tensor, dt: torch.Tensor, B: torch.Tensor, C: torch.Tensor, A: torch.Tensor, D: torch.Tensor, dt_bias: torch.Tensor, batch_indices: torch.Tensor, ssm_state: torch.Tensor, ) → torch.Tensor#

Run one decode token with full chunk-scan arithmetic.

Mutates the replay buffers and commits ssm_state when a chunk fills.

class core.ssm.ops.batch_invariant_decode.MambaBatchInvariantDecode(mixer)#

Adapter between a MambaMixer and the buffered decode.

Initialization

_get_buffers( max_requests, x, B, ) → core.ssm.ops.batch_invariant_decode.BatchInvariantDecodeBuffers#

seed(x, z, dt, B, C, cu_seqlens, batch_indices, max_requests) → None#: Seed replay buffers from the prefill tail.

step(x, z, dt, B, C, batch_indices, ssm_state) → torch.Tensor#: Run one decode step using the mixer’s flattened layouts.

core.ssm.ops.batch_invariant_decode#

Module Contents#

Classes#

Functions#

API#

`core.ssm.ops.batch_invariant_decode`#