burn_mamba/mamba3/single_ssd/ssd/serial_recalculated/
serial_recalculated.rs

1//! # Serial SSD with a custom, memory-efficient backward (Mamba-3 single-SSD)
2//!
3//! The `SerialRecalculated` path for the single-SSD pathway.  The forward is the
4//! same serial scan as [`super::super::serial`], routed through the
5//! [`Mamba3SingleSsdBackendExt`] trait so `Autodiff` backends substitute a
6//! custom backward that recomputes per-chunk intermediates rather than storing
7//! them (see [`super::backward`] / [`super::combined_backward`]).  Unlike the
8//! double-SSD form, the kernels here apply the trapezoid `gamma`/`scale` and the
9//! boundary-β seed internally, so the backward also returns `d_gamma`/`d_scale`.
10//!
11//! The default body runs under a generic backend `B`, where the high-level
12//! `Tensor` (pinned to `Dispatch`) is unavailable, so the K1–K5 math goes
13//! through the rank-tagged [`F`] primitive wrapper.  K1–K4 are mode-agnostic and
14//! reused from the double-SSD forward; only the single-SSD K5 (strict-lower
15//! intra-chunk + γ-correction) is owned here, and it is forward-only.
16
17#![allow(non_snake_case)]
18
19use crate::mamba3::double_ssd::ssd::serial_recalculated::{
20    k1_ssd_chunk_cumsum, k2_ssd_bmm, k3_ssd_chunk_state, k4_ssd_state_passing,
21};
22use crate::mamba3::single_ssd::prelude::*;
23use crate::mamba3::single_ssd::ssd::serial_recalculated::diag::y_diag_correction;
24use burn_stack::utils::fprim::{F, san};
25use burn::backend::tensor::FloatTensor;
26use burn::backend::*;
27use burn::backend::{Backend, Dispatch, backend_extension};
28use burn::tensor::Tensor;
29
30impl Mamba3SingleSsdInput {
31    /// MIMO-first single-ssd form Serial SSD with recalculated backward.
32    ///
33    /// Delegates the full K1–K5 (single-ssd) computation to
34    /// [`Mamba3SingleSsdBackendExt::single_ssd_serial_recalculated`], which can provide
35    /// a memory-efficient custom backward for supported backends (the Autodiff
36    /// wrapper) and falls back to the standard K1–K5 forward on others.
37    ///
38    /// # Returns
39    /// - `y_bnlmhp`:         `[batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim]`
40    /// - `final_state_bhpr`: `[batch, nheads, per_head_dim, state_rank]`
41    pub fn single_ssd_serial_recalculated(self) -> (Tensor<6>, Tensor<4>) {
42        let input = self;
43        input.sanity();
44        assert!(
45            input.init_state_hpr.is_none(),
46            "init_state_hpr not yet implemented for single_ssd_serial_recalculated"
47        );
48
49        let (y_bnlmhp, final_state_bhpr) =
50            <Dispatch as Mamba3SingleSsdBackendExt>::single_ssd_serial_recalculated(
51                input.v_bnlmhp.into_dispatch(),
52                input.da_bnlh.into_dispatch(),
53                input.b_bnlmhr.into_dispatch(),
54                input.c_bnlmhr.into_dispatch(),
55                input.gamma_bnlh.into_dispatch(),
56                input.scale_bnlh.into_dispatch(),
57                input.initial_state_bhpr.into_dispatch(),
58                input.siso_specialization,
59            );
60        let y_bnlmhp = Tensor::from_dispatch(y_bnlmhp);
61        let final_state_bhpr = Tensor::from_dispatch(final_state_bhpr);
62        (y_bnlmhp, final_state_bhpr)
63    }
64}
65
66/// Extends the backend for the memory-efficient single-ssd form serial SSD.
67///
68/// The default implementation runs K1–K5 using primitive tensor operations,
69/// reusing the mode-agnostic K1/K2/K3/K4 from the double-SSD forward and the
70/// single-ssd form K5 below. Backends that support a custom memory-efficient
71/// backward (the Autodiff wrapper) override this to recompute forward
72/// intermediates during backward instead of saving them.
73#[backend_extension(
74    Cpu:  cfg(feature = "backend-cpu"),
75    Cuda: cfg(feature = "backend-cuda"),
76    Rocm:  cfg(feature = "backend-rocm"),
77    Metal:  cfg(feature = "backend-metal"),
78    Vulkan:  cfg(feature = "backend-vulkan"),
79    Wgpu:  cfg(feature = "backend-wgpu"),
80    WebGpu:  cfg(feature = "backend-webgpu"),
81    Flex:  cfg(feature = "backend-flex"),
82    NdArray:  cfg(feature = "backend-ndarray"),
83    LibTorch:  cfg(any(feature = "backend-tch-cpu", feature = "backend-tch-gpu")),
84    Autodiff:  cfg(feature = "autodiff"),
85)]
86pub trait Mamba3SingleSsdBackendExt: Backend {
87    /// Memory-efficient MIMO single-ssd form serial SSD.
88    ///
89    /// # Arguments
90    /// - `v_bnlmhp`:           `[batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim]`
91    /// - `da_bnlh`:            `[batch, nchunks, chunk_len, nheads]` — pre-combined Δ·A
92    /// - `b_bnlmhr`:           `[batch, nchunks, chunk_len, mimo_rank, nheads, state_rank]`
93    /// - `c_bnlmhr`:           `[batch, nchunks, chunk_len, mimo_rank, nheads, state_rank]`
94    /// - `gamma_bnlh`:         `[batch, nchunks, chunk_len, nheads]` — `γₜ = λₜ Δₜ`
95    /// - `scale_bnlh`:         `[batch, nchunks, chunk_len, nheads]` — `scaleₜ = γₜ + (1−λₜ₊₁)Δₜ₊₁`
96    /// - `initial_state_bhpr`: `[batch, nheads, per_head_dim, state_rank]`
97    /// - `siso_specialization`: allow the specialized `mimo_rank == 1`
98    ///   γ-correction (performance-only; see
99    ///   [`Mamba3Config::siso_specialization`](crate::mamba3::mamba3::Mamba3Config::siso_specialization))
100    ///
101    /// # Returns
102    /// - `y_bnlmhp`:         `[batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim]`
103    /// - `final_state_bhpr`: `[batch, nheads, per_head_dim, state_rank]`
104    #[allow(clippy::too_many_arguments)]
105    fn single_ssd_serial_recalculated(
106        v_bnlmhp: FloatTensor<Self>,
107        da_bnlh: FloatTensor<Self>,
108        b_bnlmhr: FloatTensor<Self>,
109        c_bnlmhr: FloatTensor<Self>,
110        gamma_bnlh: FloatTensor<Self>,
111        scale_bnlh: FloatTensor<Self>,
112        initial_state_bhpr: FloatTensor<Self>,
113        siso_specialization: bool,
114    ) -> (FloatTensor<Self>, FloatTensor<Self>) {
115        // Default impl: replicate the single-ssd form K1–K5 on primitives.
116        let v_bnlmhp = F::<Self, 6>::new(v_bnlmhp);
117        let da_bnlh = F::<Self, 4>::new(da_bnlh);
118        let b_bnlmhr = F::<Self, 6>::new(b_bnlmhr);
119        let c_bnlmhr = F::<Self, 6>::new(c_bnlmhr);
120        let gamma_bnlh = F::<Self, 4>::new(gamma_bnlh);
121        let scale_bnlh = F::<Self, 4>::new(scale_bnlh);
122        let initial_state_bhpr = F::<Self, 4>::new(initial_state_bhpr);
123
124        // K1 — chunk cumulative decay.
125        let (da_cumsum_bhnl, da_chunk_end_bhn) = k1_ssd_chunk_cumsum::<Self>(da_bnlh);
126        san(&da_cumsum_bhnl);
127
128        // K2 — CB matrix on unscaled B/C.
129        let cb_bnhLMLM = k2_ssd_bmm::<Self>(c_bnlmhr.clone(), b_bnlmhr.clone());
130        san(&cb_bnhLMLM);
131
132        // K3 — chunk state on K_scaled = scaleₜ · B.
133        let scale_bnlh11 = scale_bnlh.clone().unsqueeze_dims::<6>(&[3, 5]);
134        let k_scaled_bnlmhr = b_bnlmhr.clone() * scale_bnlh11;
135        let intra_chunk_state_bnhpr =
136            k3_ssd_chunk_state::<Self>(v_bnlmhp.clone(), k_scaled_bnlmhr, da_cumsum_bhnl.clone());
137        san(&intra_chunk_state_bnhpr);
138
139        // K4 — sequential state passing.
140        let (chunk_input_state_bnhpr, final_state_bhpr) = k4_ssd_state_passing::<Self>(
141            intra_chunk_state_bnhpr,
142            da_chunk_end_bhn,
143            initial_state_bhpr,
144        );
145        san(&chunk_input_state_bnhpr);
146        san(&final_state_bhpr);
147
148        // K5 — single-ssd form chunk scan.
149        let y_bnlmhp = k5_single_ssd_chunk_scan::<Self>(
150            da_cumsum_bhnl,
151            v_bnlmhp,
152            c_bnlmhr,
153            b_bnlmhr,
154            cb_bnhLMLM,
155            gamma_bnlh,
156            scale_bnlh,
157            chunk_input_state_bnhpr,
158            siso_specialization,
159        );
160        san(&y_bnlmhp);
161
162        (y_bnlmhp.inner(), final_state_bhpr.inner())
163    }
164}
165
166burn_stack::decl_autodiff_backend_ext!(Mamba3SingleSsdAutodiffBackendExt, Mamba3SingleSsdBackendExt);
167
168// ---------------------------------------------------------------------------
169// Per-backend impls: each delegates to the trait's default (K1–K5) body. The
170// custom autodiff backward lives in `super::backward` as a separate impl.
171// ---------------------------------------------------------------------------
172burn_stack::impl_backend_ext_for_burn_backends!(Mamba3SingleSsdBackendExt);
173
174// ─── Primitive forward K5 (single-ssd) ───────────────────────────────────────
175// Primitive port of [`super::super::serial::k5_single_ssd_chunk_scan`]. Combines
176// the strict-lower intra-chunk (LOWER), the γ-weighted same-step (DIAG), and the
177// state-to-output (BLUE/Y_off) contributions. Forward-only; the backward
178// computes these gradients analytically. K1–K4 are reused from the double-SSD
179// forward (mode-agnostic).
180
181/// SingleSsd chunk scan, on primitives.
182///
183/// - **Strict lower triangular intra-chunk** (`t1 > t2`):
184///   `(cb[i,j] · scale[t2] · exp(cumA[t1] − cumA[t2])) · V[t2]`
185/// - **Same-time-step (`t1 == t2`) γ-correction**:
186///   `γ[t] · (Σₙ C[t,r_out,n] · B[t,r_in,n]) · V[t,r_in,p]`
187/// - **State-to-output (Y_off)**: `exp(cumA[t]) · C[t] · h'[n-1]`
188///
189/// `cb_bnhLMLM` is the unscaled `C · Bᵀ` from K2; `b_bnlmhr` is the unscaled
190/// K/B tensor (γ-correction matmul).
191#[allow(clippy::too_many_arguments)]
192fn k5_single_ssd_chunk_scan<B: Backend>(
193    da_cumsum_bhnl: F<B, 4>,
194    v_bnlmhp: F<B, 6>,
195    c_bnlmhr: F<B, 6>,
196    b_bnlmhr: F<B, 6>,
197    cb_bnhLMLM: F<B, 5>,
198    gamma_bnlh: F<B, 4>,
199    scale_bnlh: F<B, 4>,
200    chunk_input_state_bnhpr: F<B, 5>,
201    siso_specialization: bool,
202) -> F<B, 6> {
203    let [batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim] = v_bnlmhp.dims();
204    let [.., state_rank] = c_bnlmhr.dims();
205    let device = v_bnlmhp.device();
206    let dtype = v_bnlmhp.dtype();
207    let fused = chunk_len * mimo_rank;
208
209    // Fuse mimo_rank into chunk_len for the SSM-style matmul.
210    let v_bnLMhp = v_bnlmhp
211        .clone()
212        .reshape([batch, nchunks, fused, nheads, per_head_dim]);
213    let c_bnLMhr = c_bnlmhr
214        .clone()
215        .reshape([batch, nchunks, fused, nheads, state_rank]);
216
217    // Per-fused-step cumulative decay (interleave-expand the base grid).
218    let da_cumsum_bhnLM = da_cumsum_bhnl
219        .unsqueeze_dim::<5>(4)
220        .expand([batch, nheads, nchunks, chunk_len, mimo_rank])
221        .reshape([batch, nheads, nchunks, fused]);
222
223    // ── Y_off: exp(cumA[t]) · C[t] · h'[n-1] ────────────────────────────────
224    let exp_da_bnhLMp = da_cumsum_bhnLM
225        .clone()
226        .exp()
227        .swap_dims(1, 2) // bnhLM
228        .unsqueeze_dim::<5>(4) // bnhLM1
229        .expand([batch, nchunks, nheads, fused, per_head_dim]);
230    let c_bnhLMr = c_bnLMhr.swap_dims(2, 3);
231    let chunk_input_state_bnhrp = chunk_input_state_bnhpr.transpose();
232    let ch_bnhLMp = c_bnhLMr.matmul(chunk_input_state_bnhrp);
233    let y_off_bnhLMp = ch_bnhLMp * exp_da_bnhLMp;
234
235    // ── Y_lower: strict lower-tri intra-chunk with scale and decay ──────────
236    let da_cumsum_bnhLM = da_cumsum_bhnLM.swap_dims(1, 2); // bnhLM
237    let target_da_cumsum_bnhLMLM = da_cumsum_bnhLM
238        .clone()
239        .unsqueeze_dim::<5>(4) // bnhLM1
240        .expand([batch, nchunks, nheads, fused, fused]);
241    let source_da_cumsum_bnhLMLM = da_cumsum_bnhLM
242        .unsqueeze_dim::<5>(3) // bnh1LM
243        .expand([batch, nchunks, nheads, fused, fused]);
244    let diff_bnhLMLM = target_da_cumsum_bnhLMLM - source_da_cumsum_bnhLMLM;
245
246    // Strict-upper -inf mask on the base time grid (`t1 <= t2` → -inf), then
247    // interleave-expand to fused length so MIMO same-time blocks are zeroed.
248    let inf_upper_bnhLMLM =
249        F::<B, 2>::full([chunk_len, chunk_len], f32::NEG_INFINITY, &device, dtype)
250            .triu(0) // upper triangle INCLUDING diagonal
251            .unsqueeze_dims::<5>(&[0, 1, 2])
252            .expand([batch, nchunks, nheads, chunk_len, chunk_len])
253            .unsqueeze_dim::<6>(4)
254            .expand([batch, nchunks, nheads, chunk_len, mimo_rank, chunk_len])
255            .reshape([batch, nchunks, nheads, fused, chunk_len])
256            .unsqueeze_dim::<6>(5)
257            .expand([batch, nchunks, nheads, fused, chunk_len, mimo_rank])
258            .reshape([batch, nchunks, nheads, fused, fused]);
259    let decay_strict_bnhLMLM = (diff_bnhLMLM + inf_upper_bnhLMLM).exp();
260
261    // Per-column scale: `scale[t2]` lives on the source axis (column).
262    let scale_bnhLM = scale_bnlh
263        .transpose() // bnhl
264        .unsqueeze_dim::<5>(4) // bnhl1
265        .expand([batch, nchunks, nheads, chunk_len, mimo_rank])
266        .reshape([batch, nchunks, nheads, fused]);
267    let scale_col_bnhLMLM = scale_bnhLM
268        .unsqueeze_dim::<5>(3) // bnh1LM
269        .expand([batch, nchunks, nheads, fused, fused]);
270
271    let kernel_bnhLMLM = decay_strict_bnhLMLM * scale_col_bnhLMLM;
272    let masked_cb_bnhLMLM = cb_bnhLMLM * kernel_bnhLMLM;
273    let v_bnhLMp = v_bnLMhp.swap_dims(2, 3);
274    let y_lower_bnhLMp = masked_cb_bnhLMLM.matmul(v_bnhLMp);
275
276    // ── Y_diag: γ-weighted same-step correction ─────────────────────────────
277    let y_diag_bnlmhp = y_diag_correction::<B>(
278        v_bnlmhp,
279        b_bnlmhr,
280        c_bnlmhr,
281        gamma_bnlh,
282        siso_specialization,
283    );
284    let y_diag_bnLMhp = y_diag_bnlmhp.reshape([batch, nchunks, fused, nheads, per_head_dim]);
285    let y_diag_bnhLMp = y_diag_bnLMhp.swap_dims(2, 3);
286
287    // ── Combine and reshape ─────────────────────────────────────────────────
288    let y_bnhLMp = y_off_bnhLMp + y_lower_bnhLMp + y_diag_bnhLMp;
289    let y_bnLMhp = y_bnhLMp.swap_dims(2, 3);
290    y_bnLMhp.reshape([batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim])
291}
burn_mamba/mamba3/single_ssd/ssd/serial_recalculated/serial_recalculated.rs

burn_mamba/mamba3/single_ssd/ssd/serial_recalculated/
serial_recalculated.rs
