Crate mamba_rs 

mamba-rs

Mamba SSM (Selective State Space Model) implementation in Rust.

Provides both CPU and GPU (CUDA) paths for inference and training, including full backward pass with BPTT through recurrent state.

Based on: Gu & Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces” (NeurIPS 2024).

Module Structure

• mamba_ssm — Mamba-1 implementation (CPU + GPU)
  • cpu/ — inference, forward, backward
  • gpu/ — CUDA inference, forward, backward
• ops — shared operations (dims, BLAS, math)
• module — high-level MambaBackbone API
• config, state, weights, serialize — data types

Re-exports

pub use config::MambaConfig;

pub use mamba_ssm::cpu::inference::MambaLayerScratch;

pub use mamba_ssm::cpu::inference::MambaStepScratch;

pub use mamba_ssm::cpu::inference::mamba_block_step;

pub use mamba_ssm::cpu::inference::mamba_layer_step;

pub use mamba_ssm::cpu::inference::mamba_step;

pub use module::MambaBackbone;

pub use state::MambaLayerState;

pub use state::MambaState;

pub use weights::MambaLayerWeights;

pub use weights::MambaWeights;

Modules

config

inference

mamba_ssm

Mamba-1 SSM (Selective State Space Model).

module

High-level Mamba wrappers.

ops

Shared operations: dimensions, BLAS, math utilities.

serialize

Weight serialization via safetensors format (HuggingFace standard).

state

train

weights