mamba-rs 0.4.2

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
//! Shared operations: dimensions, BLAS, math, normalization utilities.

pub mod blas;
pub mod dims;
pub mod fast_math;
pub mod norms;

pub use dims::{MambaDims, MambaRecurrentState};