mamba-rs 0.3.1

Mamba SSM and Mamba-3 SISO in Rust with optional CUDA GPU acceleration. Inference and training (BPTT through SSM state, AdamW), CPU + GPU paths, custom CUDA kernels, CUDA Graph capture, f32 / bf16 / f16. Batch-invariant bf16 inference — per-row output is bit-identical across batch sizes.
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

//! Target network forward for Mamba-3 SISO.
//!
//! Reuses `forward_mamba3_layer_batched` with a throwaway activation buffer.
//! No backward needed — just state evolution for target Q-network.

use super::dims::Mamba3Dims;
use super::flat::Mamba3LayerFlat;
use super::forward::forward_mamba3_layer_batched;
use super::scratch::Mamba3Scratch;
use super::weights::TrainMamba3LayerWeights;

/// Target network forward — single layer. Mutates state, discards activations.
pub fn target_mamba3_layer_forward(
    temporal_flat: &mut [f32],
    layer_w: &TrainMamba3LayerWeights,
    ssm_state: &mut [f32],
    k_state: &mut [f32],
    v_state: &mut [f32],
    angle_state: &mut [f32],
    scratch: &mut Mamba3Scratch,
    acts: &mut Mamba3LayerFlat,
    dims: &Mamba3Dims,
) {
    // Reuse full forward — activations saved but not used for backward.
    // For production, a no-save variant would save ~50% bandwidth.
    forward_mamba3_layer_batched(
        temporal_flat,
        acts,
        layer_w,
        ssm_state,
        k_state,
        v_state,
        angle_state,
        scratch,
        dims,
    );
}