mamba3

Mamba-3 MIMO chunked forward and recurrent decode on NVIDIA GPUs, loadable through kernels. The reference baselines are the mamba3_MIMO_chunk_ref and mamba3_MIMO_step_ref programs from state-spaces/mamba, matched to 2.2e-6 in float32.

Upstream ships these paths as a pinned tilelang==0.1.8 kernel for the forward and a CUTLASS-DSL kernel for decode, and Mamba3(is_mimo=True) raises in its constructor when TileLang is absent rather than falling back. This is the same two operations as one compiled extension with no TileLang, TVM-FFI, CUTLASS-DSL or quack dependency.

Usage

from kernels import get_kernel

m3 = get_kernel("phanerozoic/mamba3", version=1, trust_remote_code=True)

# prefill: q, k are [B, S, R, Gqk, N]; v, z are [B, S, H, P]
y = m3.forward(q, k, v, q_bias, k_bias, mimo_v, mimo_o, angles, adt, dt, trap,
               z=z, mimo_z=mimo_z, D=D, chunk_size=16)

# decode: one token per call, state carried
st = m3.DecodeState(B, H, headdim=P, dstate=N, rank=R, nangles=Na)
y_t = st.step(q_t, k_t, v_t, q_bias, k_bias, mimo_v, mimo_o, angles_t,
              adt_t, dt_t, trap_t, z=z_t, mimo_z=mimo_z, D=D)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. angles are the raw per-step rotation rates; the cumulative sum is taken internally in float32.

API

norm_weight selects the fused head-wise RMSNorm pre-gate path. All three ops carry meta kernels, so the forward traces under torch.compile(fullgraph=True) without a graph break.

Method

The rank axis is flattened into the row index i = c*R + r, so a chunk tile is (C*R, N) against (C*R, P) and the forward is the SSD decomposition: chunk_state builds kv[n], state_passing scans it, chunk_scan computes q_inter @ states[n] + mask(Q Kᵀ) V + D v. Keys carry factor = γ_t + Δ_{t+1}(1 - λ_{t+1}), the single-SSD form of the trapezoidal recurrence, and the intra diagonal block is rescaled to γ_t. Contractions are register-blocked 4x4 on a 16x16 thread grid only where the tensor cores are not reachable. The float32 contractions split each operand into a tf32 head and tail and run hi*hi + hi*lo + lo*hi; the dropped lo*lo term is below the fp32 rounding of the sum, and a single tf32 pass would not hold the float32 contract. Rotation is applied while staging and never materialized.

Decode runs the recurrence directly, S ← αS + β(kprevᵀ vprev) + γ(kᵀv), one warp per p row with the rank axis in registers. The trapezoidal term reaches one step back, so the previous key and value are carried state.

mimo_v, mimo_o and mimo_z are (H, R, P) elementwise weights, so rank expansion is a broadcast multiply and the fold is a contraction over r alone. Rotation angles stay float32 at any input dtype, since they accumulate over the sequence.

On the bfloat16 path the operands are staged by ldmatrix and contracted by mma.sync.m16n8k16 written out directly, because nvcc's wmma never emits ldmatrix for bf16 and falls back to scalar loads and matrix-move fixups. The rank, the head dimension and the tile counts each warp owns are template arguments, so the accumulators index statically and stay in registers across the whole d loop and the epilogue, and the index arithmetic carries no integer division. Every shared tile is padded off a multiple of the 32 banks. The causal weight, the chunk decay and the D term all fold into the accumulator drain, D v_r riding the intra product on A's diagonal, so the float A tile never exists and the output tile reuses the space its operands occupied. The rank weights and the decay column are staged once per block, a rotary pair is one wide load, and the head axis varies fastest across the grid so that blocks sharing a chunk's q and k are resident together. The state scan issues its chunk loads four ahead of the carry that consumes them; the fma order is unchanged, and the writes land on the addresses just read.

Measured

RTX 6000 Ada, H=32, P=64, N=128, R=4, C=16. Forward against the torch reference program, which is what runs when TileLang is unavailable:

Decode against the eager step, with the state traffic it sustains:

At batch 64 the decode step is at this card's memory roofline. At batch 1 it is launch- and dependency-latency bound; capture it in a CUDA graph.

Against upstream's CuTe decode kernel on an H200, H=64, P=64, N=128, R=4, bfloat16, gated, eight tokens a call. Upstream's kernel takes q and k already rotated, so the comparison is against its Triton rotary step and its CuTe step together, which is what one call here does; the CuTe kernel alone is the column beside it. Both sides agree with the reference program to 5.1e-3, and with each other to 5.6e-3, which is bfloat16 rounding:

The two scale differently. This kernel is work-proportional in batch; upstream's pair is nearly flat from 1 to 256, and its rotary step is four fifths of it, so the crossover is near batch 32. Below it this kernel leads by up to 3.6x, above it upstream does, by 4x at batch 256.

Against upstream's TileLang forward rather than the reference program, on an RTX 3070 Ti Laptop at sm_86 in bfloat16. Both sides run in one process from one float32 master set, each cast to its own dtype contract:

On an H200 at sm_90, where upstream's float32 forward also fits its 168 KB dynamic shared-memory request and both dtypes can be compared:

The head-count row is part-dependent. This kernel is work-proportional, one block per (batch, head, chunk); upstream's unit of parallelism absorbs 32 heads, so it is underutilized below H=32 and saturated at it. On sm_86 that saturation is enough to take the row, and on sm_90 it is not. The float32 column leads at H=8 and trails at H=32, the one geometry either dtype loses on this part. Upstream's float32 forward requests 168,128 bytes of dynamic shared memory per block, and 180,992 at C = 64. The opt-in ceiling is 101,376 bytes on consumer and workstation Ampere and Ada. This kernel's float32 forward requests 39,104 bytes at the same geometry. The harness is benchmarks/bench_vs_upstream.py; TODO.md records the rest.

Backward

forward is an autograd function. The backward is two chunk-parallel passes around a reverse state scan, recomputing chunk intermediates rather than storing them: the first recovers O and takes the epilogue adjoint, the scan runs dSt[n] = dSt_part[n] + decay[n] dSt[n+1] in reverse, and the second produces dQ, dKf and dV before the rotation adjoint lands on q, k and the biases. Gradients cover the activations, every (H,R,P) rank weight, D, the norm weight, and the schedule itself: dt, adt, trap and the rotation rates.

It is reproducible. The grid is (batch, chunk) with the heads looped inside, so every gradient indexed by (b, s, ...) is owned by a single block and is accumulated in head order rather than by atomics; per-head parameters go to a per-block slot reduced in block order, and the two cross-chunk terms, the shifted half of the trapezoidal factor and the scan's decay adjoint, are staged and summed in a second ordered pass. No floating-point atomic appears anywhere in the backward source, and repeated runs return the forward and all 14 gradients bitwise identical.

The two chunk workspaces and the recomputed states are [B, H, Nc, N, P] each and dO is [B, H, Nc, C*R, P], so peak allocation is roughly the forward's plus that. At the geometry above:

Correctness

Every figure below is the worst case over that program's grid, measured on an RTX 6000 Ada against state-spaces/mamba's own reference programs.

Upstream's chunk reference omits the cos/sin padding its step reference performs, so its pairwise path is only defined at Na == N/2; partial rotation is checked against decode instead.

The rotation angle is a cumulative sum over the sequence and grows without bound, so it is evaluated with a range-reduced sincos rather than the approximate intrinsic, whose argument reduction would make the residual proportional to the angle. Agreement with the reference is flat in sequence length: 2.3e-6 at S = 2048 and 3.5e-6 at S = 32768, where the accumulated angle reaches 3,378 radians. The residual is set by the three-term tf32 split, which accumulates over the N/8 contraction tiles, so it is the figure to watch if N grows; MAMBA3_SCAN=fma returns the float32 scan to 1.4e-6 at a third less throughput.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+; float32 or bfloat16 values, float32 schedule and rotation angles.
• N even, Na <= N/2, mimo_rank <= 8, nheads divisible by the query-group count. Sequence length and chunk size are arbitrary; above C*R = 64 the forward switches to a row-block streamed intra term so nothing scales with (C*R)². C = 64 / mimo_rank remains the fastest setting, since the intra term is O(S*C) and grows faster than the state workspace shrinks.
• The backward has no streamed path, so its tiles stay resident and its shared memory grows as (C*R)² and C*R*P. Both entry points check the requirement against the device's opt-in maximum and raise naming the two figures rather than failing at launch. A 99 KB device admits C*R = 64 at headdim = 64; C*R = 128, or headdim = 128 at C*R = 64, needs the forward instead.
• Gradients are computed in float32 and returned in the input dtype, and the backward computes all of them whether or not each input requires grad.
• chunk_state and chunk_scan both contract on the tensor cores in both dtypes, float32 through the three-pass tf32 split and bfloat16 natively. The tensor-core kernels are instantiated for mimo_rank in {1, 2, 4, 8} and headdim in {16, 32, 64, 128}, and the scans additionally want chunk_size * mimo_rank == 64; any other geometry takes the FMA path. dispatch_paths reports which kernels a geometry selects, and MAMBA3_SCAN=fma demotes the float32 scan to FMA for an A/B in one process.
• Packaged for compute capabilities 8.0, 8.6, 8.9, 9.0, 10.0 and 12.0. The whole suite runs on 8.6 and 8.9; the others are build targets.

References

Lahoti, Li, Chen, Wang, Bick, Kolter, Dao, Gu, "Mamba-3: Improved Sequence Modeling using State Space Principles" (ICLR 2026, arXiv:2603.15569); Dao and Gu, the SSD chunked scan of Mamba-2; state-spaces/mamba as the reference implementation.

License

Apache-2.0.