mamba-rs 0.6.5

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

//! Error taxonomy for the distributed data-parallel layer.

use std::fmt;

/// Everything that can go wrong while setting up or running a
/// data-parallel world. Variants carry enough context to act on: config
/// errors are caller bugs, rendezvous/transport errors are environment
/// problems. `RankFailed` and `ReplicaDivergence` are RESERVED for the
/// training-harness lanes (supervisor-side peer diagnostics and the
/// replica-consistency probe) — no code in this crate produces them
/// yet; today a dead peer surfaces as a `Transport` error from the
/// aborted collective window and a non-zero rank exit the supervisor
/// aggregates.
#[non_exhaustive]
#[derive(Debug)]
pub enum DistError {
    /// Invalid configuration (bad device list, world size zero, an
    /// accumulation window not divisible by the world size, ...).
    Config(String),
    /// The launch environment promised a rank layout it did not deliver
    /// (missing or contradictory environment variables).
    EnvContract(String),
    /// Rendezvous failed: peers did not show up within the deadline, or
    /// the rendezvous directory is unusable.
    Rendezvous(String),
    /// A collective or transport operation failed.
    Transport(String),
    /// A peer rank exited or signalled failure; the run must stop.
    RankFailed { rank: usize, detail: String },
    /// Replica-consistency check failed: ranks that must hold identical
    /// bytes diverged (wrong binary, libm drift, memory corruption).
    ReplicaDivergence { tag: String, detail: String },
}

impl fmt::Display for DistError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DistError::Config(m) => write!(f, "dist config: {m}"),
            DistError::EnvContract(m) => write!(f, "dist env contract: {m}"),
            DistError::Rendezvous(m) => write!(f, "dist rendezvous: {m}"),
            DistError::Transport(m) => write!(f, "dist transport: {m}"),
            DistError::RankFailed { rank, detail } => {
                write!(f, "dist rank {rank} failed: {detail}")
            }
            DistError::ReplicaDivergence { tag, detail } => {
                write!(f, "dist replica divergence at {tag}: {detail}")
            }
        }
    }
}

impl std::error::Error for DistError {}

impl From<DistError> for String {
    fn from(e: DistError) -> Self {
        e.to_string()
    }
}