opus r14 track A: mm_gu3 / mm_down3 epilogue unrolling + W8 lane-constant hoisting #
the short version: only one change survives. it's a half4 HS store in the mm_gu3 epilogue, gated to F == T_IQ2_XS, and it's
0.44-0.63% faster on the production IQ2_XS gate/up. every other idea in the brief was either noise or a regression (´・ω・`)
- patch:
/tmp/opus_r14_moemm/opus_r14_moemm.patch(touches only metal/kernels/moe_mm.metal, +11 lines, applies withgit -C vendor/strata apply --check) - scratch copy:
/tmp/opus_r14_moemm/strata. every variant tried sits in/tmp/opus_r14_moemm/v_*/
method #
- bitwise check: the existing
build-metal/moe_mm_paritybinary withSTRATA_METAL_KERNELS=<variant dir>, run--v3 --dump(HS + y). that covers all 11 mm_gu3 formats x tiles {16,32,32,4 (prod), 16,64,32,2} with dn Q2_0, plus all 4 mm_down3 formats x dtiles {64,32,2 (prod), 32,32,4} x gu {Q2_0, IQ2_XS}, each at T=256 and T=512: 72 dumps, memcmp against the baseline. the baseline dumped twice comes out identical, so the dumps are deterministic. - timing: the existing
--ab <dir>paired, interleaved mode (5 alternating samples, medians), production tiles, run in both orders (base as A, then variant as A). the number reported is the median of the B-vs-A % diffs. - caveat: the live :8080 server is resident, so each run took ~11 s of wall time (allocation/paging). the timed kernels themselves are short. one cell, Q6_K at T=1024, is noisy and biased even base-vs-base (+2.3% / +8.1%), so I don't trust any Q6_K@1024 number.
variants (all 72/72 bitwise identical) #
| variant | change | screen result |
|---|---|---|
| A | unroll(full) on the accg/accu init, the SiLU epilogue, and the mm_down3 acc init |
noise everywhere (IQ2_XS -0.08/+0.03, IQ3_S +0.19/-0.18, Q2_0 dn +0.11/-0.04) -> dropped |
| B | A + per-lane scalar HS store (m loop, f = lane) | IQ2_XS -0.36%, IQ3_S +3.1..3.4% at T=512 -> dropped |
| C | A + half4 HS store (float4 from ot, 4 rows a pass) | IQ2_XS -0.37..-0.60%, IQ3_S +3..5% at T=384/512/768 (-0.4% at 256/1024), Q6_K@512 +1.1% -> regresses |
| C2/C3 | C with the store loop at unroll(disable) / unroll_count(2) | IQ3_S still +2.2..3.0% -> so it isn't the unroll; any epilogue change moves IQ3_S |
| D | lane-constant j_in: kl = ix0[0] & 15, `j_in = (k0 & 224) |
kl(QK256) /(k0 & 32) |
| Q | a real W8<T_Q2_0> (one scale load, q0 = bl[2+o], q1 = bl[6+o]) |
Q2_0 dn +0.06/+0.11, Q2_0 gu -0.03/+0.08 -> dropped (matches the old comment in the file saying the shared-scale W8 lost) |
| E (kept) | C's half4 store, gated to F == T_IQ2_XS, without A |
see below |
I didn't build the explicit 256-block outer loop. D already shows that giving the compiler a lane-constant (ib, l, hh) changes nothing, so the decode address math isn't the bottleneck.
E: final paired A/B (T = 512 / 1024 / 2048, one run each order, base ~us, median diff, per-run diffs) #
| kernel | fmt | T | base us | diff | runs |
|---|---|---|---|---|---|
| gu | IQ2_XS | 512 | 8080 | -0.63% | [-0.72, -0.54] |
| gu | IQ2_XS | 1024 | 13435 | -0.53% | [-0.55, -0.5] |
| gu | IQ2_XS | 2048 | 22962 | -0.44% | [-0.39, -0.49] |
| gu | Q6_K | 512 | 23204 | -0.18% | [0.24, -0.61] |
| gu | Q6_K | 1024 | 31518 | +1.75% | [2.01, 1.49] |
| gu | Q6_K | 2048 | 36624 | -0.08% | [-0.1, -0.06] |
| gu | IQ3_S | 512 | 18868 | -0.13% | [-0.03, -0.23] |
| gu | IQ3_S | 1024 | 18282 | +0.01% | [0.19, -0.18] |
| gu | IQ3_S | 2048 | 30964 | +0.02% | [0.22, -0.19] |
| dn | Q2_0 | 512 | 3482 | +0.05% | [0.14, -0.03] |
| dn | Q2_0 | 1024 | 4949 | +0.03% | [-0.16, 0.21] |
| dn | Q2_0 | 2048 | 8097 | -0.03% | [0.01, -0.06] |
The earlier screen of E at T=512 gave IQ2_XS -1.03/-0.51 and IQ3_S +0.11/+0.40. IQ3_S, Q6_K and Q2_0 compile to the original epilogue: the branch is a compile-time constant. Their diffs are all within the noise (Q6_K@1024 is the biased cell noted above).
correctness #
- E: 72/72 dumps memcmp == 0 against the baseline.
- default parity PASS (prod tiles 16,32,32,4 / 64,32,2): IQ2_XS rel L2 1.79e-05, IQ3_S 1.66e-05.
- the half4 store needs HS + rown_ff + n0 + f to be 8-byte aligned. n_ff is a multiple of TNSG, n0 and f are multiples of 4, and window.cpp binds p.HS as a whole buffer, so that holds.
take-away for later rounds #
mm_gu3's epilogue is sensitive to codegen per format: the same store rewrite is -0.5% on IQ2_XS and +3..5% on IQ3_S at mid T. The decode-side micro-hoisting (D, Q) is already done by the compiler. More gate/up speed has to come from the decode itself (table lookups, bytes loaded), not from address arithmetic or the epilogue.