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_ssd_chunk_scan_varlen_fwd_inplace

mamba2_ssd_chunk_scan_varlen_fwd_inplace()

max.nn.legacy.state_space.mamba2_ssd_chunk_scan_varlen_fwd_inplace(x, dt, A, B, C, D, dt_bias, ssm_pool, query_start_loc, has_initial_state, cache_indices, update_state=None)

source

Performs the Mamba-2 SSD chunked-scan forward.

Writes final states back into the SSM pool in place.

Writes final states directly into ssm_pool[cache_indices[b], ...] in place, so the graph never round-trips the whole pool.

The kernel picks its algorithm from the launch shape: long prefills run the chunked (tensor-core) scan, decode, speculative verify and short launches run the token-serial scan. Both read and write the same pool layout.

Parameters:

• x (TensorValue) – The [total_len, nheads, head_dim] SSM input (model dtype).
• dt (TensorValue) – The [total_len, nheads] per-head time deltas (model dtype).
• A (TensorValue) – The [nheads] per-head scalar (model dtype; already -exp(A_log)).
• B (TensorValue) – The [total_len, ngroups, dstate] grouped input proj (model dtype).
• C (TensorValue) – The [total_len, ngroups, dstate] grouped output proj (model dtype).
• D (TensorValue) – The [nheads] skip connection (model dtype; empty to disable).
• dt_bias (TensorValue) – The [nheads] dt bias (model dtype; empty to disable softplus).
• ssm_pool (BufferValue) – The [max_slots, nheads, head_dim, dstate] mutable state pool (fp32; bf16 on Apple GPUs — storage only, the scan accumulates in fp32). Read at ssm_pool[cache_indices[b]] when has_initial_state[b] is true; written in-place with final state.
• query_start_loc (TensorValue) – The [batch + 1] int32 cumulative sequence lengths.
• has_initial_state (TensorValue) – The [batch] bool, whether to load initial state for each sequence (empty to disable).
• cache_indices (TensorValue) – The [batch] uint32 slot indices into ssm_pool, or one [1, batch] row of a per-layer table.
• update_state (TensorValue | None) – The [batch] bool, whether each sequence writes its final state back to ssm_pool[cache_indices[b]]. None writes every sequence’s, as before the flag existed. A sequence with the flag false leaves its slot untouched, so the scan reads the pool without advancing it: a speculative verify can run on the live state and a later replay still starts from the pre-verify state.

Returns:

y, the [total_len, nheads, head_dim] output (model dtype). ssm_pool is mutated in place.

Return type:

TensorValue