mamba-rs 0.3.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. Batch-invariant bf16 inference — per-row output is bit-identical across batch sizes.
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,
};