core.ssm.ops.intermediate_extraction

Fused gather + conditional-scatter kernels for Mamba intermediate-state extraction used by prefix caching.

These replace the two-step states[indices] (dense gather) + .copy_() (scratch write) pattern with a single kernel that:

1. Reads a runtime real_count from a fixed-address GPU tensor.

2. For each slot i < real_count, gathers the source row indexed by the per-slot index/position and writes it directly into the destination scratch.

3. For each slot i >= real_count, returns immediately (no work, no write).

This is CUDA-graph safe: the launch grid is sized at capture time to the maximum possible slot count, but per-program execution is data-conditional on the runtime real_count, so padded slots cost almost nothing.

Module Contents

Functions

_scatter_intermediate_ssm_kernel	Conditional gather+scatter for SSM intermediate states.
scatter_intermediate_ssm	Gather rows of states at chunk_indices and scatter into out, for the first real_count_gpu slots only.
_scatter_intermediate_conv_kernel	Conditional gather of a length-D_CONV conv window per slot.
scatter_intermediate_conv	Gather length-d_conv conv windows from src at abs_positions and scatter (transposed) into out, for the first real_count_gpu slots only.

API

core.ssm.ops.intermediate_extraction._scatter_intermediate_ssm_kernel(

states_ptr,

chunk_indices_ptr,

real_count_ptr,

out_ptr,

state_flat,

states_row_stride,

out_row_stride,

BLOCK_N: triton.language.constexpr,

)

Conditional gather+scatter for SSM intermediate states.

Grid: (max_count, ceil(state_flat / BLOCK_N)). Each program owns one (slot, column-block) pair; programs with pid_slot >= real_count exit immediately, so padded slots produce no HBM traffic.

core.ssm.ops.intermediate_extraction.scatter_intermediate_ssm(

states: torch.Tensor,

chunk_indices: torch.Tensor,

real_count_gpu: torch.Tensor,

out: torch.Tensor,

) → None

Gather rows of states at chunk_indices and scatter into out, for the first real_count_gpu slots only.

Parameters:

• states – (num_chunks, *ssm_state_shape) chunk-scan output for one layer.

• chunk_indices – (max_count,) int64 per-slot gather index.

• real_count_gpu – int32[1] GPU tensor with the runtime real count.

• out – (max_count, *ssm_state_shape) destination scratch slice (one layer).

core.ssm.ops.intermediate_extraction._scatter_intermediate_conv_kernel(

src_ptr,

abs_positions_ptr,

real_count_ptr,

out_ptr,

seq_len,

conv_dim,

src_stride_s,

src_stride_c,

out_slot_stride,

D_CONV: triton.language.constexpr,

BLOCK_C: triton.language.constexpr,

)

Conditional gather of a length-D_CONV conv window per slot.

Reads window [abs_pos - D_CONV, abs_pos) from src_ptr (clamped into [0, seq_len)) and writes it transposed into out[slot, :, :] of shape (conv_dim, D_CONV). The transpose is folded into the write pattern to match the slot allocator’s storage layout.

core.ssm.ops.intermediate_extraction.scatter_intermediate_conv(

src: torch.Tensor,

abs_positions: torch.Tensor,

real_count_gpu: torch.Tensor,

out: torch.Tensor,

d_conv: int,

) → None

Gather length-d_conv conv windows from src at abs_positions and scatter (transposed) into out, for the first real_count_gpu slots only.

Parameters:

• src – (batch, seq_len, conv_dim) pre-conv xBC tensor. Batch is assumed to be 1 (inference); only batch index 0 is read.

• abs_positions – (max_count,) int32 extraction-window end position per slot (window is [pos - d_conv, pos)).

• real_count_gpu – int32[1] GPU tensor with the runtime real count.

• out – (max_count, conv_dim, d_conv) destination scratch slice (one layer).

• d_conv – conv window length (constexpr in the kernel).