mamba-rs 0.2.1

Mamba SSM (Selective State Space Model) in Rust with CUDA GPU support. Training + inference, forward + backward (BPTT), burn-in, custom CUDA kernels.
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
58
59
60
61
62
63
64
65
66
67
68
69

//! # mamba-rs
//!
//! Mamba SSM (Selective State Space Model) implementation in Rust.
//!
//! Provides both CPU and GPU (CUDA) paths for inference and training,
//! including full backward pass with BPTT through recurrent state.
//!
//! Based on: Gu & Dao, "Mamba: Linear-Time Sequence Modeling with
//! Selective State Spaces" (NeurIPS 2024).
//!
//! ## Module Structure
//!
//! - [`mamba_ssm`] — Mamba-1 implementation (CPU + GPU)
//!   - `cpu/` — inference, forward, backward
//!   - `gpu/` — CUDA inference, forward, backward
//! - [`mamba3_siso`] — Mamba-3 SISO implementation (CPU + GPU)
//!   - `cpu/` — inference, forward, backward, parallel
//!   - `gpu/` — CUDA kernels (38 kernels)
//! - [`ops`] — shared operations (dims, BLAS, math, norms)
//! - [`module`] — high-level MambaBackbone API
//! - [`config`], [`state`], [`weights`], [`serialize`] — Mamba-1 data types

pub mod config;
pub mod mamba3_siso;
pub mod mamba_ssm;
pub mod module;
pub mod ops;
pub mod serialize;
pub mod state;
pub mod weights;

// Re-export old paths for backward compatibility during transition.
// These will be removed once all external users migrate.
pub mod inference {
    pub use crate::mamba_ssm::cpu::inference::*;
}
pub mod train {
    pub use crate::mamba_ssm::cpu::backward;
    pub use crate::mamba_ssm::cpu::backward_ops;
    pub use crate::mamba_ssm::cpu::flat;
    pub use crate::mamba_ssm::cpu::forward;
    pub use crate::mamba_ssm::cpu::parallel;
    pub use crate::mamba_ssm::cpu::scratch;
    pub use crate::mamba_ssm::cpu::target;
    pub use crate::mamba_ssm::cpu::weights;

    // Re-export shared ops that were previously in train/
    pub use crate::ops::blas;
    pub use crate::ops::fast_math;
}

#[cfg(feature = "cuda")]
pub mod gpu {
    pub use crate::mamba_ssm::gpu::*;
}

pub use config::MambaConfig;
pub use mamba_ssm::cpu::inference::{
    MambaLayerScratch, MambaStepScratch, mamba_block_step, mamba_layer_step, mamba_step,
};
pub use module::MambaBackbone;
pub use state::{MambaLayerState, MambaState};
pub use weights::{MambaLayerWeights, MambaWeights};

// Mamba-3 SISO re-exports
pub use mamba3_siso::{
    Mamba3Config, Mamba3Dims, Mamba3LayerState, Mamba3LayerWeights, Mamba3State, Mamba3StepScratch,
    Mamba3Weights,
};