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.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 from depending on max.experimental.functional.

Submodules

• max.nn.state_space.kimi_delta

Functions

causal_conv1d_varlen_fwd	Applies a slot-indexed varlen causal depthwise conv1d for prefill and decode.
gated_delta_conv1d_fwd	Applies the causal conv1d pass, mutating a slot-indexed conv-state pool in place.
gated_delta_recurrence_fwd	Applies the gated delta recurrence pass, mutating a slot-indexed state pool in place.
gated_group_rmsnorm	Fuses the gated group-RMSNorm (HF Zamba2RMSNormGated with norm_before_gate=False) into a single dispatch.
kda_decode	Runs the KDA recurrence, mutating state_pool in place.
mamba2_ssd_chunk_scan_varlen_fwd	Performs the Mamba-2 SSD chunked-scan forward for prefill and decode.
mamba2_ssd_chunk_scan_varlen_fwd_inplace	Performs the Mamba-2 SSD chunked-scan forward, writing final states back into the SSM pool in place.
short_conv_ring_commit	Writes each sequence's last ring_len rows of x into its slot of ring.
short_conv_ring_fwd	Returns x + conv(x) over a ragged batch; reads ring, writes nothing.