Struct Mamba2SsdInput 

pub struct Mamba2SsdInput {
    pub x_bnlhp: Tensor<5>,
    pub dt_bnlh: Tensor<4>,
    pub a_decay_h: Tensor<1>,
    pub b_bnlhr: Tensor<5>,
    pub c_bnlhr: Tensor<5>,
    pub d_h: Tensor<1>,
    pub initial_state_bhpr: Tensor<4>,
    pub init_state_hpr: Option<Tensor<3>>,
}

SSD input.

All tensors are pre-processed: B/C are already GQA-expanded to per-head.

Fields

x_bnlhp: Tensor<5>

Shape

• [batch, nchunks, chunk_len, nheads, per_head_dim]

dt_bnlh: Tensor<4>

Shape

• [batch, nchunks, chunk_len, nheads]

a_decay_h: Tensor<1>

Shape

• [nheads]

b_bnlhr: Tensor<5>

B tensor, expanded to per-head.

Shape

• [batch, nchunks, chunk_len, nheads, state_rank]

c_bnlhr: Tensor<5>

C tensor, expanded to per-head.

Shape

• [batch, nchunks, chunk_len, nheads, state_rank]

d_h: Tensor<1>

Shape

• [nheads]

initial_state_bhpr: Tensor<4>

Shape

• [batch, nheads, per_head_dim, state_rank]

init_state_hpr: Option<Tensor<3>>

Shape

• [nheads, per_head_dim, state_rank]

Implementations

impl Mamba2SsdInput

pub fn ssd_minimal(self) -> (Tensor<5>, Tensor<4>)

Minimal chunkwise SSD algorithm.

Implements the four-step decomposition of the semiseparable matrix multiplication described in §4 of the paper. The sequence of length is split into nchunks = ⌈sequence/chunk_len⌉ chunks of length chunk_len.

The four steps

Step 1 — Intra-chunk outputs (Y_diag)

Within each chunk, compute the output assuming the initial hidden state is zero. This is the quadratic attention form of the SSD layer restricted to a window of chunk_len tokens (§4.1):

  Y_diag[n] = (L[n] ∘ C[n] B[n]ᵀ) · X[n]

where L[n] is the chunk_len×chunk_len 1-semiseparable mask for chunk n. This step is a batched GEMM (exploits tensor cores).

Step 2 — Chunk state (state_bnhpr)

Compute the final SSM state of each chunk assuming zero initial state (§4.1, Eq. 20):

  s[n] = Σ_{t ∈ chunk n}  exp(A_cum[end] - A_cum[t]) · B̄[t] · x[t]ᵀ

This is also a batched GEMM and is fully parallel across chunks.

Step 3 — Inter-chunk state scan (state passing)

Propagate the true hidden state across chunk boundaries using the recurrence (§4.1, Eq. 22):

  h[n] = Ā[n]_chunk · h[n-1] + s[n]

where Ā[n]_chunk = exp(Σ_{t ∈ chunk n} Δₜ · A) is the cumulative decay over the whole chunk. This step is implemented as a single batched matrix multiplication using the 1-semiseparable structure of the inter-chunk decay matrix (same segsum trick, now over chunks). The scan has length nchunks = sequence/chunk_len rather than sequence, so its cost is negligible for typical chunk sizes.

Step 4 — State-to-output (Y_off)

For each chunk n, compute the contribution of the true initial state h[n-1] to the outputs within that chunk (§4.1, Eq. 23):

  Y_off[n, t] = C[n, t]ᵀ · exp(A_cum[t]) · h[n-1]

This is again a batched GEMM.

Final output (with D skip-connection)

  Y = Y_diag + Y_off + D · X

impl Mamba2SsdInput

pub fn detached(&self) -> Self

A value-identical copy detached from any autodiff graph (used by the diagnostic wrapper so the moments branch records no backward nodes).

pub fn state_moments(&self, valid_len: usize) -> StateMoments

Exact pooled moments of every per-token SSM state (see the module header). valid_len is the unpadded token count — states at zero-pad positions are excluded. C/D are not read.

Returns raw sums; count = valid_len · per_head_dim samples per (batch, head) slice.

impl Mamba2SsdInput

pub fn ssd_serial(self) -> (Tensor<5>, Tensor<4>)

Forward pass for the Mamba-2 SSD module (serial-over-chunks form).

Returns:

• y_bnlhp.
• final_state_bhpr.

impl Mamba2SsdInput

pub fn ssd_serial_recalculated(self) -> (Tensor<5>, Tensor<4>)

Forward pass for the Mamba-2 SSD module (recompute-backward path).

Returns:

• y_bnlhp.
• final_state_bhpr.

impl Mamba2SsdInput

pub fn sanity(&self)

Run the NaN/Inf guards on every input tensor.

impl Mamba2SsdInput

pub fn run(self, path: &Mamba2SsdPath) -> (Tensor<5>, Tensor<4>)

Run the selected SSD algorithm on this input.

Dispatches by Mamba2SsdPath variant to ssd_minimal, ssd_serial, or ssd_serial_recalculated.

Returns

• y_bnlhp: [batch, nchunks, chunk_len, nheads, per_head_dim]
• final_state_bhpr: [batch, nheads, per_head_dim, state_rank]