burn_mamba/mamba3/
mod.rs

1//! # Mamba-3
2//!
3//! Mamba-3 extends Mamba-2 with three independent additions (each works alone
4//! or combined): **trapezoidal discretisation**, **data-dependent RoPE** on the
5//! B/C projections, and **MIMO** (multiple-input multiple-output) rank
6//! expansion.  See [`mamba3`] for the full combined math.
7//!
8//! ## Two SSD pathways
9//!
10//! The trapezoidal recurrence is realised by two interchangeable algorithms,
11//! selected at runtime by which **cache variant** is supplied:
12//!
13//! - [`double_ssd`] — splits the trapezoid into two standard SSD calls
14//!   (simple, easy to verify; ~2× the intra-chunk memory).
15//! - [`single_ssd`] — one SSD call in the official-kernel form
16//!   (≈ half the training memory; the cache's SSM accumulator has different
17//!   mid-sequence semantics).
18//!
19//! [`cache`] holds the enum that dispatches between them; [`ssd_path`] selects
20//! the pathway-agnostic *algorithm* (Minimal / Serial / SerialRecalculated).
21
22pub mod double_ssd;
23pub mod single_ssd;
24
25pub mod cache;
26pub(crate) mod helpers;
27pub mod mamba3;
28pub mod moments;
29pub mod quat_scan;
30pub mod rotation;
31pub mod ssd_path;
32mod step_constant;
33
34use crate::mamba3::double_ssd::prelude::*;
35use crate::mamba3::single_ssd::prelude::*;
36use burn::backend::Backend;
37
38/// Backend capability required to run Mamba-3.
39///
40/// Aggregates the per-pathway extension traits ([`Mamba3DoubleSsdBackendExt`]
41/// and [`Mamba3SingleSsdBackendExt`]); every plain Burn backend satisfies it
42/// via the default implementations, and `Autodiff<B>` additionally gets the
43/// custom memory-efficient backward.
44pub trait Mamba3BackendExt:
45    Backend + Mamba3DoubleSsdBackendExt + Mamba3SingleSsdBackendExt
46{
47}
48
49crate::decl_ssd_autodiff_backend_ext!(
50    Mamba3AutodiffBackendExt,
51    Mamba3BackendExt,
52    Mamba3DoubleSsdAutodiffBackendExt,
53    Mamba3SingleSsdAutodiffBackendExt
54);
55crate::impl_ssd_backend_ext_for_burn_backends!(Mamba3BackendExt);
56
57/// Blanket [`Mamba3BackendExt`] implementation for autodiff backends.
58#[cfg(feature = "autodiff")]
59pub mod backwards {
60    use super::*;
61    use burn::backend::autodiff::{Autodiff, checkpoint::strategy::CheckpointStrategy};
62
63    impl<B: Backend + Mamba3DoubleSsdBackendExt + Mamba3SingleSsdBackendExt, C: CheckpointStrategy>
64        Mamba3BackendExt for Autodiff<B, C>
65    {
66    }
67}
68
69/// Public re-exports for Mamba-3.
70pub mod prelude {
71    #[cfg(feature = "autodiff")]
72    pub use super::Mamba3AutodiffBackendExt;
73    pub use super::Mamba3BackendExt;
74    use super::*;
75
76    pub use cache::{Mamba3Cache, Mamba3Caches};
77    pub use mamba3::{Mamba3, Mamba3Config};
78    pub use moments::Mamba3MomentsInput;
79    pub use quat_scan::Mamba3QuatScanBackendExt;
80    pub use rotation::{RotationKind, RotationSeq, RotationState};
81    pub use ssd_path::Mamba3SsdPath;
82}
burn_mamba/mamba3/mod.rs

burn_mamba/mamba3/
mod.rs
