mamba-rs 0.5.1

Mamba SSM and Mamba-3 SISO in Rust with optional CUDA GPU acceleration. Inference and training (BPTT through SSM state, AdamW), CPU + GPU paths, custom CUDA kernels, CUDA Graph capture, f32 / bf16 / f16. Opt-in deterministic training (bit-identical runs, batch-invariant inference) with a tensor-core tier that beats cuBLAS on LLM-sized models.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

//! CPU backend for Mamba SSM.
//!
//! - `inference` — zero-alloc T=1 step (single + batch + sequence)
//! - `forward` — training forward pass (batched SGEMM)
//! - `backward` — training backward pass (BPTT)
//! - `target` — target network forward (no activation saves)

pub mod backward;
pub mod backward_ops;
pub mod flat;
pub mod forward;
pub mod inference;
pub mod parallel;
pub mod prefill;
pub mod scratch;
pub mod target;
pub mod weights;