burn_mamba/
lib.rs

1//! # burn-mamba — Mamba-1/2/3 selective state space models on Burn
2//!
3//! A minimal, readable reference implementation of the
4//! [Mamba-1](https://arxiv.org/abs/2312.00752),
5//! [Mamba-2](https://arxiv.org/abs/2405.21060), and
6//! [Mamba-3](https://arxiv.org/abs/2603.15569) SSM architectures on top of the
7//! [Burn](https://github.com/tracel-ai/burn/) deep learning framework.
8//!
9//! The goal is clarity: the official CUDA/Triton kernels are ported down to
10//! standard, portable Burn tensor operations, so the same code runs on every
11//! backend (CPU, WGPU, CUDA, Metal, LibTorch, …).  There are **no custom
12//! kernels**.
13//!
14//! ## Module families
15//!
16//! Each family lives in its own module and follows the same composition
17//! (`Network` → `Layers` → `Layer` → `Block`):
18//!
19//! - [`mamba1`] — the original selective SSM (conv1d + sequential selective
20//!   scan).
21//! - [`mamba2`] — Structured State Space Duality (SSD): the recurrence is recast
22//!   as a chunkwise, GEMM-friendly algorithm.
23//! - [`mamba3`] — SSD extended with trapezoidal discretisation, data-dependent
24//!   RoPE on B/C, and MIMO rank expansion.
25//!
26//! Shared infrastructure lives in [`modules`] (the family-generic layer/network
27//! composition, plus activations, norms and losses) and [`utils`] (virtual-layer
28//! scheduling, class tokens, and the custom-backward plumbing).
29//!
30//! ## Two execution modes
31//!
32//! Every block, layer, and network exposes both a parallel `forward()` (used
33//! for training and prompt prefill) and a recurrent `step()` (used for
34//! token-by-token decoding).  The two are mathematically equivalent: a
35//! `forward()` over a sequence equals unrolling `step()` token by token from the
36//! same initial cache — a parity property the test suites assert on outputs,
37//! final cache, and gradients.
38
39#![warn(missing_docs)]
40#![allow(clippy::let_and_return)]
41#![allow(clippy::module_inception)]
42#![allow(warnings)]
43
44/// Mamba-1: the original selective state space model.
45#[cfg(feature = "mamba1")]
46pub mod mamba1;
47/// Mamba-2: Structured State Space Duality (SSD).
48#[cfg(feature = "mamba2")]
49pub mod mamba2;
50/// Mamba-3: trapezoidal SSD with data-dependent RoPE and MIMO.
51#[cfg(feature = "mamba3")]
52pub mod mamba3;
53
54/// Convenience re-exports: `use burn_mamba::prelude::*;` brings the enabled
55/// model families and their public types into scope.
56pub mod prelude {
57    #[cfg(feature = "mamba1")]
58    pub use crate::mamba1::{self, prelude::*};
59
60    #[cfg(feature = "mamba2")]
61    pub use crate::mamba2::{self, prelude::*};
62
63    #[cfg(feature = "mamba3")]
64    pub use crate::mamba3::{self, prelude::*};
65
66    // The family-generic, runtime-selectable unified API.
67    pub use crate::modules::{
68        CacheStack, Layer, Layers, MambaBidiLayers, MambaBidiLayersConfig, MambaBlock,
69        MambaBlockConfig, MambaCaches, MambaLatentNet, MambaLatentNetConfig, MambaSsdPath,
70        MambaVocabNet, MambaVocabNetConfig, ResidualsConfig, StateMoments, StatePairing,
71    };
72    pub use crate::utils::{ClassLatent, ClassToken};
73}
74
75/// Family-generic composition (`Layer`/`Layers`/networks/bidi/caches) plus the
76/// shared neural modules (activations, norms, losses, tensor helpers).
77pub mod modules;
78/// Virtual-layer/LR scheduling, class tokens, and custom-backward plumbing.
79pub mod utils;
80
81/// When `true`, [`modules::sanity`] panics if it observes a `NaN`.
82///
83/// Compiled-in guard (off by default) for debugging numerical issues; leaving
84/// it `false` removes the check entirely.
85#[cfg(feature = "check-nan")]
86pub const DENY_NAN: bool = true;
87#[cfg(not(feature = "check-nan"))]
88pub const DENY_NAN: bool = false;
89
90/// When `true`, [`modules::sanity`] panics if it observes an `Inf`.
91///
92/// Compiled-in guard (off by default), companion to [`DENY_NAN`].
93#[cfg(feature = "check-inf")]
94pub const DENY_INF: bool = true;
95#[cfg(not(feature = "check-inf"))]
96pub const DENY_INF: bool = false;
burn_mamba/lib.rs

burn_mamba/
lib.rs
