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

Mojo package

state_space

State space model kernels for Mamba and related architectures.

This package provides GPU and CPU kernel implementations for the selective scan, causal conv1d, and SSD (state-space duality) operations used by Mamba, Mamba-2, and Gated DeltaNet models, including variable-length sequence variants for continuous-batching inference.

Both forward and update (decode-step) paths are included, and operation registrations in the *_ops modules wire these kernels into the MAX graph compiler.

Modules

causal_conv1d: Core Causal Conv1D Kernel Implementations.
causal_conv1d_ops: Causal Conv1D operation registrations for state space models.
gated_delta: Gated DeltaNet recurrence kernel for Qwen3.5: Pass 2 of two-pass prefill.
gated_delta_conv1d: Causal depthwise conv1d for the Gated DeltaNet two-pass prefill.
gated_group_rmsnorm: Fused gated group-RMSNorm for the Mamba-2 mixer (norm_before_gate=False).
mamba2_ssd_scan: Mamba-2 SSD (state-space duality) varlen prefill scan.
rms_norm_fused_residual: RMSNorm with fused residual connection for state space models.
rms_norm_fused_residual_ops: RMSNorm fused residual op registration for state space models.
selective_scan: Selective scan kernel implementations for Mamba SSM prefill and decode.
selective_scan_ops: Selective scan operation registrations for Mamba SSM.
varlen_causal_conv1d: Causal Conv1D with variable length sequence support (vLLM interface).
varlen_causal_conv1d_ops: Varlen Causal Conv1D operation registrations for Mamba SSM.
varlen_selective_scan: Variable-length selective scan kernels for Mamba SSM architecture.
varlen_selective_scan_ops: Varlen selective scan operation registrations for Mamba SSM.

Modules​

Modules