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

//! Mamba-3 SISO CPU backend.
//!
//! - `dims` — dimension calculator
//! - `inference` — T=1 recurrent step (TODO)
//! - `forward` — training forward (TODO)
//! - `backward` — training backward (TODO)

pub mod backward;
pub mod dims;
pub mod flat;
pub mod forward;
pub mod inference;
pub mod parallel;
pub mod scratch;
pub mod target;
pub mod weights;

pub use dims::Mamba3Dims;
pub use flat::{Mamba3FieldOffsets, Mamba3LayerFlat};
pub use inference::{Mamba3StepScratch, mamba3_layer_step, mamba3_step};
pub use scratch::Mamba3Scratch;
pub use weights::{TrainMamba3LayerWeights, TrainMamba3Weights};