---
library_name: kernels
license: apache-2.0
---

# 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

```python
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

| Symbol | Purpose |
|---|---|
| `forward(q, k, v, ..., chunk_size, norm_weight)` | chunked forward, `[B, S, H, P]` out, differentiable |
| `DecodeState(batch, nheads, headdim, dstate, rank, nangles)` | carried decode state |
| `DecodeState.step(...)` | one decode token, `[B, H, P]` out |
| `cumulative_angles(angles, dt)`, `chunk_decay(adt, chunk_size)` | schedule helpers |
| `dispatch_paths(mimo_rank, dstate, headdim, chunk_size, bf16)` | which kernels a geometry selects |

`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:

| seqlen | kernel | torch | speedup | peak memory |
|---|---|---|---|---|
| 512 | 0.331 ms | 3.312 ms | 10.02x | 62 vs 375 MB |
| 1024 | 0.586 ms | 7.860 ms | 13.42x | 131 vs 741 MB |
| 2048 | 1.531 ms | 17.609 ms | 11.50x | 254 vs 1472 MB |
| 4096 | 3.186 ms | 40.460 ms | 12.70x | 499 vs 2935 MB |

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

| batch | kernel | eager | speedup | state |
|---|---|---|---|---|
| 1 | 0.0245 ms | 0.754 ms | 30.8x | 86 GB/s |
| 4 | 0.0398 ms | 0.749 ms | 18.8x | 211 GB/s |
| 16 | 0.0665 ms | 0.773 ms | 11.6x | 505 GB/s |
| 64 | 0.1659 ms | 2.030 ms | 12.2x | 809 GB/s |
| 256 | 0.9697 ms | 7.492 ms | 7.7x | 554 GB/s |

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:

| batch | kernel | rotary + cute | cute alone | ratio |
|---|---|---|---|---|
| 1 | 0.812 ms | 2.901 ms | 0.646 ms | 3.57x |
| 4 | 0.813 | 2.885 | 0.645 | 3.55x |
| 16 | 1.180 | 2.894 | 0.638 | 2.45x |
| 64 | 3.563 | 2.894 | 0.742 | 0.81x |
| 256 | 12.738 | 3.051 | 1.973 | 0.24x |

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:

| S | H | R | C | kernel | tilelang | ratio |
|---|---|---|---|---|---|---|
| 512 | 8 | 4 | 16 | 0.180 ms | 0.398 ms | 2.21x |
| 1024 | 8 | 4 | 16 | 0.297 ms | 0.789 ms | 2.66x |
| 2048 | 8 | 4 | 16 | 0.544 ms | 1.628 ms | 2.99x |
| 4096 | 8 | 4 | 16 | 1.054 ms | 3.316 ms | 3.15x |
| 2048 | 8 | 1 | 64 | 0.252 ms | 0.514 ms | 2.04x |
| 2048 | 32 | 4 | 16 | 2.189 ms | 1.862 ms | 0.85x |

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:

| S | H | R | C | kernel | tilelang | ratio | | kernel f32 | tilelang f32 | ratio |
|---|---|---|---|---|---|---|---|---|---|---|
| 512 | 8 | 4 | 16 | 0.114 ms | 0.204 ms | 1.79x | | 0.183 ms | 0.281 ms | 1.54x |
| 1024 | 8 | 4 | 16 | 0.106 | 0.383 | 3.62x | | 0.308 | 0.554 | 1.80x |
| 2048 | 8 | 4 | 16 | 0.181 | 0.754 | 4.16x | | 0.598 | 1.140 | 1.91x |
| 4096 | 8 | 4 | 16 | 0.378 | 1.641 | 4.34x | | 1.236 | 2.547 | 2.06x |
| 2048 | 32 | 4 | 16 | 0.712 | 1.035 | 1.45x | | 2.217 | 1.311 | 0.59x |
| 2048 | 8 | 1 | 64 | 0.100 | 0.267 | 2.68x | | 0.211 | 0.365 | 1.73x |

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:

| seqlen | forward | forward + backward |
|---|---|---|
| 512 | 55 MB | 129 MB |
| 1024 | 109 MB | 258 MB |
| 2048 | 218 MB | 513 MB |
| 4096 | 435 MB | 1025 MB |

## 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.

| program | checks | worst |
|---|---|---|
| `tests/test_mamba3.py` decode | ranks 1 to 8, grouped-query, both rotation fractions | 3.0e-7 |
| `tests/test_mamba3.py` forward | against `mamba3_MIMO_chunk_ref` | 2.5e-6 |
| `tests/test_mamba3.py` forward vs decode | partial rotation | 1.3e-6 |
| `tests/test_mamba3.py` ragged | sequence length not a multiple of the chunk | 3.3e-6 |
| `tests/test_mamba3.py` fused norm | against the upstream fused-norm reference | 3.6e-7 |
| `tests/test_mamba3.py` bfloat16 | against the float32 kernel | 9.4e-3 |
| `tests/test_mamba3.py` long sequences | `S` to 32,768, accumulated angle to 3,378 radians | 3.5e-6 |
| `tests/test_bwd.py` | 15 raw gradients against autograd through the reference | 7.6e-7 |
| `tests/test_autograd.py` | the wrapped function end to end, cumsum adjoints included | 6.4e-7 |
| `tests/test_bwd_bf16.py` | the bfloat16 backward against the float32 one | 6.8e-3 |
| `tests/test_bigchunk.py` | the streamed intra path up to `C*R = 256` | 6.1e-7 |
| `ref/test_ref.py` | this repo's step reference against upstream's | 2.3e-7 |
| `tests/test_dispatch.py` | all five launcher paths against the step reference | 1.5e-2 |
| `tests/test_determinism.py` | forward and 14 gradients over 3 repeats | bitwise |
| `tests/test_fixes.py` | meta kernels, `fullgraph=True`, the backward memory guard | pass |

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.