perf(ddx-ad): transpose causal attention's contractions without their joins - #162

Closes #153.

What changed. The contraction rule (#126, crates/ddx-ad/src/contraction.rs) differentiates a saved SUM(e_X · e_Y) over a join without rebuilding the join. Until now it missed both of causal attention's aggregates, which this PR now covers:

• A mask. The scores' join, q JOIN k ON q.d = k.d WHERE q.t >= k.t, carries a filter across the sides. The filter reads only grouping keys (q.t, k.t), so every row of a group passes it or none does. The transpose now applies a conjunct like this to the cotangent's rows.
• Sides of several inputs. The output, SUM(w.w / n.z * v.x) over weights ⋈ v ⋈ norm, joins three inputs. Inner joins can be regrouped, so the inputs are now split by which factors of the product read them. Here that's {weights, norm} (one row per score) and {v}, and each side is rebuilt as its own join. This is jax.grad's plan: transpose p = w / z per score, then p @ v. An input that no factor reads goes with the side its join conditions tie it to.

Each side is now transposed by the ordinary region rules (its projections, windows and inputs), starting from the cotangent of its factor. That cotangent is t, the contraction of the cotangent with the other side. This replaces #126's inlining of each side's factor down to one input's columns, which couldn't pass through a window: attention's weights side is a softmax window over the scores. t is most of the work. When a side reads it more than once (several inputs, or a window), it is materialized as a step.

Where attention's time went. Measured on the issue's query at 256 tokens with 64 features, timing each step of the program, on the same machine (i7-8700, 12 threads):

So the backward steps go from about 1,020 ms to 190 ms, 5.4× less. The issue's hand-written JAX plan runs its backward in about 1.3× its forward; ddx's is now about 2× its forward steps. The three gradients agree with main's to about 1e-12 relative.

ad_perf (the run column, best of 3, same machine) shows no regression elsewhere, and nn gains too:

matmul's sides are single inputs with no window, so its plan keeps #126's shape. Its difference from main is within run-to-run noise.

What could have gone wrong, and how the tests cover it. crates/ddx-datafusion/tests/contraction_transpose.rs compares each loss with a spelling that the general rule takes, on random tables with NULLs, NaNs, ±∞, missing rows and repeated keys. It also asserts which rule ran, by whether a gradient step rebuilds the join.

• The oracle. Its old spelling, WHERE a.s + w.o >= 0, reads only grouping keys, so it is now a mask and would take the fast rule too. It now reads a join key instead: a.k IS NULL OR a.k + w.o >= 0, which also keeps every row.
• Masks. a.s >= w.o, a mask beside a one-sided filter, and NULL grouping keys under a mask, where the mask is NULL and the row is dropped.
• A divisor joined on a grouping key. SUM(p.val / z.val * v.val), differentiated in all three tables, in three join orders. Also with a fourth table that no factor reads, with repeated keys, so it repeats rows on its side. Dividing by a zero, NULL or infinite z must give the same NaNs and infinities as the general rule.
• The issue's whole query: mask, MAX window, normalizer, with q, k and v all differentiated.
• Rankings. Both sides are recomputed, the one that doesn't move too (it builds t), so the whole region's windows and LIMITs must be total, as under the general rule. A constant side read through a LIMIT is refused (CodeRabbit caught that the first version checked only the moving side).
• A side in pieces. A side is the inputs its factors read, and its own conditions need not join them. In p JOIN v ON p.j = v.t JOIN z ON v.u = z.i with SUM(p.val / z.val * v.val), the side {p, z} would be rebuilt as p × z, a cross join the query never formed: quadratic where the forward pass is linear. Each side is now joined one input at a time, each tied by a condition to those before it, and the general rule applies when no order does that. A test pins it, and it fails on the earlier commit. (Found by a /code-review pass.)
• The existing cases (perf(ddx-ad): transpose a contraction without rebuilding its join #126's matrix products, two-key joins, the two-layer network, a factor of two columns) pass unchanged.

.github/scripts/mutation_test.py's two contraction mutants are updated to the new code, and contraction_transpose.rs still kills both. A third mutant that drops the mask from the cotangent survives: a cotangent only ever has rows for groups the query kept, so the mask is never what removes a row. I kept it anyway, because it makes the rule correct without relying on that property of how cotangents are built. It costs one filter on the cotangent.

cargo test -p ddx-ad -p ddx-datafusion passes, as does cargo clippy --workspace --all-targets -- -D warnings.

Not split in two. In the issue thread I proposed one PR per rule. The second rule turned out to be a rewrite of how a side is transposed, and the mask is about ten lines inside it, so landing the first rule on code the second replaces seemed worse than one PR.

The third item in #153 (a region read by several reductions) is not needed: with these two rules, attention has no such region left, as the issue expected.

Design §4.3's paragraph on contractions and decision S19 describe the extension.

🤖 Generated with Claude Code

https://claude.ai/code/session_01CvszJhH8pn8H9G69wEPMU2

Summary by CodeRabbit

• New Features
  • Gradient calculations can now use optimized contraction for saved sums over more complex inner-join structures, including multi-input sides and causal attention.
  • Grouping-key-only join conditions and masks are supported in these cases. Unsupported query shapes continue to use the general gradient rule.
• Documentation
  • Updated design documentation to describe the expanded contraction optimization and its fallback behavior.