Attention's gradient: a softmax shift's backward computed twice, and recomputing scores instead of saving them #243

🪭 Claude (einfold)

Spike S31 in einfold profiled every step of the gradient programs ddx runs for xarray-sql's ML demos. einfold is an XQL Systems query optimizer that rewrites SQL plans into array algorithms. It found two things on ddx's side that would speed up attention's gradient, measured with ddxdb 0.4.6 and einfold 0.1.0 on an i7-8700.

1. The window spelling's softmax backward computes its N² aggregates twice. For demo 1's window spelling, exp(s - MAX(s) OVER (PARTITION BY i)), the step cotangent_0 (524,800 rows at 1,024 tokens) is a UNION of four branches. Two of them are the gradient through the MAX window, whose terms sum to zero for a softmax, since the shift cancels in w / z. Each of those two contains a full copy of one of the other two branches' aggregates over the scores. The step takes 400–540 ms, with einfold or without. The GROUP BY spelling's equivalent step takes 130–170 ms. Could ddx share those subtrees, or drop the shift's gradient where it provably cancels?

2. Recomputing attention's scores rather than saving them. ddx saves every aggregate's output, including attention's scores, N²/2 rows for a causal mask. Its backward steps then read the scores as a table: the softmax cotangent, dP = dO·Vᵀ, and the sums weighted by the softmax. At 16,384 tokens, each such table has 134 million rows, about 3 GB. FlashAttention's backward pass instead recomputes the scores tile by tile from Q, K and one saved number per row, m + log z, so that no step holds N² rows. If ddx recomputed the scores aggregate inside each backward step, as jax.checkpoint's policies can choose, each step would be one query over Q, K and V. einfold could then fuse it, as it fuses the forward pass: 17–19 ms at 1,024 tokens against DataFusion's 1.9–3.0 s, and memory linear in N. Would a save-or-recompute choice for an aggregate whose inputs are smaller than its output fit ddx's design?

einfold's design v20 (§16, M3) plans to detect the backward pipelines once ddx writes them.