`core.ssm.ops.mamba2.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.
`_decode_step_prep_kernel`	One decode entry’s bookkeeping: integer updates and row copies only.
`_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.mamba2.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.mamba2.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.mamba2.batch_invariant_decode._decode_step_prep_kernel( batch_indices_ptr, num_buffered_ptr, x_ptr, dt_ptr, B_ptr, C_ptr, x_buf_ptr, dt_buf_ptr, B_buf_ptr, C_buf_ptr, out_ptr, target_rows_ptr, chunk_flags_ptr, chunk_starts_ptr, stride_x, stride_dt, stride_B, stride_C, X_SIZE: triton.language.constexpr, DT_SIZE: triton.language.constexpr, BC_SIZE: triton.language.constexpr, OUT_SIZE: triton.language.constexpr, CHUNK_SIZE: triton.language.constexpr, BLOCK_SIZE: triton.language.constexpr, )#

One decode entry’s bookkeeping: integer updates and row copies only.

For entry e with slot s and write position w = num_buffered[s]: copy the current token’s x / dt / B / C rows into buffer row s * CHUNK_SIZE + w, set target_rows[e] = w, chunk_flags[e] = (w + 1 == CHUNK_SIZE), chunk_starts[e] = s * CHUNK_SIZE, zero out[e] and advance num_buffered[s] (0 after a crossing). Inactive entries (s < 0) only get target row -1, flag 0 and a zero output.

core.ssm.ops.mamba2.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')), use_row_kernels=True, )#

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. z is the current token’s gate per entry (decode_batch_size, nheads, headdim): the gate is row-local and only the target row is computed, so no other row’s z is needed.

When the training kernels are pinned to one config each, the row kernels (batch_invariant_decode_rows) compute the same bits with that config and less work; otherwise the training kernels run in target-rows mode.

class core.ssm.ops.mamba2.batch_invariant_decode.BatchInvariantDecodeBuffers#

Per-slot persistent state for the buffered decode scan.

x: 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

chunk_starts: 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.mamba2.batch_invariant_decode.BatchInvariantDecodeBuffers#: Allocate the per-slot decode buffers.

seed( x: 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.mamba2.batch_invariant_decode.batch_invariant_decode_buffered_scan( buffers: core.ssm.ops.mamba2.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.mamba2.batch_invariant_decode.MambaBatchInvariantDecode(mixer)#

Adapter between a MambaMixer and the buffered decode.

Initialization

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

seed(x, 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.mamba2.batch_invariant_decode#

Module Contents#

Classes#

Functions#

API#

`core.ssm.ops.mamba2.batch_invariant_decode`#