Module diag 

The same-step γ-correction shared by all three algorithms.

Same-step γ-correction (the single-SSD diagonal term)

The single-SSD recurrence scales K by scaleₜ = γₜ + (1−λₜ₊₁)Δₜ₊₁, which is the right weight for every source step s < t but not for the same step s = t, where the weight must be γₜ. The intra-chunk path therefore masks the diagonal out (strict lower triangle) and this module adds it back:

  y_diag[t, m_out, h, p] = γₜ · Σ_{m_in} (Σ_r C[t, m_out, h, r] · B[t, m_in, h, r])
                                       · V[t, m_in, h, p]

It is computed fresh (a small same-step product) rather than extracted from the block diagonal of the fused L·M CB matrix, which would need a fiddly reshape.

SISO fast path

The inner m × m Gram matrix is what makes this a pair of matmuls. At mimo_rank == 1 it collapses to the scalar Cₜ·Bₜ, so both matmuls degenerate into 1×r×1 and 1×1×p GEMMs — thousands of tiny batched products, one per (batch, nchunks, chunk_len, nheads). y_diag_correction_siso instead contracts state_rank with a reduction and folds the result (together with γₜ) in as a per-(b, n, l, h) scalar broadcast. Both branches compute the same quantity; only the op mix differs, so Mamba3Config::siso_specialization can force the general branch at mimo_rank == 1 to measure the difference.

Reference kernels:

• SISO: refs/state-spaces/mamba/mamba_ssm/ops/triton/mamba3/mamba3_siso_fwd.py
• MIMO: refs/state-spaces/mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py

Functions

y_diag_correction

The γ-weighted same-step correction y_diag, dispatching to the SISO fast path (y_diag_correction_siso) or the general MIMO path (y_diag_correction_mimo) on mimo_rank.

y_diag_correction_mimo 🔒

General MIMO y_diag: the m_out × m_in Gram matrix C · Bᵀ (contracted over state_rank) applied to V, then scaled by γₜ.

y_diag_correction_siso 🔒

SISO (mimo_rank == 1) y_diag: qk_dot is the scalar Σ_r Cₜ·Bₜ per (b, n, l, h), so both MIMO matmuls become a reduction plus broadcasts.