Crate mamba_rs 

mamba-rs

Mamba SSM and Mamba-3 SISO in Rust with optional CUDA GPU acceleration. Supports Mamba SSM (Gu & Dao, 2023) and Mamba-3 SISO (Lahoti et al., 2026) on CPU and GPU, with full inference and training pipelines.

No Python, no C++ build step and no framework dependency: the kernels compile at run time through NVRTC, and the GPU path links only the CUDA driver API and cuBLAS.

Capabilities

• Mamba SSM and Mamba-3 SISO architectures
• CPU and GPU (CUDA) paths for both
• Full training with BPTT through the recurrent SSM state + AdamW
• WeightDtype::{F32, Tf32, Bf16, F16} with f32 master state and accumulation; the storage precision decides the product precision
• CUDA Graph capture for inference and training steps
• Deterministic f32/tf32/bf16/f16 GEMM kernels for inference and training, the default since 0.7.0; the inference kernels are batch-invariant, the training kernels within one dispatch bucket
• HuggingFace safetensors loader for Mamba SSM checkpoints

GEMM modes (CUDA)

Every GPU context carries a GemmMode (mamba_ssm::gpu::GemmMode, behind the cuda feature):

• Deterministic (default): the crate’s own fixed-reduction-order kernels serve every GEMM that goes through the context; cuBLAS is never called in this mode. Model contexts use the Inference family, trainers and plain contexts the Triad family.
• CublasFast: cuBLAS with TF32 permitted for f32 operands and f32 accumulation for half operands.
• CublasPedantic: cuBLAS with pedantic f32 compute, the default of 0.6.9 and earlier.

Select the mode at construction (GpuCtx::new_with_mode, the model and trainer *_with_mode constructors, or MAMBA_RS_GEMM_MODE for the environment-reading constructors) or change it with GpuCtx::set_gemm_mode, which is refused while a graph is being captured. Storage precision (WeightDtype) and mode are the two settings; inside the deterministic mode the kernels are chosen from them and from whether the context serves a model or a trainer. The guide is https://github.com/silvermpx/mamba-rs/blob/main/docs/gemm-modes.md and the measurements are in https://github.com/silvermpx/mamba-rs/blob/main/docs/determinism-benchmarks.md.

Module Structure

• mamba_ssm — Mamba SSM (CPU + GPU forward, backward, training)
• mamba3_siso — Mamba-3 SISO (CPU + GPU forward, backward, training)
• module — high-level backbone and LM wrappers, HF integration
• ops — shared dimensions, BLAS, norms, fast-math helpers
• dist — deterministic data-parallel training (one fixed-order reduction per optimizer step)
• config, state, weights, serialize — Mamba SSM data types

References

• Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752, 2023.
• Lahoti et al., Mamba-3: Improved Sequence Modeling using State Space Principles, ICLR 2026.

Re-exports

pub use config::MambaConfig;

pub use mamba_ssm::cpu::inference::MambaLayerScratch;

pub use mamba_ssm::cpu::inference::MambaStepScratch;

pub use mamba_ssm::cpu::inference::mamba_block_step;

pub use mamba_ssm::cpu::inference::mamba_layer_step;

pub use mamba_ssm::cpu::inference::mamba_step;

pub use mamba_ssm::cpu::inference::mamba_step_no_proj;

pub use module::MambaBackbone;

pub use state::MambaLayerState;

pub use state::MambaState;

pub use weights::MambaLayerWeights;

pub use weights::MambaWeights;

pub use mamba3_siso::Mamba3Config;

pub use mamba3_siso::Mamba3Dims;

pub use mamba3_siso::Mamba3LayerState;

pub use mamba3_siso::Mamba3LayerWeights;

pub use mamba3_siso::Mamba3State;

pub use mamba3_siso::Mamba3StepScratch;

pub use mamba3_siso::Mamba3Weights;

pub use mamba_ssm::gpu::dtype::WeightDtype;

Modules

config

dist

Deterministic data-parallel training.

gpu

gpu3

Convenience re-exports for the Mamba-3 SISO GPU path.

hf

HuggingFace checkpoint loading for Mamba LLM inference.

inference

mamba3_siso

Mamba-3 SISO (Single-Input Single-Output) implementation.

mamba_ssm

Mamba SSM (Selective State Space Model).

module

High-level Mamba wrappers.

ops

Shared operations: dimensions, BLAS, math, normalization utilities.

serialize

Weight serialization via safetensors format (HuggingFace standard).

state

train

weights

Constants

VERSION

Crate version, for consumers stamping numeric-route provenance (checkpoint sidecars, serve fingerprints) without parsing Cargo.lock.