opus_r21_gr — gr.metal / gr2.metal GR_NTOK loop unrolling + program-scope plane tables #

patch: /tmp/opus_r21_gr/opus_r21_gr.patch (git -C vendor/strata apply --check clean, against the current working tree, incl. the uncommitted fuse_parity.cpp). scratch tree: /tmp/opus_r21_gr/strata (built: strata_window_cli, fuse_parity, gr_parity, gr2_parity, comb_fold_parity). baseline kernels for A/B: /tmp/opus_r21_gr/base_kernels. device: Apple M5 (24 GB Air). the live :8080 server was not touched; no full-model runs.

what changed #

gr.metal

• gr_q2s[8] / gr_p4s[4] are program-scope constant tables now (they were function-local arrays in gr_up_f and gr_up3); gr2.metal uses the same tables (it is always compiled after gr.metal, it already needs GR_NTOK).
• grd_acc: the token loop body is factored into gr_dot<UNIT> (load16 + unit_combine), shared with gr_up_f. the token loop is unrolled only when GR_NTOK <= 6; for 7 and 8 it stays rolled. reason: with full unroll, gr_down at ntok 6..8 measured consistently 1-2% slower (ntok 8: 0/24 pair wins for 8_8 and 8_4); ntok 2..5 gained 5-14%. GR_NTOK is a function constant, so the branch folds at pipeline specialization (no runtime cost). i also tried 4-token chunks with a simdgroup_barrier between them; no help at ntok 8.
• gr_down_f: unroll(full) on the acc init, the per-token simd_sum loop, the KS reduction.
• gr_up_f: unroll(full) on the acc init, the token loop (full unroll at all ntok here; it wins everywhere, up to -18%), the gate loop, the k<16 shuffle loop; the stream mix is if (hc == 4) unrolled-4 else the old loop (same fma order).
• gr_norm: unroll(full) on the 8-partial sum.

gr2.metal (gr_down3, gr_down3b, gr_up3)

• unroll(full) on the RPS, TT, UNIT/16-element copy, k<8, i<5, e<8 and k<16 shuffle loops.
• the t < T loops over acc[RPS][TT] became t < TT with if (t >= T) break;, so acc[k][t] gets a constant index (T = min(TT, GR_NTOK - t0) is not a compile-time constant). same additions in the same order, so bitwise.
• gr_up3: same hc == 4 stream-mix split and gr_q2s/gr_p4s as gr_up_f.

fuse_parity.cpp: new mode fuse_parity --gr-unroll BASE_DIR NTOK [bench] — compiles quant/units/gr/gr2 twice (BASE_DIR copies vs the tree's kernels) and checks the same entry points byte for byte on the same inputs: gr_norm (apply 0/1), gr_norm_comb, gr_down_{q8_8_4, q8f_8_4, q8_8_8, q8f_8_8, 8_4, 8_8}, gr_up_q8, gr_up, all 18 gr_down3/gr_down3b instantiations (apply 0/1, comb 0/1), gr_up3_{30,8} — 67 checks per ntok. bench adds 8 alternating pairs × 64 reps for gr_down_q8_8_4 / q8f_8_4 / q8_8_8 and gr_up_q8. ~1-1.5 s per invocation.

parity #

• --gr-unroll ntok 1..8: 67/67 bitwise PASS at every ntok (memcmp == 0).
• non-vacuity: against a mutated baseline (eps + 1e-3, sigmoid 1.001) every one of the 67 checks FAILs.
• default fuse_parity (gemv_cat, router_route, gr_norm_comb): ALL PASS. gr_parity (CPU fp64 ref, planes == xprep bitwise) ALL PASS, gr2_parity ALL PASS, comb_fold_parity ALL PASS (that one runs 9.5 s, T=1024 prompt-size — sorry, over the 5 s budget; nothing else was).

A/B timings (µs per dispatch, median over 3 runs × 8 pairs = 24 alternating pairs; ratio = median of per-pair ratios) #

notes:

• ntok 1 is flat (as expected: a 1-trip loop needs no unroll). wins are at ntok 2..5 for gr_down (-4..-14%) and at every ntok ≥ 2 for gr_up (-6..-18%). the 194 decode-window dispatches are gr_down + gr_up at the window's ntok, so at ntok 2-4 on the M5 8_4 path that is ≈3.4 µs (ntok 2) to ≈7.5 µs (ntok 4) saved per down+up pair, ×97 ≈ 0.3-0.7 ms per window in isolated-kernel terms (end-to-end not measured: no full-model runs allowed).
• the pairs ran while the live server was up, so there is some contention noise; medians/win counts are stable.
• side observation (not in scope): gr_down scales badly with ntok (23 µs at 1 → 117-139 µs at 8, ~5x for 8x tokens on a weight-bound kernel). at ntok 8 every row re-reads all 8 tokens' xn (327 KB) per lane stride; staging xn in threadgroup memory or a token-tiled layout looks like the next lever for mtp_prefill / fold.