# 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)

| kernel | ntok | base | variant | Δ | variant wins |
|---|---|---|---|---|---|
| `gr_down_q8_8_4` | 1 | 23.42 | 23.53 | +0.1% | 10/24 |
| `gr_down_q8_8_4` | 2 | 26.39 | 24.57 | -6.9% | 24/24 |
| `gr_down_q8_8_4` | 4 | 46.85 | 44.34 | -5.4% | 24/24 |
| `gr_down_q8_8_4` | 8 | 117.19 | 117.30 | +0.1% | 11/24 |
| `gr_down_q8_8_8` | 1 | 23.73 | 23.67 | -0.2% | 13/24 |
| `gr_down_q8_8_8` | 2 | 30.07 | 25.98 | -13.4% | 24/24 |
| `gr_down_q8_8_8` | 4 | 54.68 | 46.92 | -14.2% | 24/24 |
| `gr_down_q8_8_8` | 8 | 139.25 | 139.15 | -0.0% | 14/24 |
| `gr_down_q8f_8_4` | 1 | 24.41 | 24.44 | +0.1% | 10/24 |
| `gr_down_q8f_8_4` | 2 | 26.98 | 25.52 | -5.5% | 24/24 |
| `gr_down_q8f_8_4` | 4 | 46.94 | 45.05 | -4.1% | 24/24 |
| `gr_down_q8f_8_4` | 8 | 119.59 | 119.62 | +0.0% | 9/24 |
| `gr_up_q8` | 1 | 24.08 | 23.65 | -1.6% | 24/24 |
| `gr_up_q8` | 2 | 26.02 | 24.42 | -6.0% | 24/24 |
| `gr_up_q8` | 4 | 34.22 | 29.20 | -14.7% | 24/24 |
| `gr_up_q8` | 8 | 57.36 | 47.03 | -18.0% | 24/24 |

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.