opus_r34_gr — Round 34 Track A: multi-row / coalesced gr_down (ntok 2..8) + Window::gr_read dispatch #
Patch: /tmp/opus_r34_gr/opus_r34_gr.patch (git -C vendor/strata apply --check OK; touches metal/kernels/gr.metal,
metal/src/engine/window.cpp (Window::gr_read dispatch line only; gr_read_m unchanged) and metal/tests/fuse_parity.cpp).
Scratch tree /tmp/opus_r34_gr/strata, baseline kernels /tmp/opus_r34_gr/base_kernels, full logs
/tmp/opus_r34_gr/final_bench.txt (+ final_bench_run1.txt), every experimental variant kept in
/tmp/opus_r34_gr/experiments_gr.metal / experiments_fuse_parity.cpp. Device: Apple M5 10-core GPU, live server sharing the GPU
(a few late-run samples were trashed by it; medians of 16 paired ratios used throughout).
tl;dr #
The big-idea candidates (j-outer token loop, NR = 4, larger RS, threadgroup-staged coalesced xn) are all bitwise PASS but
slower at every ntok. The only win is small: keep the packed Q8_0 NR = 2 token loop rolled at every ntok for KS = 4, and
let gr_read switch to _r2 from ntok 4 instead of 5.
Shipped (bitwise PASS ntok 1..8 in --gr-unroll, default fuse_parity ALL PASS) #
gr.metal:grd_acc<F, NR, ROLL>;gr_down_fsetsROLL = F == Q8_0 && NR == 2 && KS == 4(the 4_4_r2 / q8f_4_4_r2 down rows). Same fmas, same order, just nounroll(full)over tokens. KS = 8 (8_8_r2, the MTP head) measured +2..6% slower rolled, so it keeps the unrolled loop (now ±0.2%).window.cppWindow::gr_read:r2 = ... ntok > (m5_gr ? 3 : 4).4_4_r2is bitwise4_4(same KS), so the switch is free for parity.gr_read_m(KS = 8) stays at> 4: 8_8 beats 8_8_r2 at ntok 4 (45.5 vs 46.7 µs).fuse_parity --gr-unroll: benches for the dispatch choice (4_4vs4_4_r2,q8f_4_4vsq8f_4_4_r2,8_8vs8_8_r2).
Final paired timings (final_bench.txt, median of 16 pair ratios) #
| kernel (base → new) | ntok | base µs | new µs | median Δ |
|---|---|---|---|---|
| 4_4_r2 → 4_4_r2 (rolled) | 4 | 41.34 | 39.55 | -4.4% |
| 4_4_r2 → 4_4_r2 | 5 | 49.50 | 48.18 | -2.7% |
| 4_4_r2 → 4_4_r2 | 6 | 58.43 | 57.44 | -1.7% |
| 4_4_r2 → 4_4_r2 | 7 / 8 | 68.14 / 79.77 | 68.14 / 79.79 | 0.0% (already rolled) |
| q8f_4_4_r2 → same | 4 / 5 / 6 | 41.24 / 49.51 / 58.96 | 39.96 / 48.48 / 57.69 | -3.0 / -2.1 / -1.9% |
| 8_8_r2, 4_4, q8f_4_4, 8_8 | 1..8 | ±0.2% (unchanged code) | ||
| engine ntok 4: 4_4 → 4_4_r2 (rolled) | 4 | 40.64 | 39.63 | -2.5% (run 1: -0.8%) |
| engine ntok 3: 4_4 → 4_4_r2 | 3 | 31.67 | 33.03 | +4.3% → stays 4_4 |
| engine ntok 2: 4_4 → 4_4_r2 | 2 | 23.84 | 23.82 | -0.3% (noise) → stays 4_4 |
| engine ntok 5..8: old r2 → rolled r2 | 5/6 | 49.50/58.43 | 48.18/57.44 | -2.7/-1.7% |
Net at the 3-draft verify (ntok 4, 97 calls): ~0.3-1.0 µs/call → ~0.03-0.1 ms per verify step. Small, but free and bitwise.
Measured and rejected (all vs base gr_down_q8_4_4 NR = 1 at the same ntok, median Δ; all bitwise PASS unless noted) #
| variant | 2 | 3 | 4 | 5 | 6 | 7 | 8 |
|---|---|---|---|---|---|---|---|
| 4_4_r2 (old, unrolled ≤6) | -0.9 | +4.4 | +1.8 | -11.0 | -22.4 | -30.0 | -33.5 |
| 4_4_j (NR 1, j-outer) | +1.5 | +0.5 | +25.9 | +5.5 | +6.9 | +5.2 | +8.9 |
| 4_4_r2j (NR 2, j-outer, task idea #1) | +10.5 | +8.5 | +13.4 | +56.0 | +78.7 | +32.6 | +52.8 |
| 4_4_r4 (NR 4 packed, 81 TGs) | +32.6 | +21.2 | +17.5 | +28.6 | +28.7 | +9.9 | +0.0 |
| 4_4_r4j | +16.6 | +20.3 | +42.5 | +40.9 | +49.8 | +34.4 | +37.8 |
| 8_4 (RS 2) | +2.1 | +2.3 | +2.2 | -0.3 | -3.7 | -0.1 | -0.6 |
| 8_4_r2 | +2.4 | +8.6 | +5.7 | -7.9 | -21.7 | -30.1 | -34.2 |
| 8_4_r4j | +22.6 | +26.0 | +50.0 | +46.6 | +48.0 | +38.2 | +36.0 |
| 16_4 (RS 4) | +8.1 | +5.8 | +3.8 | -1.0 | -10.7 | -12.9 | -15.6 |
| 16_4_r2 | +20.4 | +21.0 | 0.0 | -15.8 | -30.3 | ||
| 32_4 (RS 8) | +17.4 | +13.3 | +7.7 | -8.3 | |||
4_4_u2 (unroll_count(2) tokens) |
+9.5 | +9.7 | +0.2 | ||||
| 4_4_r2u2 | +16.2 | +27.2 | -14.1 | ||||
| 4_4_st (per-simdgroup threadgroup-staged, coalesced xn) | +96.3 | +49.9 | +32.4 | ||||
| 4_4_r2st (same, NR 2) | +30.5 | -0.9 | -9.0 | ||||
| probe 4_4_pc (coalesced xn, WRONG math) | -1.9 | -15.1 | -0.1 | -11.2 | -22.0 | -29.9 | -22.9 |
| probe 4_4_r2pc (coalesced, WRONG math) | +4.3 | -2.5 | -7.5 | -14.4 | -23.9 | -34.3 | -36.2 |
Notes:
- j-outer (task idea #1) is bitwise identical but 5-80% slower. Holding
s[NR][ntok][4]live across the j loop costs more registers than the 2-float4-per-row it saves; the CSE-hoist theory doesn't pay off in practice (the compiler already widens the packedchar4per use well enough — rolling the token loop is what helps). - NR = 4 / more RS: fewer, fatter simdgroups lose at ntok ≤ 6 (latency/occupancy: 81 TGs × 4 SGs isn't enough in flight).
Bigger threadgroups (16_4, 32_4) only help where NR = 1 is already much worse than
_r2. - Coalesced xn: the wrong-math probe bounds the gain: ≤ 0% at ntok 4 for NR 1 (confirms r33), -7.5% for NR 2 at ntok 4,
-2..-3% vs
_r2at 7-8. A real transpose through threadgroup memory (8 coalesced loads → padded 9-float4/unit tile →simdgroup_barrier→ 8 tile reads per token) costs far more than that (+30..+96%), so staging is a dead end here. Cross-simdgroup cooperative staging needs RS > 1, which is itself slower (above), plus two threadgroup barriers per token.
Where ntok 4 is stuck #
gr_down_q8_4_4 at ntok 4 is ~40 µs vs a ~23 µs DRAM floor. Coalescing doesn't move it (probe), halving xn traffic (NR 2) moves
it ~2.5% only when rolled, and more/fatter simdgroups lose: it looks latency-bound per simdgroup (3 serial u-iterations of
DRAM weight load → L2 xn loads, with lanes 16-31 idle on the third). The bitwise constraint pins the unit→lane map (each lane's
acc order and the simd_sum tree), so the remaining lever is a structurally different grid (e.g. split-K across
threadgroups with an ordered second-pass reduce) — out of this file/line scope.