burn_mamba/mamba3/single_ssd/ssd/
diag.rs

1//! # Same-step γ-correction (the single-SSD diagonal term)
2//!
3//! The single-SSD recurrence scales `K` by `scaleₜ = γₜ + (1−λₜ₊₁)Δₜ₊₁`, which is
4//! the right weight for every source step `s < t` but *not* for the same step
5//! `s = t`, where the weight must be `γₜ`. The intra-chunk path therefore masks
6//! the diagonal out (strict lower triangle) and this module adds it back:
7//!
8//! ```text
9//!   y_diag[t, m_out, h, p] = γₜ · Σ_{m_in} (Σ_r C[t, m_out, h, r] · B[t, m_in, h, r])
10//!                                        · V[t, m_in, h, p]
11//! ```
12//!
13//! It is computed fresh (a small same-step product) rather than extracted from
14//! the block diagonal of the fused `L·M` CB matrix, which would need a fiddly
15//! reshape.
16//!
17//! ## SISO fast path
18//!
19//! The inner `m × m` Gram matrix is what makes this a pair of matmuls. At
20//! `mimo_rank == 1` it collapses to the **scalar** `Cₜ·Bₜ`, so both matmuls
21//! degenerate into `1×r×1` and `1×1×p` GEMMs — thousands of tiny batched
22//! products, one per `(batch, nchunks, chunk_len, nheads)`.
23//! `y_diag_correction_siso` instead contracts `state_rank` with a reduction
24//! and folds the result (together with `γₜ`) in as a per-`(b, n, l, h)` scalar
25//! broadcast. Both branches compute the same quantity; only the op mix differs,
26//! so [`Mamba3Config::siso_specialization`](crate::mamba3::mamba3::Mamba3Config::siso_specialization)
27//! can force the general branch at `mimo_rank == 1` to measure the difference.
28//!
29//! Reference kernels:
30//! - SISO: `refs/state-spaces/mamba/mamba_ssm/ops/triton/mamba3/mamba3_siso_fwd.py`
31//! - MIMO: `refs/state-spaces/mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py`
32
33#![allow(non_snake_case)]
34
35use burn::prelude::*;
36
37/// The γ-weighted same-step correction `y_diag`, dispatching to the SISO fast
38/// path (`y_diag_correction_siso`) or the general MIMO path
39/// (`y_diag_correction_mimo`) on `mimo_rank`.
40///
41/// `siso_specialization` is
42/// [`Mamba3Config::siso_specialization`](crate::mamba3::mamba3::Mamba3Config::siso_specialization):
43/// `false` keeps the general branch even at `mimo_rank == 1`.
44///
45/// # Shapes
46/// - `v_bnlmhp`: `[batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim]`
47/// - `b_bnlmhr`, `c_bnlmhr`: `[batch, nchunks, chunk_len, mimo_rank, nheads, state_rank]`
48/// - `gamma_bnlh`: `[batch, nchunks, chunk_len, nheads]`
49/// - returns `y_diag_bnlmhp`: `[batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim]`
50pub fn y_diag_correction(
51    v_bnlmhp: Tensor<6>,
52    b_bnlmhr: Tensor<6>,
53    c_bnlmhr: Tensor<6>,
54    gamma_bnlh: Tensor<4>,
55    siso_specialization: bool,
56) -> Tensor<6> {
57    let [.., mimo_rank, _nheads, _per_head_dim] = v_bnlmhp.dims();
58    if mimo_rank == 1 && siso_specialization {
59        y_diag_correction_siso(v_bnlmhp, b_bnlmhr, c_bnlmhr, gamma_bnlh)
60    } else {
61        y_diag_correction_mimo(v_bnlmhp, b_bnlmhr, c_bnlmhr, gamma_bnlh)
62    }
63}
64
65/// SISO (`mimo_rank == 1`) `y_diag`: `qk_dot` is the scalar `Σ_r Cₜ·Bₜ` per
66/// `(b, n, l, h)`, so both MIMO matmuls become a reduction plus broadcasts.
67pub(crate) fn y_diag_correction_siso(
68    v_bnlmhp: Tensor<6>,
69    b_bnlmhr: Tensor<6>,
70    c_bnlmhr: Tensor<6>,
71    gamma_bnlh: Tensor<4>,
72) -> Tensor<6> {
73    let qk_dot_bnlmh1: Tensor<6> = (c_bnlmhr * b_bnlmhr).sum_dim(5);
74    let gamma_bnl1h1 = gamma_bnlh.unsqueeze_dims::<6>(&[3, 5]);
75    v_bnlmhp * qk_dot_bnlmh1 * gamma_bnl1h1
76}
77
78/// General MIMO `y_diag`: the `m_out × m_in` Gram matrix `C · Bᵀ` (contracted
79/// over `state_rank`) applied to `V`, then scaled by `γₜ`.
80pub(crate) fn y_diag_correction_mimo(
81    v_bnlmhp: Tensor<6>,
82    b_bnlmhr: Tensor<6>,
83    c_bnlmhr: Tensor<6>,
84    gamma_bnlh: Tensor<4>,
85) -> Tensor<6> {
86    // c_bnlmhr [b, n, l, m, h, r] -> c_bnlhmr [b, n, l, h, m, r]
87    // b_bnlmhr [b, n, l, m, h, r] -> b_bnlhrm [b, n, l, h, r, m]
88    let c_bnlhmr = c_bnlmhr.swap_dims(3, 4);
89    let b_bnlhrm = b_bnlmhr.permute([0, 1, 2, 4, 5, 3]);
90    // qk_dot_bnlhmM [b, n, l, h, m_out, m_in]
91    let qk_dot_bnlhmM = c_bnlhmr.matmul(b_bnlhrm);
92
93    // V in [b, n, l, h, m_in, p] layout for the next matmul, then (qk_dot) · V.
94    let v_bnlhmp = v_bnlmhp.swap_dims(3, 4);
95    let y_d_bnlhmp = qk_dot_bnlhmM.matmul(v_bnlhmp);
96
97    // Multiply by γₜ (per (batch, nchunks, chunk_len, nheads)), back to bnlmhp.
98    let gamma_bnlh11 = gamma_bnlh.unsqueeze_dims::<6>(&[4, 5]);
99    (y_d_bnlhmp * gamma_bnlh11).swap_dims(3, 4)
100}
101
102// ---------------------------------------------------------------------------
103// Tests — SISO fast path ≡ MIMO-general path
104// ---------------------------------------------------------------------------
105
106#[cfg(all(test, feature = "_dev-test"))]
107mod tests;
burn_mamba/mamba3/single_ssd/ssd/diag.rs

burn_mamba/mamba3/single_ssd/ssd/
diag.rs
