mamba-rs 0.2.0

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

//! CUDA GPU backend for Mamba-1 SSM.
//!
//! - `inference` — GPU T=1 step with CUDA Graphs
//! - `forward` — GPU training forward pass
//! - `backward` — GPU training backward pass
//! - `device` — CUDA context + cuBLAS
//! - `buffers` — GPU memory management
//! - `kernels` — NVRTC kernel compilation
//! - `blas` — cuBLAS SGEMM wrappers
//! - `launch` — kernel launch helpers

pub mod backward;
pub mod blas;
pub mod buffers;
pub mod context;
pub mod device;
pub mod forward;
pub mod inference;
pub mod kernels;
pub mod launch;
pub mod weights;