mamba-rs 0.5.0

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

//! HuggingFace checkpoint loading for Mamba LLM inference.
//!
//! Supports both HF-native (`-hf`) and original (`mamba_ssm`) config formats.
//! Handles safetensors (single + multi-shard), bf16 upcast, key remapping.

pub mod bf16;
pub mod config_json;
pub mod embed;
pub mod keys;
pub mod load;

pub use config_json::ModelFamily;
pub use load::load_hf;