Module moments 

Physical-frame state moments for Mamba-3 (serial chunkwise, exact)

The Mamba-3 state is complex (paper Props. complex-to-real / rope-trick): the real cache state h̃ₜ (ssm_bhpr) carries the rotations absorbed into B̃/C̃, and the physical state — what raw, un-rotated C reads — is the per-token de-rotation cₜ = Dₜ†h̃ₜ (see RotationSeq::derotate_states). The shipped observable/penalty is the Hermitian PR of M_phys = Σₜ cₜᴴcₜ (StateMoments::pr_complex).

Unlike the Mamba-2 moments (mamba2/ssd/moments.rs), no closed form exists: in M_phys[a,b] the token index couples to the matrix entry through the relative phase e^{i(φ_{t,b} − φ_{t,a})}, so the per-token phases do not factor out of the Gram kernel — every exact factorisation reproduces the per-token state as an intermediate. Exact M_phys therefore costs materialised states, affordable chunk-locally in the SerialRecalculated discipline. Per chunk (serial over n, one small (m2, m1) accumulator pair carried):

1. Chunk-local decay dₜ = exp(a_cum[t]) and 1-semiseparable mask L[t,j] = exp(a_cum[t] − a_cum[j]) (as SSD Step 1).
2. One chunk of cache-frame states, folding the combined-injection channel axis M into the write index: h̃[t] = dₜ·h₋ + Σ_{j≤t,m} L[t,j] · x̂[j,m] ⊗ b̂[j,m] — a broadcast product plus one batched GEMM; the transients are [batch, nheads, l, l·M, per_head_dim] and one chunk of states [batch, nheads, l, per_head_dim, state_rank].
3. De-rotate the states’ r axis to the physical frame (RotationSeq::derotate_states — θ-differentiable, which is what lets a PR penalty shape the rotation itself).
4. Accumulate the plain real m2/m1 sums, masked to valid_len.

The combined injections realise the trapezoid as one scalar-decay system with a 2·mimo_rank channel axis: x̂ = concat_m(v_γ, v_β), b̂ = concat_m(B̃, B̃_prev) (the β stream shift-before-chunking, its first element from the cache’s k_state/v_state), shared log-decay da. The initial state is the cache’s ssm_bhpr (+ optional learnable h₀), counted exactly once — never the single-ssd kernel’s seed-augmented initial state, whose boundary-β term the β stream already carries.

The chunk carry h₋ is read from the unmasked last position (zero pads: Δ = 0 ⇒ identity decay, zero write, so the state is carried through unchanged); the valid_len mask alone excludes pad garbage from the sums (pad rotation values are irrelevant — masked before accumulation).

Plain autodiff over the serial loop retains every chunk’s states (the full per-token trajectory) — the study-scale mode. The at-scale execution model is a custom recompute backward in the SerialRecalculated pattern (module backward, milestone 3).

Modules

backward 🔒

Custom autodiff node for the physical-frame state moments

recalculated 🔒

Physical-frame moments with a custom, memory-efficient backward

Structs

Mamba3MomentsInput

Inputs of the physical-frame state-moments computation — the combined (γ + β) injections of one chunked Mamba-3 forward, plus the per-token cumulative rotation. Built at either pathway’s pre-SSD seam; C/D are never read. See the module header for the math.

Traits

Mamba3MomentsBackendExt

Extends the backend with the physical-frame state-moments computation.