For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

mamba2_replay_fold

mamba2_replay_fold()

max.nn.legacy.state_space.mamba2_replay_fold(conv_pool, ssm_pool, conv_rows, ssm_rows, replay_bounds, projected, convolved, A, dt_bias)

source

Commits each request’s accepted verify rows to every Mamba-2 layer.

For each layer and request, the conv window becomes the last width - 1 accepted conv inputs and the SSM state folds in the accepted rows, all in one launch. The SSM state lands exactly where replaying those rows through mamba2_ssd_chunk_scan_varlen_fwd_inplace() would put it. Every request continues its live state. GPU only.

Parameters:

• conv_pool (BufferValue) – The [max_slots, conv_dim, width - 1] conv windows, updated in place.
• ssm_pool (BufferValue) – The [max_slots, nheads, head_dim, dstate] SSM states, updated in place.
• conv_rows (TensorValue) – The [num_layers, batch] uint32 conv-window rows.
• ssm_rows (TensorValue) – The [num_layers, batch] uint32 SSM-state rows.
• replay_bounds (TensorValue) – The [2 * batch] int32 accepted (start, end) rows of each request.
• projected (Sequence[TensorValue]) – Each layer’s in_proj output rows, [rows, nheads * head_dim + conv_dim + nheads]: the gate, the conv input and dt.
• convolved (Sequence[TensorValue]) – Each layer’s conv output rows, [rows, conv_dim]: the scan’s x, B and C.
• A (Sequence[TensorValue]) – Each layer’s [nheads] decay, already -exp(A_log).
• dt_bias (Sequence[TensorValue]) – Each layer’s [nheads] dt bias.

Return type:

None