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
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

//! Mamba-3 SISO CUDA GPU backend.
//!
//! - `kernels` — NVRTC compilation + 43 kernel handles
//! - `mamba3_gpu` — GPU training forward + backward
//! - `inference` — GPU T=1 step + CUDA Graph
//! - `weights` — GPU weight/gradient storage

pub mod backward;
pub mod backward_mixed;
pub mod forward;
pub mod forward_mixed;
pub mod inference;
pub mod kernels;
pub mod mamba3_gpu;
pub mod state;
pub mod trainer;
pub mod training_graph;
pub mod weights;
pub mod weights_mixed_train;

pub use inference::GpuMamba3Backbone;
pub use kernels::Mamba3Kernels;
pub use weights::{
    GpuMamba3Grads, GpuMamba3LayerGrads, GpuMamba3LayerWeights, GpuMamba3LayerWeightsInf,
    GpuMamba3Weights, GpuMamba3WeightsInf,
};