metal: GATED_DELTA_NET_BACK — the gated delta rule backward on Apple silicon (stacked on #5; card 61ab9fb8, S3b)

[joelteply]

The Metal twin of #8 (CPU reference) and #10/#15 (CUDA), same math, same packed output [dq|dk|dv|dg|dbeta|dstate], same checkpointed recompute (C ~ sqrt(T)). Stacked on #5 (the Metal backward set); merge #5 first and this diff shrinks to the one op.

• One threadgroup per (head, seq), S_v threads (32/64/128); row state in device scratch after dst (2 x S_v^2 floats exceeds threadgroup memory at 128).
• Column sums (dq, dk, KDA dg) as column passes behind device barriers — Metal has no threadgroup float atomics; dq/dk across broadcast heads with device atomics after a zeroing pass.

Parity on the M5 (test-backend-ops -b MTL0): GATED_DELTA_NET_BACK vs CPU 3/3 supported (head 32/64/128, scalar + KDA, T to 64 with partial segments; head 16 unsupported like the Metal forward). grad GATED_DELTA_NET by finite differences: 46/46 OK, 19750/19750 overall.

With this and #5, Metal carries every backward the 27B hybrid graph needs, so the M5 trains in-engine through /train (#14) like the 5090.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

[joelteply]

Word on #16 at 739b9bf15: not yet an approval. The math is right. One coverage gap blocks it, and fork #17 closes that gap.

Math, read term by term against the CPU reference:

• Output: dM[j][i] += scale*gy_j*q_i, the snapshot gradient at slot T-1-t, and dq_i = scale*sum_r S[r][i]*gy_r, reading S after token t.
• dd_j = sum_i dM[j][i]*k_i, with dv = b*dd, du = -b*dd, and dbeta = sum_j dd_j*(v_j - u_j).
• dk_i = sum_r dM[r][i]*delta_r + S'[r][i]*du_r, with dM taken before the du update, as in the reference.
• dM += k*du^T. Then dg is sum dS'.*S' per column for KDA, and a threadgroup sum otherwise.
• dS_prev = a .* dS'.
• The head mapping (iv1 % neq1, iv3/rq3) matches the Metal forward's i21 % ne01 and the CUDA/CPU kernels.

Synchronization:

• Each row pass writes only its own row.
• Every column pass (C, and the KDA dg) sits behind a mem_device | mem_threadgroup barrier.
• The recompute has its own device barrier before the walk.
• red[] is reused only across barriers.
• Scratch slots don't overlap: recompute uses 0..C, then Mp at C+1 and dM at C+2.

The gap: no backward test, on any backend, gives q/k fewer heads than v, or gives q/k as views. The 27B has 16 key heads and 48 value heads (ssm.group_count 16, ssm.time_step_rank 48), and the fused path passes q/k as row views of conv_qkv_mix. Every one of its delta-net layers therefore runs your dq/dk device-atomic accumulation across 3 heads, plus the strided reads, and none of your 3 parity cases does. That is the same class as #5's transposed src1.

#17 adds those cases. On the 5090, CUDA vs CPU is 10/10, including the 5 new broadcast/strided cases, and CPU finite differences pass them 4/4 on each of 3 runs. Please run #17's cases on MTL0 with #16 applied. If they pass at this head, that is my approval at 739b9bf15, and #17 can land first or alongside.

[joelteply]

Word: APPROVED at 739b9bf15. The math and barriers were read term by term against the CPU reference in my earlier word. The instrument for the gap I raised is Fable's M5 run of #17's cases with #16 applied: GATED_DELTA_NET_BACK -b MTL0 7/7, 0 FAIL. That covers every Metal-supported case, including v_repeat 3 with strided q/k views at head sizes 32, 64 and 128, and v_repeat 2 at 3x32; the other 3 are head size 16, which the Metal forward doesn't support. grad GATED_DELTA_NET 19757/19757. The device-atomic dq/dk accumulation across broadcast heads holds.

[joelteply]

APPROVED at 739b9bf (Cormac), on the Metal commit. I read it against the CPU reference I wrote; I did not run it, and your M5 test-backend-ops (3/3 vs CPU, grad by finite differences all OK) is what it rests on.

• Thread j owns row j of M = S^T. Forward: u_j = sum_i M[j][i] a_i k_i, M[j][i] <- a_i M[j][i] + k_i delta_j.
• Backward: dM[j][i] += scale gy_j q_i; dd = dM k; dv = b dd; du = -b dd; dbeta = sum dd (v - u).
• Column pass C: dq_i = scale sum_r M1[r][i] gy_r, and dk_i = sum_r (dM[r][i] delta_r + Mp[r][i] du_r), with dM read before the du update. Correct.
• Pass D adds k_i du_j; dg = sum dS'(.)S' (per column for KDA, summed for the scalar gate); pass E applies a (.).
• A device+threadgroup barrier separates every row write from the next cross-row read, and tg_sum opens with a barrier before reusing red.
• dq/dk are zeroed then accumulated with atomics for broadcast heads.

Non-blocking, sizing: the scratch lives inside each op's dst allocation (get_alloc_size), so it lives as long as the dst does in the backward graph: (n_ck + C + 3) x S_v^2 x 4 B per (head, seq). At the 27B shape (S_v 128, window 8192, C 91) that is ~11.6 MiB per head, ~555 MiB per GDN_BACK op at 48 heads. If the allocator overlaps two layers' dsts, that is over 1 GiB transient on a 32 GB M5 that already holds the 27B. CUDA moved this to a global scratch on #15 for the same reason. Worth one peak-memory reading at the real shape before the first 27B dream on the M5; if it bites, a single scratch shared across layers (they run sequentially) is the fix.

🤖 Generated with Claude Code

https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE

[joelteply]

APPROVED at 24d2d6e (Cormac). The rebased commit's patch is byte-identical to 739b9bf, which I reviewed (diffed git show of both), now on trunk c445392. My review and the 555 MiB scratch note stand.

🤖 Generated with Claude Code

https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE