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

//! Mamba SSM (Selective State Space Model).
//!
//! Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" (2023).
//!
//! ## Structure
//! - `cpu/` — CPU inference + training (forward, backward)
//! - `gpu/` — GPU inference + training (CUDA kernels, CUDA Graphs)

pub mod cpu;

#[cfg(feature = "cuda")]
pub mod gpu;