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 module

mamba2_ssd_scan

Mamba-2 SSD (state-space duality) varlen prefill scan.

Target hardware family: B200 (sm_100), but the kernel is architecture-agnostic (plain elementwise + SIMD-over-dstate recurrence, no MMA/TMA/swizzle). bf16 in/out, fp32 accumulate, states fp32.

This is the varlen (ragged, query_start_loc) Mamba-2 SSD forward prefill op, matching mamba_chunk_scan_combined semantics for the NemotronHMamba2Mixer (Nemotron-H). It differs from the Mamba-1 ops in varlen_selective_scan.mojo in three ways that the math requires:

A is a per-head SCALAR (nheads,) (shared across all head_dim channels and all dstate), not a per-channel (dim, dstate) diagonal.
B/C are GROUPED (total_len, ngroups, dstate); nheads/ngroups heads share each group (group_id = h // (nheads // ngroups)).
dt is per-head (total_len, nheads) + per-head dt_bias (nheads,), broadcast across head_dim; softplus is applied to dt + dt_bias.

Reference math (the source of truth for parity) is the HF torch_forward SSD path (segment_sum / reshape_into_chunks / chunk-state recurrence). The SSD chunked scan is a parallelism reformulation of the linear recurrence below; in fp32 they are numerically equivalent. This first implementation carries the state sequentially per (head, head_dim) channel (one thread per channel, SIMD over dstate), mirroring the tiling style of varlen_selective_scan_fwd_gpu. A chunk-tiled rewrite (segment_sum + matmul) is a follow-up perf slice; this op is gated on CORRECTNESS first.

Per-token recurrence (per head h, head_dim channel p, group g):

dt_t     = softplus(dt[t, h] + dt_bias[h])            # scalar per (t, h)
dA_t     = exp(A[h] * dt_t)                           # scalar per (t, h)
state_n  = state_n * dA_t + dt_t * B[t, g, n] * x[t, h, p]   # vector over n
y[t,h,p] = sum_n C[t, g, n] * state_n  +  D[h] * x[t, h, p]

State resets to zero (or to initial_states) at each query_start_loc boundary -- no cross-sequence bleed. final_states (batch, nheads, head_dim, dstate) is written at each sequence end.

`comptime` values

`LOG2E`

comptime LOG2E = 1.4426950408889634

`MAX_DSTATE`

comptime MAX_DSTATE = 256

`Strides1D`

comptime Strides1D = IndexList[Int(1)]

`Strides2D`

comptime Strides2D = IndexList[Int(2)]

`Strides3D`

comptime Strides3D = IndexList[Int(3)]

`Strides4D`

comptime Strides4D = IndexList[Int(4)]

Functions

mamba2_ssd_chunk_scan_varlen_fwd_cpu: CPU reference for the Mamba-2 SSD varlen prefill scan.
mamba2_ssd_chunk_scan_varlen_fwd_gpu: GPU kernel: Mamba-2 SSD varlen prefill scan, one thread per (head, channel).
mamba2_ssd_chunk_scan_varlen_fwd_inplace_cpu: CPU reference: Mamba-2 SSD varlen scan with in-place SSM-pool write.
mamba2_ssd_chunk_scan_varlen_fwd_inplace_gpu: GPU kernel: Mamba-2 SSD varlen prefill scan with in-place SSM-pool write.
mamba2_ssd_chunk_scan_varlen_fwd_inplace_gpu_dstate_split: GPU kernel: Mamba-2 SSD varlen in-place scan, cooperative DSTATE-split.

comptime values​

LOG2E​

MAX_DSTATE​

Strides1D​

Strides2D​

Strides3D​

Strides4D​

Functions​