sdpa bwd: lead square A100 backward window graphs from a measured study - #2352

Problem

sdpa_bwd_sm80 is a trail-only row, so on an A100 a dense backward under a left window runs on the backend. For square graphs that is the slow choice: on a fresh holdout the default ran at a geomean of 3.46x (bf16) and 2.79x (f16) over the faster of FROST and backend, and up to 35x.

Change

• placement_table.py: 16 generated always entries for sdpa_bwd_sm80 on cc 8.0 only. They cover bf16 and f16, d64 and d128, dense bottom-right causal with a left window, MHA and GQA, with S on and off the 128-key tile. They come from gen.py study.json --margin 0.05 --emit on train, run from each dtype's study directory. Quotes are the only edit to the emitted lines.
• bwd/heuristics.py: SQUARE_ONLY_LEAD_ROWS = {"sdpa_bwd_sm80"}. In _measured_lead, a row in that set lifts only when s_q == s_kv.
  • This is a hand rule under Rule S12, for a feature the class key excludes. class_key does not see squareness, and an earlier mixed square / non-square study on the A100 measured q < kv graphs in these same classes up to 1.45x slower on FROST (b12 40/8 d64, q 128, kv 8960, window 8000: 7.0 vs 4.8 ms). That study failed its gate and is diagnostic only; it is why this one measures square graphs alone.
  • I did not add squareness to class_key, because that would rekey every row's entries.
• The new test checks that the square graphs of these classes lead, and that q < kv graphs of the same classes (for which lookup still answers True), padded graphs, graphs without a window and cc 8.6 keep trailing.

Measurement

• Hardware and software: A100-SXM4-40GB, cuDNN 9.27.0, torch 2.14.1+cu130. Clocks were not locked, and no other GPU processes were running.
• Tree: develop 6903d1a + this commit (python tree sha a74720f5c8d0c147 in both runs).
• Train: 128 square cases per dtype, b 1-32, S 128-16250, windows 16-2048. Train regret at the default was 2.854 (bf16) and 2.921 (f16); with the fit, 1.000 and 1.001 (max 1.070, 2 of 128 f16 cases above 1.05).
• Holdout: 128 square cases with other heads, batches, lengths and windows (b 1-24, S 128-16484, windows 24-8000), plus 32 q < kv cases in the same classes for routing. No case appears in the earlier mixed study. Grids and gate were frozen before anything was timed, and no entry changed after the holdout.
• Timing: three repeats per square case, each with a fresh graph and tensors. Four legs are timed independently and interleaved: first backend plan, first FROST plan, first plan on this tree (auto), and first plan with the case's entry removed (default). Each leg keeps its median.

Regret is a leg's time over the faster of FROST and backend, geomean (max).

• Every timed holdout case got faster. The slowest is 0.975x of the default (f16 b24 28/4 d128 S 128, window 8000: 0.418 vs 0.429 ms).
• Largest time saved: 1671.9 → 51.2 ms, bf16 b12 48/6 d128 S 16484, window 96.
• Noise: repeat spread has median 0.007, p90 0.011, max 0.026. The default and backend legs were the same plan on every case, and their ratio ranges from 0.996 to 1.036.
• Routing: in every repeat, auto was the FROST row on all 128 square cases. All 32 q < kv cases match an entry (lookup is True) and ran the same plan as the default. There were no errors.

How the holdout was run

The first run trained, fitted and started the holdout, then hit its rental deadline after 167 of 416 holdout lines. Those lines were never scored. A second run on another A100-SXM4-40GB host ran the whole holdout from scratch on the same frozen case file, with the same fit. Its driver is the first one with the train loop replaced by checksums of the first run's generator output and a check that the applied tree reproduces the same python sha. The frozen-file check passed at its start. Both drivers, their diff and the unscored partial file are in the records.

Not measured

• Cold cache, and other cc 8.0 parts (A100 PCIe, A30). The key is cc 8.0, so the entries apply to every cc 8.0 part, but only the A100-SXM4-40GB was timed.
• q > kv graphs: the backend rejects them under a bottom-right mask, and the guard keeps them trailing anyway.
• Padded and trimmed graphs keep trailing through _measured_lead, so they were not timed.

Testing

• test_placement_table.py, test_sdpa_fwd_placement.py, test_sdpa_fwd_heuristics.py and test_sdpa_bwd_heuristics.py pass on the A100 at this commit (600 passed, in both runs), and black is clean on the three changed files.
• This branch merges cleanly with current develop (819bb56) and with perf(sdpa placement): lead the SM100 row on B300 dense d192 and windowed d64 prefill #2346. On the develop merge, run on a host with no GPU, test_placement_table.py and the two _measured_lead tests pass, and test_sdpa_fwd_placement.py / test_sdpa_fwd_heuristics.py fail exactly where develop itself fails there (no DSL or GPU). The rest of test_sdpa_bwd_heuristics.py needs the DSL and was not run on that merge.
• It has a text conflict with perf(sdpa placement): lead the SM120 row on GB10 dense d64 and d128 windowed decode #2345 and perf(sdpa placement): lead the SM120 backward on RTX 5090 dense d64 and d128 #2349 at the end of TABLE in placement_table.py, because all three append there. I have not run the placement tests on a merge with either.

Records (machine-readable results, gate, frozen cases, train data, generator output and drivers): https://gist.github.com/0z5a/c19c233368d38573122414fa24bb5b3f