5.1 kB 51 lines
eb15e062
Round 31 blitz: mm_gu3 IQ2_S/IQ2_XXS/IQ2_XS/IQ3_S tab+hh & shift (+6.8..+8.2% / -57.9 ms in kprof), moe_down/gemv Q2_0R 0x6400 half-mantissa decode (-7.6..-17.4%), dq_f16 IQ4 tab init & perm=4 inverse contiguous write
dawn
⋯
audits/opus_r31_moemm.md
sourcerendered
download raw history
at main Markdown

opus R31 Track A: mm_gu3 IQ2_S / IQ2_XXS / IQ2_XS / IQ3_S decode arithmetic #

Patch: /tmp/opus_r31_moemm/opus_r31_moemm.patch (only metal/kernels/moe_mm.metal; git -C vendor/strata apply --check passes). Scratch tree: /tmp/opus_r31_moemm/strata, build /tmp/opus_r31_moemm/build, baseline kernels /tmp/opus_r31_moemm/base_kernels. Exhaustive checker: /tmp/opus_r31_moemm/chk.{metal,cpp} (binary /tmp/opus_r31_moemm/chk).

What ships (all integer and pointer rewrites, no FP change) #

  1. W8<T_IQ2_S>::block: the lane's hh = (kl >> 2) & 1 comes from mm_gu3 (fixed for the whole K loop). The grid word is read as (tab + hh)[2 * i], and the sign byte as s >> (hh << 2) instead of hh ? s >> 4 : s. Measured separately: the shift alone gives about +4.7%, and tab + hh adds about +3.5% on top.
  2. W8<T_IQ2_XXS>: there's a new block(bl, k, d, tab, hh), and mm_gu3 runs IQ2_XXS with the 256-wide super-block loop (cached dg/du) the way IQ2_S does, with the lane hh and tab + hh. The sign4_iq2xs select stays. run() is a thin wrapper, so the other callers don't change.
  3. W8<T_IQ2_XS>: same change as IQ2_XXS (block with cached d, lane hh, tab + hh).
  4. W8<T_IQ3_S>::run: s >> (4 * hh) instead of the select.
  5. mm_gu3: constexpr bool sb = IQ2_S || IQ2_XXS || IQ2_XS picks the 256 super-block loop. The QUANT_CTAB path is kept inline (ctab ? grid[i] >> 32*hh : th[2*i]).

Results: moe_mm_parity --v3 --tile 16,32,32,4 --dtile 16,64,32,2 --dn Q2_0 --ab base --gu-bench #

diff = (base - new) / new, so a positive number means new is faster. Medians are from paired, interleaved runs. The machine was shared with the live server, so I give the quiet-run ranges.

gate/up T=512 T=1024 T=2048 bitwise
IQ2_S +6.6..+11.7% +7.9..+8.9% +7.9..+8.8% (25.1 ms vs 27.1 ms) PASS
IQ2_XXS +3.0..+5.9% +4.4..+5.7% +5.8..+7.6% PASS
IQ3_S +4.6..+5.4% +2.4..+2.6% +2.0..+3.0% PASS
IQ2_XS +3.5..+3.8% +3.5..+5.4% +6.6..+6.9% PASS

All of them hit HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS. I also checked 72 extra A/B parity runs: 4 formats × tiles {16,32,32,4 / 16,64,32,2 / 32,64,32,2 / 16,16,32,8 / 16,32,16,4 / 16,64,64,2} × {T=256, T=1000, T=512 with QUANT_CTAB=1}. All passed bitwise.

Estimated full-model kprof savings, extrapolated from the T=2048 standalone numbers (I didn't run a full-model check, per the rules): IQ2_S 763.8 ms × ~7.4% ≈ -56 ms, IQ2_XXS 220.1 × ~6% ≈ -13 ms, IQ3_S 46.7 × ~2.5% ≈ -1 ms. That's roughly -70 ms total (about 2% of prefill GPU time).

Exhaustive GPU checks (chk): 65,536 fp16 d × 16 nibbles × v in 1..63 × all 16 sign masks (1.06 G evaluations) #

  • (d*0.125f) * float(2n+1) vs d * fma(n, 0.25f, 0.125f): bit-identical for every d. IQ3_XXS (d*0.25f)*float(2n+1) vs d*fma(n,.5,.25) is also identical. This is expected: scaling by a power of 2 is exact, the product v*sc (≤ 6+5+11 bits) is exact in fp32, so there's a single rounding to fp16 either way.
  • These w4_signed rewrites are bit-identical to the original for every d, including NaN/Inf: the fp32 sign-XOR (as_type<uint4>(float4(uchar4)) ^ (s << {31..28})), and the fp16 0x6400 mantissa trick with sign XOR on the fp16 codes.
  • Sign-XOR on the fp16 output is identical for every finite d but differs on NaN-d sign bits, so it is not 100% bit-exact.

Tried and rejected (all bitwise PASS, all slower; measured on IQ2_S T=2048 unless noted) #

variant result
w4_signed fp32 sign-XOR -11.3% (IQ2_XXS -16%)
w4_signed fp16 0x6400 codes + sign-XOR -12% (IQ2_XXS -15.5%)
w4_signed 0x6400 codes, keep select -11.2%
w4_signed select on the scale (s&1 ? -sc : sc) -3.4% (IQ2_XXS ±0)
scale as fma(float(n), d/4, d/8), d/4 and d/8 cached -9.5% (d/8 computed at use: -2.6%; with tab+hh: -13%)
pre-scaled dg = d*0.125, `sc = dg * float(((scb<<1)&30) 1)`
tab + hh with the select kept for the sign -2.5%
index pre-doubled `(p0<<1) ((qh<<(9-2l))&0x600)`
l taken from the lane as well ±0 vs shipped (dropped)
IQ2_XXS sign4_iq2xs as a shift `(j parity<<7) >> 4*hh`
IQ3_S cached d via the 256 loop + lane hh +1.5% (the shift alone is better: +2.8%; lane hh without caching: same as shift alone)

Takeaway: on the M5 the decode is ALU-issue bound (each extra op per W4 costs about 1%). d * fma(float(nibble), 0.25f, 0.125f) with constant operands, and float(uchar) plus the select-negate in w4_signed, are already the cheapest forms. The wins came from removing per-lookup integer work (+ hh, the hh ? select on the sign byte) and per-call scale loads. Selects are not uniformly bad, though: sign4_iq2xs's select beats the shift form.

Not done: IQ3_XXS (not in this model's hot list). The non-coop_k4 (M4) path still calls W4<F>::run, which is unchanged; I couldn't run it on the M5.