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
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

//! Persistent recurrent state for Mamba-3 SISO layers.
//!
//! Four state components per layer (no conv state — Mamba-3 has no conv1d):
//! - SSM hidden state: `[nheads * headdim * d_state]`
//! - K state (previous B after RoPE): `[nheads * d_state]`
//! - V state (previous x): `[nheads * headdim]`
//! - Angle state (RoPE cumulative angles): `[nheads * num_rope_angles]`

use super::config::Mamba3Config;

/// Persistent recurrent state for one Mamba-3 SISO layer.
#[derive(Clone)]
pub struct Mamba3LayerState {
    /// SSM hidden state: `[nheads * headdim * d_state]`.
    pub ssm_state: Vec<f32>,
    /// K state (post-RoPE, post-bias B from previous step): `[nheads * d_state]`.
    pub k_state: Vec<f32>,
    /// V state (x from previous step): `[nheads * headdim]`.
    pub v_state: Vec<f32>,
    /// RoPE cumulative angle state: `[nheads * num_rope_angles]`.
    /// Wraps to `[0, 2*pi)` after each step.
    pub angle_state: Vec<f32>,
}

impl Mamba3LayerState {
    /// Create a zeroed layer state.
    pub fn zeros(nheads: usize, headdim: usize, d_state: usize, num_rope_angles: usize) -> Self {
        Self {
            ssm_state: vec![0.0; nheads * headdim * d_state],
            k_state: vec![0.0; nheads * d_state],
            v_state: vec![0.0; nheads * headdim],
            angle_state: vec![0.0; nheads * num_rope_angles.max(1)],
        }
    }

    /// Reset all state to zero (new sequence boundary).
    pub fn reset(&mut self) {
        self.ssm_state.fill(0.0);
        self.k_state.fill(0.0);
        self.v_state.fill(0.0);
        self.angle_state.fill(0.0);
    }
}

/// Full recurrent state for the Mamba-3 SISO backbone (all layers).
#[derive(Clone)]
pub struct Mamba3State {
    pub layers: Vec<Mamba3LayerState>,
}

impl Mamba3State {
    /// Create zeroed state for all layers from config.
    pub fn zeros(cfg: &Mamba3Config) -> Self {
        let nh = cfg.nheads();
        let hd = cfg.headdim;
        let ds = cfg.d_state;
        let na = cfg.num_rope_angles();
        Self {
            layers: (0..cfg.n_layers)
                .map(|_| Mamba3LayerState::zeros(nh, hd, ds, na))
                .collect(),
        }
    }

    /// Reset all layers (new sequence boundary).
    pub fn reset(&mut self) {
        for layer in &mut self.layers {
            layer.reset();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_state_zeros() {
        let cfg = Mamba3Config::default();
        let state = Mamba3State::zeros(&cfg);
        assert_eq!(state.layers.len(), cfg.n_layers);
        let l = &state.layers[0];
        assert_eq!(l.ssm_state.len(), cfg.nheads() * cfg.headdim * cfg.d_state);
        assert_eq!(l.k_state.len(), cfg.nheads() * cfg.d_state);
        assert_eq!(l.v_state.len(), cfg.nheads() * cfg.headdim);
        assert_eq!(
            l.angle_state.len(),
            cfg.nheads() * cfg.num_rope_angles().max(1)
        );
        assert!(l.ssm_state.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn test_state_reset() {
        let cfg = Mamba3Config::default();
        let mut state = Mamba3State::zeros(&cfg);
        state.layers[0].ssm_state[0] = 42.0;
        state.layers[0].angle_state[0] = std::f32::consts::PI;
        state.reset();
        assert_eq!(state.layers[0].ssm_state[0], 0.0);
        assert_eq!(state.layers[0].angle_state[0], 0.0);
    }
}