Is this page helpful?

core.ssm.ops.mamba2.batch_invariant_decode_rows#

Row kernels for batch-invariant Mamba decode.

The buffered decode replays the training chunk-scan kernels for the one row a step consumes. Those kernels still compute a whole training tile per (entry, head): the chunk scan a BLOCK_SIZE_M-row M tile, CB a full M tile, and chunk state plus state passing launch programs for every entry although only entries that fill their chunk (about one in chunk_size) commit a boundary state.

The kernels here compute only what the step consumes, with the training kernels’ statements and tile decomposition so each stored value is bitwise the training one:

  • _chunk_scan_row_kernel: the chunk scan on the 16-row MMA tile holding the target row, with the training tile’s BLOCK_SIZE_N / BLOCK_SIZE_K and causal bound K_MAX. A tensor-core MMA computes each output row from that row of A alone, so with the same instruction (mma.sync), k decomposition and accumulation order, the target row’s bits do not depend on how many other rows share the tile.

  • _bmm_row_kernel: CB with the training kernel’s config, loading and storing only the target row.

  • _chunk_state_passing_rows_kernel: chunk state (training config) and the state passing update exp(dA) * S + new fused, run only for entries that cross a chunk boundary. The update is one fma, as the state-passing kernel compiles it.

Each kernel takes its tile from the config the training autotuner launches with (configs holds exactly one entry once batch-invariant mode pins it). When a kernel is not pinned to one config, or the chunk-scan tile does not use mma.sync, the caller keeps the training kernels.

Module Contents#

Functions#

_pinned_config

The one config kernel’s autotuner launches with, or None while it still chooses.

row_kernel_configs

Training configs (chunk scan, CB, chunk state) the row kernels mirror, or None.

_chunk_scan_row_kernel

_chunk_scan_fwd_kernel in target-rows mode for the target row only.

_bmm_row_kernel

_bmm_chunk_fwd_kernel in target-rows mode, loading and storing only the target row.

_chunk_state_passing_rows_kernel

_chunk_state_fwd_kernel then _state_passing_fwd_kernel for crossing entries.

_launch_options

_block_sizes

chunk_scan_rows

_chunk_scan_fwd(..., target_rows=, chunk_starts=) for the target rows only.

bmm_rows

_bmm_chunk_fwd(C, B, ..., output_dtype=fp32, target_rows=, chunk_starts=).

_chunks_per_program

About one crossing per program: 1 below 32 entries, 4 at 64, 16 from 256.

chunk_state_passing_rows

Commit the boundary state of every flagged entry into states[slot].

Data#

ROW_BLOCK_SIZE_M

API#

core.ssm.ops.mamba2.batch_invariant_decode_rows.ROW_BLOCK_SIZE_M#

16

core.ssm.ops.mamba2.batch_invariant_decode_rows._pinned_config(kernel)#

The one config kernel’s autotuner launches with, or None while it still chooses.

core.ssm.ops.mamba2.batch_invariant_decode_rows.row_kernel_configs(C: torch.Tensor, chunk_size: int)#

Training configs (chunk scan, CB, chunk state) the row kernels mirror, or None.

None keeps the training kernels: an autotuner not pinned to one config, a dstate the chunk scan splits over several dots, or a chunk-scan tile large enough for warpgroup / tcgen05 MMAs (num_warps >= 4 and BLOCK_SIZE_M >= 64), whose instruction the 16-row tile cannot reproduce (on GB200 every such candidate differs from the row kernel in fp32; every mma.sync candidate matches it bitwise).

core.ssm.ops.mamba2.batch_invariant_decode_rows._chunk_scan_row_kernel(
cb_ptr,
x_ptr,
z_ptr,
out_ptr,
dt_ptr,
dA_cumsum_ptr,
C_ptr,
D_ptr,
initstates_ptr,
chunk_starts_ptr,
slots_ptr,
target_rows_ptr,
chunk_size: triton.language.constexpr,
hdim: triton.language.constexpr,
dstate: triton.language.constexpr,
nheads_ngroups_ratio: triton.language.constexpr,
stride_cb_chunk: triton.language.int64,
stride_cb_head: triton.language.int64,
stride_cb_csize_m: triton.language.int64,
stride_cb_csize_k: triton.language.constexpr,
stride_x_seqlen: triton.language.int64,
stride_x_head: triton.language.int64,
stride_x_hdim: triton.language.constexpr,
stride_z_seqlen: triton.language.int64,
stride_z_head: triton.language.int64,
stride_z_hdim: triton.language.constexpr,
stride_out_seqlen: triton.language.int64,
stride_out_head: triton.language.int64,
stride_out_hdim: triton.language.constexpr,
stride_dt_chunk: triton.language.int64,
stride_dt_head: triton.language.int64,
stride_dt_csize: triton.language.constexpr,
stride_dA_cs_chunk: triton.language.int64,
stride_dA_cs_head: triton.language.int64,
stride_dA_cs_csize: triton.language.constexpr,
stride_C_seqlen: triton.language.int64,
stride_C_head: triton.language.int64,
stride_C_dstate: triton.language.constexpr,
stride_init_states_batch: triton.language.int64,
stride_init_states_head: triton.language.int64,
stride_init_states_hdim: triton.language.int64,
stride_init_states_dstate: triton.language.constexpr,
stride_D_head: triton.language.constexpr,
HAS_D: triton.language.constexpr,
D_HAS_HDIM: triton.language.constexpr,
HAS_Z: triton.language.constexpr,
TILE_SIZE_M: triton.language.constexpr,
BLOCK_SIZE_M: triton.language.constexpr,
BLOCK_SIZE_N: triton.language.constexpr,
BLOCK_SIZE_K: triton.language.constexpr,
BLOCK_SIZE_DSTATE: triton.language.constexpr,
)#

_chunk_scan_fwd_kernel in target-rows mode for the target row only.

The training kernel’s statements on a BLOCK_SIZE_M-row tile placed over the target row, with K_MAX from the training tile (TILE_SIZE_M rows) so the same zero-weight products are summed. Only the target row of C and cb is loaded; the other rows are zero and never reach the stored row.

core.ssm.ops.mamba2.batch_invariant_decode_rows._bmm_row_kernel(
a_ptr,
b_ptr,
out_ptr,
chunk_starts_ptr,
target_rows_ptr,
chunk_size: triton.language.constexpr,
K: triton.language.constexpr,
ngroups: triton.language.constexpr,
stride_a_seqlen: triton.language.int64,
stride_a_head: triton.language.int64,
stride_ak: triton.language.constexpr,
stride_b_seqlen: triton.language.int64,
stride_b_head: triton.language.int64,
stride_bk: triton.language.constexpr,
stride_out_chunk: triton.language.int64,
stride_out_head: triton.language.int64,
stride_outm: triton.language.int64,
stride_outn: triton.language.constexpr,
dot_dtype: triton.language.constexpr,
BLOCK_SIZE_M: triton.language.constexpr,
BLOCK_SIZE_N: triton.language.constexpr,
BLOCK_SIZE_K: triton.language.constexpr,
)#

_bmm_chunk_fwd_kernel in target-rows mode, loading and storing only the target row.

core.ssm.ops.mamba2.batch_invariant_decode_rows._chunk_state_passing_rows_kernel(
x_ptr,
b_ptr,
dt_ptr,
dA_cumsum_ptr,
chunk_starts_ptr,
chunk_flags_ptr,
slots_ptr,
states_ptr,
nchunks,
hdim: triton.language.constexpr,
dstate: triton.language.constexpr,
chunk_size: triton.language.constexpr,
nheads_ngroups_ratio: triton.language.constexpr,
stride_x_seqlen: triton.language.int64,
stride_x_head: triton.language.int64,
stride_x_hdim: triton.language.constexpr,
stride_b_seqlen: triton.language.int64,
stride_b_head: triton.language.int64,
stride_b_dstate: triton.language.constexpr,
stride_dt_head: triton.language.int64,
stride_dt_chunk: triton.language.int64,
stride_dt_csize: triton.language.constexpr,
stride_dA_cs_head: triton.language.int64,
stride_dA_cs_chunk: triton.language.int64,
stride_dA_cs_csize: triton.language.constexpr,
stride_states_batch: triton.language.int64,
stride_states_head: triton.language.int64,
stride_states_hdim: triton.language.int64,
stride_states_dstate: triton.language.constexpr,
CHUNKS_PER_PROGRAM: triton.language.constexpr,
BLOCK_SIZE_M: triton.language.constexpr,
BLOCK_SIZE_N: triton.language.constexpr,
BLOCK_SIZE_K: triton.language.constexpr,
)#

_chunk_state_fwd_kernel then _state_passing_fwd_kernel for crossing entries.

For each entry with a chunk flag, the chunk-state tile (training statements and tile), then the state-passing update of the slot’s state on the same fp32 values. A program owns CHUNKS_PER_PROGRAM entries of one (head, tile) and skips entries without a flag.

core.ssm.ops.mamba2.batch_invariant_decode_rows._launch_options(config)#
core.ssm.ops.mamba2.batch_invariant_decode_rows._block_sizes(config)#
core.ssm.ops.mamba2.batch_invariant_decode_rows.chunk_scan_rows(
cb,
x,
z,
dt,
dA_cumsum,
C,
out,
slots,
target_rows,
chunk_starts,
initial_states,
D,
config,
)#

_chunk_scan_fwd(..., target_rows=, chunk_starts=) for the target rows only.

core.ssm.ops.mamba2.batch_invariant_decode_rows.bmm_rows(C, B, chunk_size, target_rows, chunk_starts, config)#

_bmm_chunk_fwd(C, B, ..., output_dtype=fp32, target_rows=, chunk_starts=).

Only each entry’s target row of the result is written.

core.ssm.ops.mamba2.batch_invariant_decode_rows._chunks_per_program(nchunks: int) → int#

About one crossing per program: 1 below 32 entries, 4 at 64, 16 from 256.

core.ssm.ops.mamba2.batch_invariant_decode_rows.chunk_state_passing_rows(
x,
B,
dt,
dA_cumsum,
chunk_starts,
chunk_flags,
slots,
states,
config,
)#

Commit the boundary state of every flagged entry into states[slot].

_chunk_state_fwd(..., chunk_flags=) followed by _state_passing_fwd(..., dst_states=states, dst_indices=slots, dst_flags=chunk_flags).