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

//! Mamba-3 SISO CUDA GPU backend.
//!
//! - `kernels` — NVRTC compilation + 43 kernel handles
//! - `mamba3_gpu` — GPU training forward + backward
//! - `inference` — GPU T=1 step + CUDA Graph
//! - `weights` — GPU weight/gradient storage

pub mod inference;
pub mod kernels;
pub mod mamba3_gpu;
pub mod weights;

pub use inference::GpuMamba3Backbone;
pub use kernels::Mamba3Kernels;
pub use weights::{
    GpuMamba3Grads, GpuMamba3LayerGrads, GpuMamba3LayerWeights, GpuMamba3LayerWeightsInf,
    GpuMamba3Weights, GpuMamba3WeightsInf,
};