Module serial_recalculated 

Serial SSD with a custom, memory-efficient backward (Mamba-2)

This is the SerialRecalculated path. The forward is the same five-kernel serial scan as super::super::serial, but it is routed through the Mamba2BackendExt trait so that Autodiff backends can substitute a custom backward that recomputes the per-chunk intermediates instead of storing them (see super::backward / super::combined_backward), trading a little extra compute for ~⅓ less training memory.

Every plain (non-autodiff) backend uses the trait’s default body, which simply replays the super::super::serial kernels K1–K5. The crate::impl_ssd_backend_ext_for_burn_backends! / crate::decl_ssd_autodiff_backend_ext! macros wire up the per-backend impls and the autodiff marker trait.

Traits

Mamba2AutodiffBackendExt

Marker for an autodiff-capable backend that also implements the SSD extension trait (so the custom memory-efficient backward is available).

Mamba2BackendExt

Extends the backend and wraps it for burn.

Functions

k1_ssd_chunk_cumsum 🔒

Primitive port of crate::mamba2::ssd::serial::k1_ssd_chunk_cumsum.

k2_ssd_bmm 🔒

Primitive port of crate::mamba2::ssd::serial::k2_ssd_bmm.

k3_ssd_chunk_state 🔒

Primitive port of crate::mamba2::ssd::serial::k3_ssd_chunk_state (lean: returns only the chunk-end state).

k4_ssd_state_passing 🔒

Primitive port of crate::mamba2::ssd::serial::k4_ssd_state_passing.

k5_ssd_chunk_scan 🔒

Primitive port of crate::mamba2::ssd::serial::k5_ssd_chunk_scan.