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 module

max.nn.legacy.state_space

Python wrappers for the state_space Mojo kernel package.

The module layout mirrors max/kernels/src/state_space/. All ops here are registered as graph-compiler builtins, so no custom_extensions plumbing is needed. The legacy Mamba-1 eager wrappers intentionally stay in the Mamba architecture (their sole consumer) to keep max.nn.legacy from depending on max.functional.

Submodules

max.nn.legacy.state_space.kimi_delta	Python wrapper for the KDA (Kimi Delta Attention) recurrence op.

Classes

GatedDeltaReplayInputs	One layer's per-token inputs to the two state kernels.

Functions

causal_conv1d_varlen_fwd	Applies a slot-indexed varlen causal depthwise conv1d.
gated_delta_conv1d_fwd	Applies the causal conv1d pass.
gated_delta_conv1d_verify_fwd	Runs gated_delta_conv1d_fwd() in a speculative verify graph.
gated_delta_recurrence_fwd	Applies the gated delta recurrence pass.
gated_delta_recurrence_verify_ring_fwd	Runs the recurrence over a speculative verify window into a ring.
gated_delta_state_fold	Applies a verify's accepted ring records to the state pool.
gated_group_rmsnorm	Fuses the gated group-RMSNorm into a single dispatch.
kda_chunk	Runs the KDA chunk-parallel prefill recurrence, mutating state_pool.
kda_chunk_is_registered	Reports whether the current graph's kernel library has kda_chunk.
kda_chunk_supports_head_dims	Reports whether kda_chunk is compiled for this head-dim pair.
kda_decode	Runs the KDA recurrence, mutating state_pool in place.
mamba2_replay_fold	Commits each request's accepted verify rows to every Mamba-2 layer.
mamba2_ssd_chunk_scan_varlen_fwd_inplace	Performs the Mamba-2 SSD chunked-scan forward.
short_conv_ring	Returns x + conv(x) over a ragged batch.
verify_width_operand	Returns the operand the speculative state ops read the verify width from.