# mamba-rs
Mamba SSM and Mamba-3 SISO in Rust, on the CPU and on NVIDIA GPUs.
Inference and training for both architectures, with the crate's own CUDA
kernels compiled at run time through NVRTC. No Python, no C++ build step,
no framework dependency; the GPU path links only the CUDA driver API and
cuBLAS.
## What's new in 0.7.0
0.7.0 is a big performance release: new deterministic GEMM kernels, now
the default, a pass over every Mamba kernel around them, and the two
scan and weight-gradient routes moved to their faster families. A
training step of the release shapes runs 1.13 to 1.63 times faster than
0.6.9 on an RTX 6000 Ada.
- **Every GPU context has a `GemmMode`.** `Deterministic` (the default)
runs the crate's own kernels and never calls cuBLAS; `CublasFast` and
`CublasPedantic` select cuBLAS explicitly. In the deterministic mode the
same inputs give the same bits run after run, from an eager launch and
from a captured graph, and for serving the same bits for a row at any
batch size. To keep the exact numbers 0.6.9 produced, construct with
`GemmMode::CublasPedantic`.
- **New deterministic training kernels** (the Triad family: forward, weight
gradient and input gradient) and **new deterministic serving kernels**
(the Inference family, formerly `Fixed`), measured kernel by kernel on an
RTX 6000 Ada and an RTX 5090 against cuBLAS Fast and cuBLAS Pedantic, and
against the 0.6.9 kernels: on the large training and serving shapes the
new kernels are 1.26 to 1.43 times faster on average, single kernels up
to 3 times.
- **Deterministic TF32**, a new precision setting for f32: the products run
on the tensor cores in TF32 with one fixed rounding and a fixed
summation order, the accumulation stays f32, and the bits are
reproducible like the rest of the deterministic mode. 0.6.9 had no such
setting. Stream-K weight-gradient kernels serve the deep reductions by
default.
- **Explicit-mode constructors** beside every environment-reading one, and
a recorded numeric route on every captured graph.
The full list is in [CHANGELOG.md](CHANGELOG.md); the mode guide is
[docs/gemm-modes.md](docs/gemm-modes.md); the numbers are in
[docs/determinism-benchmarks.md](docs/determinism-benchmarks.md).
What comes next: the scan, convolution and norm kernels around the GEMMs
are the next target, the GEMM kernels keep moving toward cuBLAS Fast in
the 0.7.x releases, and the architectures on the portable kernels today
(SM80, SM86, Hopper, datacenter Blackwell, CC 12.1) get measured kernels
in later releases.
## Features
- Mamba SSM (Gu and Dao, 2023) and Mamba-3 SISO (Lahoti et al., 2026).
- CPU and GPU paths for both, with cross-path parity tests on shared
weights.
- Inference and training: full backward pass through the recurrent SSM
state, AdamW, CUDA Graph capture for inference steps, prefill and
training steps.
- f32, bf16 and f16 storage through one `WeightDtype` selector; every
kernel accumulates in f32.
- Three GEMM modes, deterministic by default. See below.
- Bring-your-own-loss training: `trainer.forward()` returns the full
temporal output on the host, any loss gradient computed in Rust goes into
`trainer.backward_step()`, bit-identical to the fused `step()`.
- Full-sequence CPU prefill for both architectures, and a GPU prompt
prefill for Mamba-3 with a captured-graph twin.
- Deterministic data-parallel training over NCCL (`dist`): one reduction
per optimizer step in a fixed order, so the bits do not depend on the
transport or the GPU permutation. The fixed-order path is pinned against
an in-process oracle on one GPU; the NCCL collective path was validated
on two RTX 5090s.
- Bit-continuous resume: optimizer state and the carried recurrence export
and import, so a resumed run lands where the unbroken run would.
- State dimensions up to 256, sized at compile time from the config.
- HuggingFace safetensors loader for Mamba SSM checkpoints (130m, 370m,
1.4b and 2.8b validated).
- MSRV 1.97.
## Cargo features
| *(default)* | pure-Rust scalar GEMM | correctness work only; much slower than a BLAS |
| `gemm-blas` | the [`gemm`] crate's BLAS-class CPU GEMM (with rayon) | any serious CPU use |
| `accelerate` | Apple Accelerate GEMM (macOS) | macOS deployments |
| `cuda` | GPU inference and training (NVRTC-compiled kernels) | needs the CUDA toolkit |
| `hf` | safetensors and HuggingFace checkpoint loaders | LM checkpoints |
| `cli` | the `mamba-generate` binary (tokenizers and hf-hub) | text generation from the command line |
| `nccl` | data-parallel transport (pinned NCCL binding) | multi-GPU training |
| `qualification` | the hardware and toolkit instruments under `tools/qualification/` | maintainers measuring kernels on a chosen board |
| `cuda-cublaslt-qualification` | cuBLASLt in the vendor-comparison harness | maintainers only; production routing does not use cuBLASLt |
## GEMM modes
| reproducible results: the same bits run to run, eager or graph, and for serving at any batch size | `GemmMode::Deterministic` (default) |
| the fastest vendor path, TF32 permitted for f32 | `GemmMode::CublasFast` |
| the numbers 0.6.9 produced, or the vendor's most careful f32 accumulation as a reference | `GemmMode::CublasPedantic` |
Storage precision and mode are separate choices: `WeightDtype` decides how
the weights are stored, `GemmMode` decides who multiplies. Every GPU entry
point has a plain constructor that reads `MAMBA_RS_GEMM_MODE`
(`deterministic`, `cublas-fast`, `cublas-pedantic`; default
`deterministic`) and a `*_with_mode` twin that takes the mode as its last
argument and ignores the environment.
```rust
use mamba_rs::mamba_ssm::gpu::GemmMode;
use mamba_rs::mamba_ssm::gpu::context::GpuCtx;
use mamba_rs::mamba_ssm::gpu::device::GpuDevice;
let device = GpuDevice::new(0)?;
let ctx = GpuCtx::new_with_mode(&device, GemmMode::Deterministic)?;
assert_eq!(ctx.gemm_mode(), GemmMode::Deterministic);
ctx.set_gemm_mode(GemmMode::CublasPedantic)?; // refused while a graph is being captured
```
Inside the deterministic mode a model context uses the Inference kernels
and a trainer the Triad kernels; tensor cores, deterministic TF32 and
stream-K are settings on the context. What each mode guarantees, what it
does not, the environment variables and the architecture coverage are in
[docs/gemm-modes.md](docs/gemm-modes.md).
## Use cases and API choice
### Reinforcement learning and small custom models
Latency-critical, typically `d_model` up to 256, often batch 1 for actor
rollouts. Both CPU and GPU paths apply; at these sizes the CPU step and the
GPU step are within a factor of two of each other, so the choice depends
on where the rest of the program lives.
- Inference: `mamba_step` (CPU) or `GpuMambaBackbone::step` (GPU)
- Training: `parallel_mamba_forward` and `parallel_mamba_backward` (CPU,
rayon-parallel over the batch) or `MambaTrainer::step` (GPU, one
captured graph for forward, backward, AdamW and sync)
CPU training is practical where GPU launch overhead dominates (`d_model`
up to 128, batch up to 8); GPU training scales well from batch 32.
### Large language models
Throughput-critical, `d_model` from 768, token-by-token decoding of a
HuggingFace checkpoint. GPU only in practice: a 2.8b model on the CPU
decodes at single-digit tokens per second whatever the implementation.
- Inference: `GpuMambaLM::from_hf_with_dtype` and `generate`
- Fine-tuning: `MambaTrainer::new_full` on the HuggingFace backbone
weights (Mamba SSM only; no public Mamba-3 SISO checkpoint exists)
The CPU `MambaLM` path runs end to end but exists for CPU-versus-GPU
parity tests, not for serving.
### Sequence classification, embeddings and custom heads
Whole-sequence reads with a caller-defined loss (document classifiers,
distillation, contrastive embeddings). GPU training uses the
forward/backward split; CPU serving uses the prefill.
- Training: `MambaTrainer::forward`, a host-side loss,
`MambaTrainer::backward_step` (see `examples/custom_loss.rs`)
- CPU serving: `MambaBackbone::forward_prefill` and
`forward_mamba3_backbone_prefill` (see `examples/cpu_prefill.rs`), or
`prefill_batch` and `prefill3_batch` for batches of sequences
### Sharing weights across paths
All paths consume the same `MambaWeights` or `Mamba3Weights`. A training
run's `MambaTrainer::snapshot_master()` loads directly into
`GpuMambaBackbone`, `GpuMambaLM` or the CPU `MambaBackbone` without
conversion.
## Quick start (CPU)
### Mamba SSM
```rust
use mamba_rs::{MambaConfig, MambaState, MambaStepScratch, MambaWeights, mamba_step};
let cfg = MambaConfig::default();
let weights = MambaWeights::init(&cfg, input_dim, 42);
let mut state = MambaState::zeros(cfg.n_layers, cfg.d_inner(), cfg.d_state, cfg.d_conv);
let mut scratch = MambaStepScratch::new(&cfg);
let mut output = vec![0.0f32; cfg.d_model];
mamba_step(&input, &mut output, &weights, &mut state.layers, &mut scratch, &cfg, input_dim);
```
### Mamba-3
```rust
use mamba_rs::mamba3_siso::config::Mamba3Config;
use mamba_rs::mamba3_siso::cpu::inference::{Mamba3StepScratch, mamba3_step};
use mamba_rs::mamba3_siso::state::Mamba3State;
use mamba_rs::mamba3_siso::weights::Mamba3Weights;
let cfg = Mamba3Config::default();
let weights = Mamba3Weights::init(&cfg, input_dim, 42);
let mut state = Mamba3State::zeros(&cfg);
let mut scratch = Mamba3StepScratch::new(&cfg);
let mut output = vec![0.0f32; cfg.d_model];
mamba3_step(&mut output, &input, &mut scratch, &weights, &mut state.layers, &cfg);
```
Enable `gemm-blas` (or `accelerate` on macOS) for any CPU work beyond a
correctness check.
## Quick start (GPU inference)
```toml
[dependencies]
mamba-rs = { version = "0.7", features = ["cuda"] }
```
```rust
use mamba_rs::gpu::inference::GpuMambaBackbone;
use mamba_rs::mamba_ssm::gpu::GemmMode;
use mamba_rs::WeightDtype;
let mut gpu = GpuMambaBackbone::new_with_dtype_and_mode(
0, &weights, cfg, input_dim, batch, WeightDtype::Bf16, GemmMode::Deterministic,
)?;
gpu.capture_graph()?; // optional: one graph launch per step
gpu.step(&input, &mut output)?;
gpu.reset()?;
```
`GpuMamba3Backbone` has the same constructor. `new_with_dtype` reads
`MAMBA_RS_GEMM_MODE` instead.
### HuggingFace LM inference
```rust
use mamba_rs::module::gpu_lm::GpuMambaLM;
use mamba_rs::module::sample::SampleParams;
use mamba_rs::mamba_ssm::gpu::GemmMode;
use mamba_rs::WeightDtype;
use std::path::Path;
let mut lm = GpuMambaLM::from_hf_with_dtype_and_mode(
Path::new("./mamba-130m-hf"), 0, WeightDtype::Bf16, GemmMode::Deterministic,
)?;
lm.capture_graph()?;
let tokens = lm.generate(&[1, 2, 3, 4, 5], &SampleParams::default())?;
```
bf16 against f32 on the four `state-spaces/mamba-*-hf` checkpoints: 15 of
15 greedy tokens match and the KL divergence of the final logits is at
most 1.1e-3 (`tests/gpu_bf16_parity.rs`). Batch 1 against batch 32 on the
same prompt in the deterministic mode: KL about 8e-11
(`tests/extreme_edge_coverage.rs`), the batch invariance of the Inference
kernels measured through a whole model; both tests need a local checkpoint
and run with `--ignored`.
## Quick start (GPU training)
`MambaTrainer` and `Mamba3Trainer` run forward, backward, AdamW and the
master-weight sync behind one `step()` call. The `WeightDtype` argument
selects the f32 or the mixed bf16/f16 engine.
```rust
use mamba_rs::mamba_ssm::gpu::GemmMode;
use mamba_rs::mamba_ssm::gpu::trainer::{MambaTrainer, TrainSessionCfg};
use mamba_rs::WeightDtype;
let session = TrainSessionCfg {
input_dim,
batch: 2,
seq_len: 64,
lr: 3e-4,
weight_decay: 1e-2,
};
let mut trainer = MambaTrainer::new_full_with_mode(
/* gpu_ordinal */ 0,
&cpu_weights, cfg, session,
WeightDtype::Bf16,
GemmMode::Deterministic,
)?;
trainer.capture_graph()?; // optional; one graph launch per step after this
let metrics = trainer.step(&input, &d_temporal_upstream)?;
// metrics.step, metrics.graph_replayed, metrics.loss_scale (f16), metrics.overflow_skipped (f16)
let master = trainer.snapshot_master()?; // CPU-side MambaWeights for checkpointing
```
f16 training activates the dynamic loss scaler automatically;
`metrics.loss_scale` and `metrics.overflow_skipped` report its state.
### Custom losses: the forward/backward split
The fused `step()` needs the loss gradient up front; the split lets you
compute it from the actual forward output:
```rust
use mamba_rs::mamba_ssm::gpu::trainer::BackwardOpts;
let mut temporal = vec![0.0f32; batch * seq_len * cfg.d_model];
trainer.forward(&input, &mut temporal)?; // full temporal output, f32, on the host
let d_temporal = my_loss_grad(&temporal); // any host-side loss
let m = trainer.backward_step(
&d_temporal,
BackwardOpts::default().with_clip_max_norm(1.0),
)?; // backward, clipping, AdamW
// gradient accumulation: .with_accumulate_only(true) on the non-applying
// micro-batches; the fused step() refuses while a window is open.
```
The split always runs eagerly, because a caller-side loss cannot live
inside a captured graph, and it is bit-identical to the fused `step()`.
`examples/custom_loss.rs` is a complete training loop.
## Quick start (CPU prefill)
```rust
use mamba_rs::inference::PrefillMode;
use mamba_rs::module::MambaBackbone;
let backbone = MambaBackbone::init(cfg, input_dim, 42);
let mut state = backbone.alloc_state();
let mut scratch = backbone.alloc_prefill_scratch(seq_len);
let mut out = vec![0.0f32; seq_len * backbone.config().d_model];
// One batched GEMM pass over the whole prompt instead of T step calls.
backbone.forward_prefill(&prompt, &mut out, &mut state, &mut scratch,
seq_len, PrefillMode::Parallel);
// `out` holds the post-norm output at every position; `state` is
// positioned after the prompt, so forward_step continues from it.
```
Mamba-3 has the same surface (`forward_mamba3_backbone_prefill` and
`Mamba3PrefillScratch`).
## Serialization
```rust
use mamba_rs::serialize;
serialize::save(Path::new("model.safetensors"), backbone.weights(), cfg, input_dim)?;
let (weights, cfg, input_dim) = serialize::load(Path::new("model.safetensors"))?;
// Mamba-3
use mamba_rs::mamba3_siso::serialize::{save_mamba3, load_mamba3};
save_mamba3(Path::new("m3.safetensors"), &weights, &cfg, input_dim)?;
let (weights, input_dim) = load_mamba3(Path::new("m3.safetensors"), &cfg)?;
```
## Performance
Measured on an RTX 6000 Ada (SM89, driver 595.45.04) and an RTX 5090
(CC 12.0, driver 595.84 for the serving tables and 595.58.03 for the
training tables), both on CUDA 13.2. Speedups are cuBLAS time
divided by mamba-rs time; above 1.0 the deterministic kernel is faster.
Exact f32 is compared with cuBLAS Pedantic, which performs the same
arithmetic; bf16 and f16 with cuBLAS Fast on the native half-precision
tensor-core kernels; deterministic TF32 with cuBLAS Fast TF32.
Serving kernels (the Inference family, five shapes, bias off and on,
geometric mean, eager path):
| BF16 → BF16 | cuBLAS Fast | 1.19× | 1.24× |
| F16 → F16 | cuBLAS Fast | 1.17× | 1.23× |
| BF16 → F32 | cuBLAS Fast | 0.83× | 1.29× |
| F32, deterministic TF32 | cuBLAS Fast TF32 | 0.90× | 1.13× |
| F32, exact | cuBLAS Pedantic | 1.00× | 1.08× |
Training kernels (the Triad family, the large shapes, geometric mean,
eager path; the small d128 shapes are launch-bound and slower than
cuBLAS on both boards):
| BF16 | cuBLAS Fast | 1.07× | 1.06× |
| F16 | cuBLAS Fast | 1.08× | 1.06× |
| F32, deterministic TF32 | cuBLAS Fast TF32 | 0.84× | 1.07× |
| F32, exact | cuBLAS Pedantic | 0.94× | 1.17× |
Whole training step on the RTX 6000 Ada, 0.7.0 against 0.6.9, same shapes
and settings in both trees (bf16 with tensor cores, ms per step): d128
2.35 → 1.94, d256 9.54 → 7.66, d768 22.17 → 13.60, d1536 13.07 → 9.45.
The step is dominated by the scan and the other non-GEMM kernels, so the
whole-step gain is smaller than the kernel gain.
The per-kernel tables for both boards, the old-versus-new kernel
comparison, the whole-model comparison and the measurement protocol are in
[docs/determinism-benchmarks.md](docs/determinism-benchmarks.md). The
whole-model tables of earlier releases stay on the
[Mamba SSM](docs/mamba1-benchmarks.md) and
[Mamba-3 SISO](docs/mamba3-benchmarks.md) benchmark pages with their
release, board and comparator labels.
## Testing
Every target is declared in `Cargo.toml` and has a lane in
`qual/lanes.toml`; a host test keeps the two in step.
- **Regressions** (`tests/`): bit parity within a numeric route (eager and
graph, run and run), tolerance parity across routes (CPU and GPU, f32
and half, scalar and tensor-core), gradient checks, source and dispatch
contracts, batch invariance and determinism gates.
`cargo test --release --features cuda` runs the CUDA gate;
`cargo test --no-default-features` runs the host part on any machine.
Tests marked `contract` contain arms that need a checkpoint or a
specific board and run with `-- --ignored`; tests marked `record`
write evidence and never run automatically.
- **Benches** (`benches/`): timing instruments with no verdict,
`cargo bench --features cuda --bench <name>`, optionally followed by
`-- <instrument>`.
- **Qualification tools** (`tools/qualification/`): hardware, toolkit and
inventory instruments that need a specific board, built with
`--features "cuda hf qualification"` and run by name with `-- --ignored`.
`qual/run.sh <lane>` runs or lists one lane;
[docs/release-qualification.md](docs/release-qualification.md) describes
the release order.
## Documentation
For users:
- [GEMM modes](docs/gemm-modes.md): the three modes, which to choose, how
to set them, what is guaranteed, environment variables, architecture
coverage
- [GEMM benchmarks](docs/determinism-benchmarks.md): kernel-by-kernel
timings on both boards against cuBLAS Fast and Pedantic, the 0.6.9
comparison, the protocol
- [Mamba SSM architecture](docs/mamba1-architecture.md) and
[benchmarks](docs/mamba1-benchmarks.md)
- [Mamba-3 SISO architecture](docs/mamba3-architecture.md) and
[benchmarks](docs/mamba3-benchmarks.md)
- Rustdoc: `GemmMode`, `GpuCtx::new_with_mode`, `GpuCtx::set_gemm_mode`
and the `*_with_mode` constructors carry the API contract
For contributors:
- [Performance playbook](docs/performance-playbook.md): how kernels are
measured, changed and admitted
- [Release qualification](docs/release-qualification.md): test lanes,
package inspection, evidence and the release order
## Roadmap
- Multi-GPU inference for models larger than one device (pipeline
sharding), beside the data-parallel training that ships now.
- Reduced-precision tiers (fp8, int8) under the same bit discipline as
the f32, bf16 and f16 paths.
- The Mamba-2 generation beside Mamba-1 and Mamba-3, with the same
determinism and testing discipline.
## Citation
```bibtex
@article{mamba,
title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
author={Gu, Albert and Dao, Tri},
journal={arXiv preprint arXiv:2312.00752},
year={2023}
}
@inproceedings{mamba3,
title={Mamba-3: Improved Sequence Modeling using State Space Principles},
author={Lahoti, Aakash and Li, Kevin Y. and Chen, Berlin and Wang, Caitlin and Bick, Aviv and Kolter, J. Zico and Dao, Tri and Gu, Albert},
booktitle={International Conference on Learning Representations},
year={2026}
}
```
## License
Dual-licensed under MIT or Apache-2.0.