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
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

//! Shared CUDA Graph capture helper.
//!
//! Wraps the begin_capture / run-body / end_capture dance in a single
//! function so call sites don't repeat the ~25 lines of error-handling
//! boilerplate (pre-sync, end-on-error to restore stream, ok_or on the
//! end_capture Option).
//!
//! Used by every captured pipeline in the crate:
//! - M1 inference (`inference::GpuInferenceEngine::capture_graph`)
//! - M3 mixed + mixed-native inference (`inference::GpuInferenceMixed*::capture_graph*`)
//! - M1 training step (`training_graph::GpuMambaTrainingStepGraph::capture`)
//!
//! ## Why pre-sync
//! Without it, a freshly-allocated buffer used inside the body can race
//! against an in-flight HtoD that was issued before capture started. The
//! "130m race lesson" — see `GpuMambaBackboneMixedActs::new`.
//!
//! ## Why end-on-error
//! `cuStreamBeginCapture` puts the stream into a sticky capture mode. If
//! the body errors and we return without `end_capture`, every subsequent
//! op on the stream silently fails. So we ALWAYS call `end_capture`,
//! discard its result on the error path, then propagate the body error.

use std::sync::Arc;

use cudarc::driver::{CudaGraph, CudaStream};

/// Capture all CUDA work issued by `body` on `stream` into a CUDA Graph.
///
/// Mirrors the pattern in `inference::GpuInferenceEngine::capture_graph`,
/// extracted so it doesn't need to be reimplemented per pipeline.
pub fn capture_into_graph<F>(stream: &Arc<CudaStream>, body: F) -> Result<CudaGraph, String>
where
    F: FnOnce() -> Result<(), String>,
{
    stream
        .synchronize()
        .map_err(|e| format!("pre-capture sync: {e:?}"))?;

    stream
        .begin_capture(
            cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
        )
        .map_err(|e| format!("begin_capture: {e:?}"))?;

    let body_result = body();

    // ALWAYS end capture, even on body error — a stream stuck in capture
    // mode silently breaks every subsequent op.
    let end_result = stream.end_capture(
        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
    );

    // Combine both error paths: if BOTH body and end_capture failed, we want
    // the caller to see both — otherwise an `?`-shortcircuit on body_result
    // would silently drop a stream-corrupting end_capture failure.
    match (body_result, end_result) {
        (Ok(()), Ok(Some(g))) => Ok(g),
        (Ok(()), Ok(None)) => Err("end_capture returned no graph (empty body?)".to_string()),
        (Ok(()), Err(e)) => Err(format!("end_capture: {e:?}")),
        (Err(b), Ok(_)) => Err(format!("body: {b}")),
        (Err(b), Err(e)) => Err(format!(
            "body: {b}; end_capture ALSO failed (stream may be in invalid state): {e:?}"
        )),
    }
}