mamba-rs 0.7.2

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

//! Mamba SSM CPU inference — zero-allocation single-step loop.
//!
//! ```bash
//! cargo run --release --example inference
//! ```
//!
//! Demonstrates the minimal CPU path: allocate state and scratch once,
//! then call `forward_step` per token. Used for RL actors and similar
//! small-model latency-critical workloads where GPU overhead dominates.

use mamba_rs::{MambaBackbone, MambaConfig};

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights.
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba SSM: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}