mamba-rs

Mamba SSM and Mamba-3 SISO in Rust with optional CUDA GPU acceleration. Inference and training for both, with custom CUDA kernels.

Pure Rust + CUDA. Kernels compile at runtime via NVRTC.

Features

• Two architectures — Mamba SSM (Gu & Dao, 2023) and Mamba-3 SISO (Lahoti et al., ICLR 2026).
• CPU + GPU — both paths exposed, with a cross-path parity test on shared weights.
• Inference + training — full backward pass with BPTT through the recurrent SSM state; AdamW optimizer; CUDA Graph capture for both.
• f32 / bf16 / f16 — a single WeightDtype selector at construction. Compute stays f32 (upcast-in-kernel, f32 accumulators) regardless of storage dtype.
• Deterministic inference & training (opt-in) — MAMBA_RS_BATCH_INVARIANT=1 / ctx.set_batch_invariant(true) routes every GEMM through custom deterministic kernels (kernels/sgemm_bi.cu, gemm_batch_invariant.cu): inference logits are bit-identical across batch sizes (KL ≈ 1e-11), and f32 / bf16 / f16 training is bit-identical across runs. Default path is cuBLAS for maximum throughput.
• Tensor-core deterministic tier (opt-in) — MAMBA_RS_BI_TENSOR_CORES=1 / ctx.set_bi_tensor_cores(true) on top of the flag above swaps the training GEMM triad for mma.sync tensor-core kernels: still fully deterministic (own numeric contract), at-or-near cuBLAS parity even on d128/d256 models and faster than cuBLAS from d_model ≥ 768 (0.70× of PEDANTIC per step at d1536 bf16).
• Bring-your-own-loss training split — trainer.forward() returns the full batch * seq_len * d_model post-norm_f temporal output on the host; compute ANY loss gradient in plain Rust and feed it to trainer.backward_step() (global-norm clipping, exact gradient accumulation, LR schedules via set_lr, reference-faithful AdamW no-decay groups). Bit-identical to the fused step() — both compose the same eager phase bodies. Same API on Mamba3Trainer.
• Full-sequence CPU prefill — forward_prefill runs a whole prompt/page through the training forward's batched-SGEMM pipeline (no activation tape) instead of T per-step dispatches, then hands the recurrent state to the step path (prefill-then-decode). Serial mode is the deterministic reference; PrefillMode::Parallel parallelizes every phase and stays bit-equal. Both architectures.
• HuggingFace loader — safetensors, synthetic + real Mamba SSM checkpoints (130m / 370m / 1.4b / 2.8b validated).
• Standalone — no framework dependency. MSRV 1.97.

Cargo features

Use cases and API choice

The crate targets two workloads. Pick the entry point that matches yours.

Reinforcement learning / small custom models

Latency-critical, typically d_model ≤ 256, often batch = 1 for actor rollouts. Both CPU and GPU paths are supported; CPU is competitive at these sizes (~87 µs/step on Ada Xeon vs 79 µs/step on RTX 6000 Ada).

• Inference — mamba_step (CPU) or GpuMambaBackbone::step (GPU)
• Training — parallel_mamba_forward / parallel_mamba_backward (CPU, Rayon-parallel batch) or MambaTrainer::step (GPU, CUDA-Graph- captured forward + backward + AdamW + sync)

CPU training works for model sizes where GPU overhead dominates (d_model ≤ 128, batch ≤ 8); GPU training scales well to batch ≥ 32.

Large language models

Throughput-critical, d_model ≥ 768, sequence-level decoding with a HuggingFace checkpoint. GPU-only in practice — a 2.8b model on CPU is single-digit tokens/sec regardless of implementation.

• Inference — GpuMambaLM::from_hf_with_dtype + generate
• Fine-tuning — MambaTrainer::new_full accepting the HF backbone weights (Mamba SSM only; no public Mamba-3 SISO checkpoint exists yet)

The CPU MambaLM path compiles and runs end-to-end, but exists for CPU↔GPU parity testing (tests/hf_batch_parity.rs), not for production LLM serving.

Sequence classification / embeddings / custom heads

Whole-sequence reads with a caller-defined loss (document classifiers, distillation, contrastive embedding). GPU training rides the forward/backward split; CPU serving rides the prefill.

• Training — MambaTrainer::forward + host-side loss + MambaTrainer::backward_step (see examples/custom_loss.rs)
• CPU serving — MambaBackbone::forward_prefill / forward_mamba3_backbone_prefill (see examples/cpu_prefill.rs); batches of sequences via prefill_batch / prefill3_batch

Sharing weights across paths

All paths consume the same MambaWeights / Mamba3Weights struct. A training run's MambaTrainer::snapshot_master() output loads directly into GpuMambaBackbone, GpuMambaLM, or the CPU MambaBackbone without conversion.

Quick start (CPU)

Mamba SSM

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

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);

Quick start (GPU inference)

[dependencies]
mamba-rs = { version = "0.5", features = ["cuda"] }

GpuMambaBackbone::new_with_dtype and the symmetric Mamba-3 constructor take WeightDtype::{F32, Bf16, F16} — the rest of the API is unchanged.

use mamba_rs::gpu::inference::GpuMambaBackbone;
use mamba_rs::WeightDtype;

let mut gpu = GpuMambaBackbone::new_with_dtype(0, &weights, cfg, input_dim, batch, WeightDtype::Bf16)?;
gpu.capture_graph()?; // optional; ~2× decode speedup
gpu.step(&input, &mut output)?;
gpu.reset()?;

HuggingFace LM inference

use mamba_rs::module::gpu_lm::GpuMambaLM;
use mamba_rs::module::sample::SampleParams;
use mamba_rs::WeightDtype;
use std::path::Path;

let mut lm = GpuMambaLM::from_hf_with_dtype(
    Path::new("./mamba-130m-hf"), 0, WeightDtype::Bf16,
)?;
lm.capture_graph()?;
let tokens = lm.generate(&[1, 2, 3, 4, 5], &SampleParams::default())?;

bf16 vs f32 on all four cached state-spaces/mamba-*-hf checkpoints: 15/15 greedy match, KL ≤ 1.6e-3. Batch=1 vs batch=32 on the same prompt: KL ≈ 2e-11 (bit-identical up to f32 roundoff of the fixed reduction tree).

Quick start (GPU training)

MambaTrainer / Mamba3Trainer wrap the full forward + backward + AdamW + sync pipeline behind a single .step() call. One dispatch struct per architecture; an internal enum selects the f32 or mixed (bf16/f16) inner engine based on the WeightDtype constructor argument.

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(
    /* gpu_ordinal */ 0,
    &cpu_weights, cfg, session,
    WeightDtype::Bf16,
)?;
trainer.capture_graph()?; // optional; one cuGraphLaunch 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

Mamba3Trainer mirrors the same API. f16 training activates the dynamic loss scaler automatically; metrics.loss_scale / metrics.overflow_skipped report its state each step.

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:

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 readback (f32)
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 + clip + AdamW
// gradient accumulation: .with_accumulate_only(true) on the non-applying
// micro-batches (the fused step() refuses while a window is open).

Always eager (a caller-side loss cannot live inside a captured graph), and bit-identical to the fused step() — both compose the same phase bodies. See examples/custom_loss.rs for a complete training loop.

Quick start (CPU prefill)

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-SGEMM pass over the whole prompt instead of T step dispatches.
backbone.forward_prefill(&prompt, &mut out, &mut state, &mut scratch,
                         seq_len, PrefillMode::Parallel);
// `out` holds the post-norm_f output at EVERY position (pooling-ready);
// `state` is positioned after the prompt — forward_step continues from it.

Enable gemm-blas (or accelerate on macOS) — the default scalar GEMM is a correctness fallback, not a serving configuration. Mamba-3 has the same surface (forward_mamba3_backbone_prefill + Mamba3PrefillScratch).

Serialization

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 (RTX 6000 Ada)

LLM throughput — mamba-130m-hf, greedy decode, CUDA Graph, RTX 6000 Ada

On f32 both paths run on CUDA cores (no Tensor Core route), so the gap is small. On bf16/f16 cuBLAS routes through Tensor Cores (TF32-style accumulation) and wins ~7 % on per-token latency, at the cost of M=1 vs M=N algorithm-selection drift (KL ≈ 1e-3 on adversarial prompts). The batch-invariant path keeps b=1 ≡ b=N per slot (KL ≈ 1e-11).

Enable the batch-invariant path when cross-batch bit-identity matters (KL ≈ 1e-11 between b=1 and b=N per slot): set MAMBA_RS_BATCH_INVARIANT=1 or call ctx.set_batch_invariant(true).

Deterministic training — cost per step (RTX 6000 Ada, MambaTrainer)

With the batch-invariant flag on, every training GEMM (forward, dW, dX) runs on custom fixed-reduction-order kernels: two runs with the same seed/inputs produce bit-identical weights, on every dtype. The optional tensor-core tier keeps full determinism under its own numeric contract (mma.sync f32 accumulation instead of the scalar FMA chain) and turns the determinism overhead into a speedUP on LLM-sized models:

trainer.ctx().set_batch_invariant(true);   // bit-identical runs, scalar contract
trainer.ctx().set_bi_tensor_cores(true);   // + tensor-core tier (own contract)

GEMM-level tensor-core speedups vs the scalar deterministic tier: forward 3.2–6.4×, dW 4.0–5.6×, dX 3.5–5.1× (bf16, M=2048-class shapes). Two bit-identical tile families (128×128 and 64×64, shape-routed) cover everything from d128 RL models to LLM projections. Full tables and contracts: deterministic GEMM benchmarks.

Per-step latency (default config: d_model=128, 3 layers)

Detailed tables: Mamba SSM benchmarks, Mamba-3 SISO benchmarks.

Testing

52 test files, 360+ individual tests:

• Correctness: bit-parity across CPU ↔ GPU, eager ↔ CUDA Graph, f32 ↔ bf16/f16
• Gradient checks: finite-difference vs analytical on every weight tensor
• Real checkpoints: 30-step training convergence + inference on state-spaces/mamba-130m-hf for all three dtypes
• Batch invariance: KL < 1e-4 across batch sizes 1 / 4 / 16 / 32 at bf16
• Determinism: bit-identical training across runs (f32/bf16/f16, scalar and tensor-core tiers), typed-GEMM bit-parity vs the f32 reference across a 60-shape dispatch-gate boundary sweep
• Long-sequence stability: 1024-token generation + T=1024 M3 training
• CUDA Graph: replay determinism, pointer-stability assertions

Run the fast suite:

cargo test --release --features cuda

Full suite including HuggingFace-backed tests (needs the HF cache):

cargo test --release --features "cuda hf" -- --include-ignored

Documentation

• Mamba SSM architecture
• Mamba-3 SISO architecture
• Mamba SSM benchmarks
• Mamba-3 SISO benchmarks
• Deterministic GEMM benchmarks — tiers, contracts, full measurement tables (training step, tensor-core GEMM level, fallback tax), reproduction commands

Citation

@inproceedings{mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  booktitle={International Conference on Learning Representations},
  year={2024}
}

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