mamba-rs 0.5.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. Opt-in deterministic training (bit-identical runs, batch-invariant inference) with a tensor-core tier that beats cuBLAS on LLM-sized models.
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

//! bf16 → f32 conversion utilities.

use half::bf16;
use half::slice::HalfFloatSliceExt;

/// Convert raw bf16 bytes to f32 vec.
///
/// `raw_bytes` must have length divisible by 2 (each bf16 is 2 bytes).
pub fn bf16_bytes_to_f32(raw_bytes: &[u8]) -> Result<Vec<f32>, String> {
    if !raw_bytes.len().is_multiple_of(2) {
        return Err(format!(
            "bf16 buffer length {} is not even",
            raw_bytes.len()
        ));
    }
    let bf16_slice: &[bf16] = bytemuck::cast_slice(raw_bytes);
    let mut out = vec![0.0f32; bf16_slice.len()];
    bf16_slice.convert_to_f32_slice(&mut out);
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_bf16_zero_roundtrip() {
        let zero_bf16 = bf16::from_f32(0.0);
        let bytes = bytemuck::bytes_of(&zero_bf16);
        let result = bf16_bytes_to_f32(bytes).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0], 0.0);
    }

    #[test]
    fn test_bf16_known_values() {
        let vals = [1.0f32, -1.0, 0.5, 2.75];
        let bf16_vals: Vec<bf16> = vals.iter().map(|&v| bf16::from_f32(v)).collect();
        let bytes: &[u8] = bytemuck::cast_slice(&bf16_vals);
        let result = bf16_bytes_to_f32(bytes).unwrap();
        assert_eq!(result.len(), 4);
        for (got, &expected) in result.iter().zip(&vals) {
            assert!(
                (got - expected).abs() < 0.02,
                "bf16 roundtrip: {got} vs {expected}"
            );
        }
    }

    #[test]
    fn test_bf16_nan_inf() {
        let vals = [f32::INFINITY, f32::NEG_INFINITY, f32::NAN];
        let bf16_vals: Vec<bf16> = vals.iter().map(|&v| bf16::from_f32(v)).collect();
        let bytes: &[u8] = bytemuck::cast_slice(&bf16_vals);
        let result = bf16_bytes_to_f32(bytes).unwrap();
        assert!(result[0].is_infinite() && result[0] > 0.0);
        assert!(result[1].is_infinite() && result[1] < 0.0);
        assert!(result[2].is_nan());
    }

    #[test]
    fn test_bf16_simd_matches_scalar() {
        let vals: Vec<f32> = (0..257).map(|i| i as f32 * 0.01 - 1.28).collect();
        let bf16_vals: Vec<bf16> = vals.iter().map(|&v| bf16::from_f32(v)).collect();
        let bytes: &[u8] = bytemuck::cast_slice(&bf16_vals);

        let simd_result = bf16_bytes_to_f32(bytes).unwrap();
        let scalar_result: Vec<f32> = bf16_vals.iter().map(|v| v.to_f32()).collect();

        assert_eq!(simd_result.len(), scalar_result.len());
        for (i, (s, r)) in simd_result.iter().zip(&scalar_result).enumerate() {
            assert_eq!(s, r, "mismatch at index {i}");
        }
    }

    #[test]
    fn test_bf16_odd_length_errors() {
        let result = bf16_bytes_to_f32(&[0u8, 1, 2]);
        assert!(result.is_err());
    }
}