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
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

//! CUDA GPU backend for Mamba SSM.
//!
//! - `inference` — GPU T=1 step with CUDA Graphs
//! - `forward` — GPU training forward pass
//! - `backward` — GPU training backward pass
//! - `device` — CUDA context + cuBLAS
//! - `buffers` — GPU memory management
//! - `kernels` — NVRTC kernel compilation
//! - `blas` — cuBLAS SGEMM wrappers
//! - `launch` — kernel launch helpers

pub mod adamw;
pub mod backward;
pub mod backward_mixed;
pub mod blas;
pub mod buffers;
pub mod context;
pub mod device;
pub mod dtype;
pub mod forward;
pub mod forward_mixed;
pub mod grad_clip;
pub mod graph_capture;
pub mod inference;
pub mod kernels;
pub mod launch;
pub mod loss_scaler;
pub mod prefill;
pub mod sgemm_bi;
pub mod trainer;
pub mod training_graph;
pub mod weights;
pub mod weights_mixed_train;

pub use dtype::WeightDtype;