burn_mamba/mamba2/ssd/
mod.rs

1//! # Mamba-2 chunkwise SSD algorithms
2//!
3//! Three exact reformulations of the same Structured State Space Duality scan,
4//! all agreeing on values **and** gradients (asserted by the `ssd_path` tests):
5//!
6//! - [`minimal`] — mostly batched matmuls + a `segsum` mask; plain autodiff
7//!   backward.
8//! - [`serial`] — a serial loop over chunks (mirrors the Triton kernels K1–K5);
9//!   plain autodiff backward.
10//! - [`serial_recalculated`] — the same serial loop with a **custom,
11//!   memory-efficient backward** that recomputes intermediates (saves ~⅓
12//!   training memory).
13//!
14//! [`ssd_path`] holds the [`Mamba2SsdPath`] selector and the [`Mamba2SsdInput`]
15//! bundle whose `run()` dispatches to one of the three.
16
17/// Matmul/`segsum` SSD with plain autodiff backward.
18pub mod minimal;
19/// Closed-form per-token state moments (for state-PR diagnostics/penalties).
20pub mod moments;
21/// Serial-over-chunks SSD with plain autodiff backward.
22pub mod serial;
23/// Serial-over-chunks SSD with a custom recompute backward.
24pub mod serial_recalculated;
25/// The [`Mamba2SsdPath`] selector and [`Mamba2SsdInput`] dispatch bundle.
26pub mod ssd_path;
27
28#[cfg(feature = "autodiff")]
29pub use serial_recalculated::Mamba2AutodiffBackendExt;
30pub use serial_recalculated::Mamba2BackendExt;
31pub use ssd_path::{Mamba2SsdInput, Mamba2SsdPath};
burn_mamba/mamba2/ssd/mod.rs

burn_mamba/mamba2/ssd/
mod.rs
