mamba-rs 0.7.4

Mamba SSM and Mamba-3 SISO in Rust with optional CUDA acceleration: inference and training (BPTT through the SSM state, AdamW) on CPU and GPU, custom NVRTC-compiled kernels, CUDA Graph capture, f32 / bf16 / f16 storage, deterministic batch-invariant GEMMs by default with explicit cuBLAS Fast and Pedantic modes.
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

const HALF_F32OUT_S3_SOURCE: &str =
    include_str!("../../../../kernels/gemm_bi_inference/sm80/half_f32out_s3.cu");
const F32_M128N64_TAIL_COPYPLAN_SOURCE: &str =
    include_str!("../../../../kernels/gemm_bi_inference/sm80/f32_m128n64_tail_copyplan.cu");
const RETAINED_EXPORT_DEFINITIONS: [&str; 3] = [
    "void nn_sm89_tc128_f32out_s3_bf16(",
    "void nn_sm89_tc128_f32out_s3_f16(",
    "void nn_sm89_f32_m128n64_tail_copyplan(",
];

pub(crate) fn compiler_supported(
    device_cc: Option<(i32, i32)>,
    target: &str,
    state_cap: usize,
    nvrtc: (i32, i32),
) -> bool {
    // The retained members are sm_80-tier PTX: every board that composes
    // the portable overlay carries them, the Ada board by its frozen
    // evidence and every other board by the first-use proof. The CC 12
    // family composes no overlay: its Fixed module stays byte-identical to
    // the one its copy-plan cohorts were minted on.
    device_cc.is_some_and(|(major, _)| major >= 8 && major != 12)
        && super::super::gemm_bi_triad::modules::fixed_portable_overlay_composed(target)
        && matches!(state_cap, 16 | 64)
        && matches!(nvrtc, (12, 8) | (13, 0) | (13, 2))
}

pub(crate) fn compose_fixed_source(
    mut composed: String,
    device_cc: Option<(i32, i32)>,
    target: &str,
    state_cap: usize,
    nvrtc: (i32, i32),
) -> Result<String, String> {
    if !compiler_supported(device_cc, target, state_cap, nvrtc) {
        return Ok(composed);
    }
    for definition in RETAINED_EXPORT_DEFINITIONS {
        if composed.contains(definition) {
            return Err(format!(
                "retained Inference source already defines {definition}"
            ));
        }
    }

    append_fragment(
        &mut composed,
        "kernels/gemm_bi_inference/sm80/half_f32out_s3.cu",
        HALF_F32OUT_S3_SOURCE,
    );
    append_fragment(
        &mut composed,
        "kernels/gemm_bi_inference/sm80/f32_m128n64_tail_copyplan.cu",
        F32_M128N64_TAIL_COPYPLAN_SOURCE,
    );
    Ok(composed)
}

fn append_fragment(composed: &mut String, logical_path: &str, fragment: &str) {
    composed.push_str("\n#line 1 \"");
    composed.push_str(logical_path);
    composed.push_str("\"\n");
    composed.push_str(fragment);
    if !fragment.ends_with('\n') {
        composed.push('\n');
    }
}