Module ssd 

Mamba-2 chunkwise SSD algorithms

Three exact reformulations of the same Structured State Space Duality scan, all agreeing on values and gradients (asserted by the ssd_path tests):

• minimal — mostly batched matmuls + a segsum mask; plain autodiff backward.
• serial — a serial loop over chunks (mirrors the Triton kernels K1–K5); plain autodiff backward.
• serial_recalculated — the same serial loop with a custom, memory-efficient backward that recomputes intermediates (saves ~⅓ training memory).

ssd_path holds the Mamba2SsdPath selector and the Mamba2SsdInput bundle whose run() dispatches to one of the three.

Re-exports

pub use serial_recalculated::Mamba2AutodiffBackendExt;

pub use serial_recalculated::Mamba2BackendExt;

pub use ssd_path::Mamba2SsdInput;

pub use ssd_path::Mamba2SsdPath;

Modules

minimal

Matmul/segsum SSD with plain autodiff backward.

moments

Closed-form per-token state moments (for state-PR diagnostics/penalties).

serial

Serial-over-chunks SSD with plain autodiff backward.

serial_recalculated

Serial-over-chunks SSD with a custom recompute backward. Serial SSD with a custom, memory-efficient backward.

ssd_path

The Mamba2SsdPath selector and Mamba2SsdInput dispatch bundle.