mamba-rs 0.7.1

Mamba SSM and Mamba-3 SISO in Rust with optional CUDA acceleration: inference and training (BPTT through the SSM state, AdamW) on CPU and GPU, custom NVRTC-compiled kernels, CUDA Graph capture, f32 / bf16 / f16 storage, deterministic batch-invariant GEMMs by default with explicit cuBLAS Fast and Pedantic modes.
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

//! CUDA GPU backend for Mamba SSM.
//!
//! Every context carries a [`GemmMode`] (re-exported here), `Deterministic`
//! by default: a model context then multiplies on the Inference kernels, a
//! trainer or a plain context on the Triad kernels; the two cuBLAS modes are
//! explicit alternatives.
//!
//! - `context` — the GPU context: stream, kernels, cuBLAS handle and the
//!   GEMM route (mode, family, numeric policies)
//! - `device` — CUDA device and cuBLAS handle
//! - `buffers` — GPU memory management
//! - `kernels`, `kernel_identity` — NVRTC compilation and the frozen
//!   identity of every compiled kernel
//! - `blas` — the GEMM dispatch boundary: the deterministic families in the
//!   default mode, cuBLAS in the two vendor modes
//! - `gemm_bi_inference`, `gemm_bi_triad` — the two deterministic families
//! - `inference` — the decode step with CUDA Graph capture
//! - `prefill` — the GPU prompt prefill
//! - `forward`, `backward`, `forward_mixed`, `backward_mixed` — the f32 and
//!   the mixed-precision training passes
//! - `trainer`, `training_graph` — the training step and its captured graph
//! - `adamw`, `grad_clip`, `loss_scaler` — the optimizer, clipping and the
//!   f16 loss scaler
//! - `weights`, `weights_mixed_train`, `dtype` — device weights and the
//!   storage dtype
//! - `graph_capture`, `launch` — capture and 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(crate) mod diagnostics;
pub mod dtype;
mod fold_transport;
pub mod forward;
pub mod forward_mixed;
pub mod gemm_bi_inference;
pub mod gemm_bi_triad;
mod gemm_mode;
pub mod grad_clip;
pub mod graph_capture;
pub mod inference;
pub mod kernel_identity;
pub mod kernels;
pub mod launch;
pub mod loss_scaler;
pub mod prefill;
pub mod trainer;
pub mod training_graph;
pub mod weights;
pub mod weights_mixed_train;

pub use dtype::WeightDtype;
pub use gemm_mode::GemmMode;