S1-31: SSM_SCAN_BACK — Mamba has a backward

[dillon-blake]

Stacked on #24 (S1-30). Mamba has a backward.

The selective scan's VJP: a reverse recurrence, walking t from n_t-1 down to 0 and carrying the state gradient backwards. Five gradients out of one op (d_s, d_x, d_dt, d_B, d_C), packed exactly as ggml_ssm_scan itself packs y with the final states. Both branches: Mamba-2 (one scalar decay per head) and Mamba-1 (one per state).

Two things that are easy to get wrong, and both are silent

1. grad is the WHOLE packed gradient — y and the final states — and the state region is not zero. MODE_GRAD's objective sums over the packed dst, so d(sum)/d(s_final) is 1, and a kernel that seeded the reverse recurrence with zeros would disagree with the finite difference and be right to. It seeds ds at t = n_t.

That isn't hypothetical: seeding with zeros is one of the five mutations below and it measures MAA 0.59. In training the region genuinely is zero — the cached state feeds nothing downstream of the loss — so honouring it costs nothing there, and it makes cross-ubatch BPTT nearly free later.

2. The forward overwrites its state in place, so s_{t-1} is gone by the time the backward needs it — and it does need it, for d(dt) via the dA path. The states are recomputed and stored, all n_t + 1 of them. Store-all: correct, O(n_t) memory, and the honest starting point for a CPU oracle. Checkpoint-every-K is the optimization and must be pinned bit-for-bit against this.

Threaded by sequence, not by head

dB and dC accumulate over every head in a group, so a head-partitioned kernel would have several threads writing the same (i0, g, t) and would need atomics — neither deterministic nor free. One thread per sequence owns every output it touches. Parallelism is n_seqs, which is small; correctness and determinism (ADR-0002) come first in the kernel the GPU ports will be measured against.

The oracle is a float64 finite difference of ggml's own forward, and it has to be

A reverse recurrence has many ways to be subtly wrong — a dropped dA path, a state read one token late, a missing seed — and a hand-written reference would share my derivation's bugs. So each of the five gradients is checked by perturbing every input element and re-running ggml_ssm_scan itself, under a non-uniform objective that includes the packed state region.

All five exact to ~1e-5 (float32 FD precision), in both branches.

Verification

max_maa_err = 5e-2, measured:

	MAA
worst FD noise, 20 runs	2.1e-2
drop the dA path from d(dt)	0.57
seed ds with zeros instead of the packed state gradient	0.59
drop y's own contribution to dS	0.44
use s_{t-1} where s_t belongs in dC	0.51
drop the softplus derivative	0.38

Five mutations injected, five caught. 2.4x above the noise, 7.5–12x below every defect. The noise is high because the recurrence is exponential in dt*A, so a finite difference of it amplifies its own rounding — which is precisely why the float64 reference is the oracle and this is the wiring-plus-sanity check.

Before S1-29b these tests called ggml_set_param zero times and every shape was above grad_nmax, so grad -o SSM_SCAN reported OK while the op had no backward at all.