opus_r22_gr — gr_down_f two rows a simdgroup + direct float4 dot; gr_up_f lo staging (rejected) #
patch: /tmp/opus_r22_gr/opus_r22_gr.patch (git -C vendor/strata apply --check clean against the current working tree,
incl. the uncommitted fuse_parity.cpp). scratch tree: /tmp/opus_r22_gr/strata (built: fuse_parity, gr_parity;
window.cpp compiled standalone with build-metal's flags → OK). baseline kernels: /tmp/opus_r22_gr/base_kernels.
device: Apple M5 (24 GB Air). :8080 not touched, no full-model runs. files touched: gr.metal, window.cpp
(Window::gr_read only), fuse_parity.cpp.
what changed #
gr.metal
grd_acc<F, NR>: NR consecutive rows per simdgroup. both rows' units are decoded, then each xn float4 is read once (straight from device memory, nofloat xs[UNIT]copy) and FMA'd into both rows' 4 partial sums. per-row the FMA chain is exactly unit_combine's (s0/s1 over the first half, s2/s3 over the second, same order), so the result is bitwise the old kernel for any SG / NR with the same KS.gr_down_f<SG, KS, NR, F, FI>: row0 = (tg.x * RS + rs) * NR,red[KS*RS*NR*8], the final reduce runs on lanes 0..NR*8 of ks 0. NR must divide hc_lr and hc (320, 4) so a simdgroup's rows are all w_down or all w_inject.- new entry points:
gr_down{,_q8,_q8f}_4_4_r2,gr_down{,_q8,_q8f}_8_8_r2. all old names stay (NR = 1), so token_chain / prompt_chain / gr_parity / window.cpp:2518 (mtp path) are untouched and still valid. gr_up_f/gr_dot(xs + unit_combine): unchanged,gr_dotis now gr_up's only.
window.cpp gr_read: M5 → *_4_4 (128 threads, 324 threadgroups; bitwise = the old 8_4, same KS = 4) at
ntok ≤ 4, *_4_4_r2 (128 threads, 162 threadgroups) at ntok ≥ 5. non-M5 → *_8_8 at ≤ 4, *_8_8_r2 at ≥ 5 (also
bitwise = old 8_8). q4/q5 unchanged. static_assert(LR % 2 == 0 && HC % 2 == 0).
fuse_parity.cpp --gr-unroll: a Down entry now has a baseline name and a variant name with their own rows per
threadgroup and threadgroup sizes, so a variant entry point is checked byte-for-byte against the kernel it replaces.
14 gr_down checks (q8/q8f/bf16 × 4_4, 4_4_r2 vs 8_4; × 8_8, 8_8_r2 vs 8_8; q4_8_8, q5_8_8 — new Q4_0 / Q5_0 weights).
75 checks per ntok, ~3 s with bench.
parity #
--gr-unrollntok 1..8: 75/75 bitwise PASS at every ntok (memcmp == 0).- non-vacuity: mutating one FMA in grd_acc (
* 1.0000001f) → all 14 gr_down checks FAIL. gr_parityALL PASS for: default (bf16 8_8), GR_Q8, GR_Q8+GR_INJ_F32, and each of those with GR_DOWN=4_4.- default
fuse_parityALL PASS (that one runs 5.6 s; it's the pre-existing gemv_cat/router/comb suite, not mine).
A/B (µs per dispatch; median of 24 alternating base/variant pairs × 64 reps; Δ = median per-pair ratio; wins) #
| kernel | ntok 1 | ntok 2 | ntok 3 | ntok 4 | ntok 5 | ntok 6 | ntok 7 | ntok 8 |
|---|---|---|---|---|---|---|---|---|
| gr_down_q8_4_4 | 23.6→23.1 -1.9% (22/24) | 24.4→23.8 -2.3% (24/24) | 32.6→31.5 -3.2% (24/24) | 44.2→40.6 -8.1% (24/24) | 59.4→55.6 -6.5% (24/24) | 78.9→75.2 -4.5% (17/24) | 96.8→97.3 +0.6% (1/24) | 117.8→119.7 +1.6% (0/24) |
| gr_down_q8_4_4_r2 | 23.4→23.3 -0.5% (20/24) | 24.4→23.5 -3.8% (24/24) | 32.6→32.9 +0.8% (1/24) | 44.3→41.3 -6.7% (24/24) | 59.4→49.5 -16.8% (24/24) | 78.8→58.4 -26.1% (24/24) | 96.7→79.5 -17.8% (24/24) | 117.6→89.5 -23.8% (24/24) |
| gr_down_q8f_4_4_r2 | 24.4→24.3 -0.4% (19/24) | 25.4→24.6 -3.1% (24/24) | 33.3→33.4 +0.5% (3/24) | 44.9→41.2 -8.3% (24/24) | 60.1→49.6 -17.7% (24/24) | 79.1→58.6 -26.3% (24/24) | 98.7→80.5 -18.4% (24/24) | 119.8→90.4 -24.6% (24/24) |
| gr_down_q8_8_8 | 23.7→23.4 -1.5% (16/24) | 25.9→24.9 -4.0% (24/24) | 35.5→34.8 -2.0% (24/24) | 46.6→45.2 -3.1% (24/24) | 65.8→62.4 -4.9% (24/24) | 91.3→83.5 -8.6% (22/24) | 114.5→114.7 +0.3% (5/24) | 139.0→140.0 +0.7% (0/24) |
| gr_down_q8_8_8_r2 | 23.5→23.8 +1.3% (3/24) | 26.0→25.7 -1.0% (24/24) | 35.5→34.7 -2.3% (24/24) | 46.6→46.6 -0.2% (16/24) | 65.6→58.0 -11.7% (24/24) | 91.2→70.3 -23.0% (24/24) | 114.4→99.8 -12.8% (24/24) | 139.0→113.9 -18.0% (24/24) |
| gr_up_q8 | 23.7→23.6 -0.1% (15/24) | 24.5→24.4 +0.0% (12/24) | 26.2→26.1 +0.0% (10/24) | 29.2→29.2 +0.0% (12/24) | 32.8→32.9 +0.1% (11/24) | 37.1→37.1 +0.1% (10/24) | 42.0→42.0 -0.0% (13/24) | 47.2→47.2 +0.0% (11/24) |
baseline is always the old kernel the variant replaces (q8 rows: gr_down_q8_8_4; q8_8_8* rows: old q8_8_8).
shipped on M5 (q8, the production path): ntok 1..4 → q8_4_4: -1.9 / -2.3 / -3.2 / -8.1 %; ntok 5..8 →
q8_4_4_r2: -16.8 / -26.1 / -17.8 / -23.8 % (all 24/24 except ntok 1 at 22/24). q8f (Unsloth inject F32) same shape.
what was tried and not shipped #
- NR = 2 at small ntok:
4_4_r2vs4_4head-to-head: ntok 1 +0.9%, 2 -2.0%, 3 +3.5%, 4 +1.3%, 5 -11%, 6 -22%. so the crossover is at 5; ntok 2 is a tiny r2 win on q8 but a loss on bf16/8_8, kept the one simple rule (> 4). - NR = 4 (
8_4_r4,4_4_r4): +5..47% everywhere (q[4][32] + acc[32] → register spill / occupancy). - SG = 8 / 16 with NR = 2 (
8_4_r2,16_4_r2): worse than4_4_r2at ≤ 4, roughly equal past 4. - unroll threshold for NR = 2: full unroll at 7-8 (
u8): ntok 8 +19%;u4: loses ntok 5-6. kept ≤ 6. - NR = 1 direct float4 at 7-8: +0.6..1.6% (0/24) — doesn't matter, r2 takes those.
- gr_up direct float4 dot: +2..12% at ntok ≥ 4 → reverted (gr_up keeps xs + unit_combine).
- gr_up lo staged in threadgroup memory (
los[320*8], float4 copy + barrier): +1..5.5% at every ntok, 0-2/16 wins → rejected. lo is 1.3-10 KB and every simdgroup reads the same bytes, so it is already an L1 hit; the barrier only costs. - KS changes (e.g. KS = 8 on M5) were not explored for shipping: they change the K-split sum order → not bitwise.
notes #
- why NR = 2 isn't a 2x at ntok 8: 89 µs at ntok 8 vs 41 at ntok 4 is still ~2x per token-doubling; halving xn reads took ~25%, the rest is FMA issue + the rolled token loop. next lever is probably the per-token loop structure (acc in registers past 6 tokens without spill) rather than more row reuse.
- end-to-end not measured (no full-model runs allowed). gr_down runs ~97×/window at the window's ntok (+96 in fold,
+2 mtp_prefill — the mtp path at window.cpp:2518 still uses
q8_8_8NR=1, outside my lines; switching it to the same rule would pick up the same win). - the non-M5
8_8_r2choice was measured only on M5 hardware running that kernel; unverified on M1-M4.