Skip to main content

burn_mamba/mamba3/single_ssd/ssd/
mod.rs

1//! Standard MIMO-first SSD kernels specialised to the single-pass form.
2//!
3//! Three exact reformulations (Minimal / Serial / SerialRecalculated) that
4//! agree on values and gradients; selected via [`Mamba3SsdPath`] and dispatched
5//! by [`Mamba3SingleSsdInput::run`].  These take the raw `v` plus `gamma`/`scale`
6//! so the kernel applies the trapezoid weights and boundary seed internally.
7//!
8//! [`Mamba3SsdPath`]: crate::mamba3::ssd_path::Mamba3SsdPath
9
10/// Matmul/`segsum` MIMO-first SSD with plain autodiff backward.
11pub mod minimal;
12/// Serial-over-chunks SSD with plain autodiff backward.
13pub mod serial;
14/// Serial-over-chunks SSD with a custom recompute backward.
15pub mod serial_recalculated;
16/// The [`Mamba3SingleSsdInput`] bundle and its `run()` dispatch.
17pub mod ssd_path;
18
19#[cfg(feature = "autodiff")]
20pub use serial_recalculated::Mamba3SingleSsdAutodiffBackendExt;
21pub use serial_recalculated::Mamba3SingleSsdBackendExt;
22pub use ssd_path::Mamba3SingleSsdInput;