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

//! Mamba-3 SISO CPU backend.
//!
//! - `dims` — dimension calculator
//! - `inference` — T=1 recurrent step
//! - `prefill` — full-sequence batched inference forward (no tape)
//! - `forward` — training forward (F1-F7)
//! - `backward` — training BPTT backward (B1-B8)
//! - `parallel` — rayon batch forward + backward
//! - `scratch` — reusable per-sample activation scratch
//! - `flat` — packed flat layer layout for GPU upload
//! - `weights` — training weight structs

pub mod backward;
pub mod dims;
pub mod flat;
pub mod forward;
pub mod inference;
pub mod parallel;
pub mod prefill;
pub mod scratch;
pub mod weights;

pub use dims::Mamba3Dims;
pub use flat::{Mamba3FieldOffsets, Mamba3LayerFlat};
pub use inference::{Mamba3StepScratch, mamba3_layer_step, mamba3_step};
pub use prefill::{
    Mamba3PrefillScratch, forward_mamba3_backbone_prefill, forward_mamba3_backbone_prefill_mode,
    prefill3_batch,
};
pub use scratch::Mamba3Scratch;
pub use weights::{TrainMamba3LayerWeights, TrainMamba3Weights};