Struct Mamba3SsdInput 

pub struct Mamba3SsdInput<B: Backend> {
    pub v_bnlrhp: Tensor<B, 6>,
    pub da_bnlh: Tensor<B, 4>,
    pub b_bnlrhn: Tensor<B, 6>,
    pub c_bnlrhn: Tensor<B, 6>,
    pub initial_state_bhpr: Tensor<B, 4>,
    pub init_state_hpr: Option<Tensor<B, 3>>,
}

MIMO-first SSD input.

All tensors are pre-processed: B/C are already QK-normed, RoPE-applied, bias-added, and expanded to per-head (not per-group). V is already scaled by the trapezoidal coefficient (γ or β). The combined log-decay da = Δ·A is pre-computed. D skip is handled by the caller.

Fields

v_bnlrhp: Tensor<B, 6>

Value tensor, already scaled by trapezoidal coefficient (γ or β).

Shape

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

da_bnlh: Tensor<B, 4>

Pre-combined log-decay Δ·A (negative).

Shape

• [batch, nchunks, chunk_len, nheads]

b_bnlrhn: Tensor<B, 6>

Key/B tensor: QK-normed, RoPE-applied, bias-added, expanded to per-head, per-rank.

Shape

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

c_bnlrhn: Tensor<B, 6>

Query/C tensor: same processing as B.

Shape

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

initial_state_bhpr: Tensor<B, 4>

Initial SSM hidden state.

Shape

• [batch, nheads, per_head_dim, state_rank]

init_state_hpr: Option<Tensor<B, 3>>

Optional learnable initial state (broadcast over batch).

Shape

• [nheads, per_head_dim, state_rank]

Implementations

impl<B: Backend> Mamba3SsdInput<B>

pub fn sanity(&self)