burn_mamba/mamba3/double_ssd/mod.rs
1//! # Double-SSD pathway (VikramLex-style)
2//!
3//! Realises the Mamba-3 trapezoidal recurrence as **two standard SSD calls**
4//! that reuse the Mamba-2-like kernels:
5//!
6//! - γ-SSM: `hᵞₜ = αₜ hᵞₜ₋₁ + γₜ Bₜ xₜ` (current token)
7//! - β-SSM: `hᵝₜ = αₜ hᵝₜ₋₁ + βₜ Bₜ₋₁ xₜ₋₁` (previous token)
8//! - `hₜ = hᵞₜ + hᵝₜ`.
9//!
10//! Simple and easy to verify, at the cost of ~2× the intra-chunk and
11//! chunk-state memory of the [`single_ssd`](crate::mamba3::single_ssd) pathway.
12
13/// The double-SSD cache (`ssm`/`k_state`/`v_state`/`cum_angle`; no conv cache).
14pub mod cache;
15/// `forward_double_ssd` / `step_double_ssd` and the RoPE helpers.
16pub mod double_ssd;
17/// The standard SSD kernels reused by both the γ and β passes.
18pub mod ssd;
19
20/// Public re-exports for the double-SSD pathway.
21pub mod prelude {
22 use super::*;
23 pub use cache::{
24 Mamba3DoubleSsdCache, Mamba3DoubleSsdCacheConfig, Mamba3DoubleSsdCaches,
25 Mamba3DoubleSsdCachesConfig,
26 };
27 #[cfg(feature = "autodiff")]
28 pub use ssd::Mamba3DoubleSsdAutodiffBackendExt;
29 pub use ssd::Mamba3DoubleSsdBackendExt;
30 pub use ssd::Mamba3DoubleSsdInput;
31}