opus r23 quant: ksg parity signs, dq_f16 / DW / unroll (Round 23 Track A) #

Patch: /tmp/opus_r23_quant/opus_r23_quant.patch (git -C vendor/strata apply --check OK). Files touched: metal/kernels/quant.metal (shipped kernel change), metal/tests/dense_mm_parity.cpp (new A/B mode). Device: Apple M5 (24 GB). Scratch: /tmp/opus_r23_quant/{base,strata} (base = vendor at 296b59d + working tree).

TL;DR #

  • ksigns_iq2xs[i] == i | (popcount(i) & 1) << 7 holds for all 128 entries (checked vs ggml-common.h). But computing it is not a uniform win. On M5 it depends on format and on where the lookup would otherwise go (constant vs threadgroup), so ksg now computes the sign byte only where it was measured faster:
    • IQ2_XS with QUANT_CTAB (production decode moe_gu, since window.cpp sets ctab for the routed gate/up): -2.5..-5.5% gate/up time across NE=1..4, ntok 1..8 (ratios 0.945..0.995, 2 runs)
    • IQ3_XXS with the threadgroup table (production prompt mm_gu3_18_16_32_32_4, which runs with ctab off): -1.6..-3.9% (5 paired runs, quiet GPU; A/A control within +-0.1%). Its 128-byte sign copy is gone, and fmt<T_IQ3_XXS>::TAB drops from 288 to 256 words.
    • everywhere else it stays a lookup, because computing it measured slower: IQ2_XXS moe_gu ctab +2..+13%, IQ3_XXS moe_gu ctab +0..+6%, IQ2_XS mm_gu3 +1.5..2% (matches opus_r8/opus_r12, which rejected popcount for IQ2_XS mm).
    • known trade-off: IQ3_XXS moe_gu with ctab off (only reachable with STRATA_METAL_CTAB=0) is +5..7% slower. The default decode path (ctab on) for IQ3_XXS is unchanged (0.998..1.006).
  • dq_f16 is at the memory roofline. IQ4_XS 2560x10240 moves 67 MB (52 MB of fp16 out + 15 MB in) in ~509 us = ~132 GB/s on M5. None of the decode-side ideas helped, so none ship:
    • IQ4 DW without the table/barrier, using a 4-register packed codebook: slower, 1% on big shapes and 10-22% on small ones
    • IQ4 DW reading kvalues_iq4nl straight from constant memory: ~1% slower / neutral
    • a.perm guard (no integer division when perm == 0): 1.00x
    • 16-byte uint4 stores in place of the half4 pairs: 1.00x
  • #pragma clang loop unroll(full) on the fixed-trip loops (load16, X::load, every Dec<F>, unit_combine, every deq<F> loop, put loops): net loss on gemv. IQ4_XS 2560x10240 was +8..9% (best tile, 2 sweeps), Q6_K 2560x10240 +2.8%, Q4_K/Q3_K/Q5_K/Q4_0/BF16 within +-1..2%. Not shipped. The compiler already unrolls these where it pays; forcing it costs registers.

Bitwise verification #

Every output buffer bitwise: I added an instrumentation hook to the scratch base mtl_rt.mm (not in the patch): with STRATA_ABHASH=file, after every Context::run/run_graph it hashes every buffer bound to the dispatches. Each test binary runs twice, with STRATA_METAL_KERNELS=base/kernels and then =candidate/kernels, and the hash logs are compared. A/A runs first confirmed the tests are deterministic. With the final patch, all 31 cases gave the same hashes (OK) and the same exit codes (0):

  • dense_mm_parity (360 hashed buffer states), gemv_parity (26,316)
  • expert_parity --gu {Q2_0,IQ3_XXS,IQ2_XXS,IQ2_XS,IQ2_S,IQ3_S,IQ4_XS,Q4_K,Q5_K,Q6_K,Q8_0} --ntok 3, --dn {Q2_0,IQ4_NL,Q5_1,Q8_0} --ntok 8
  • moe_mm_parity --v3 --tile 16,32,32,4 --dtile 64,32,2 --T 48 --gu {Q2_0,IQ3_XXS,IQ2_XXS,IQ2_XS,IQ2_S,IQ3_S,IQ4_XS,Q4_K,Q8_0}, --dn {IQ4_NL,Q5_1,Q8_0}, --gu IQ2_XS --tile 16,64,64,2 (mm_gu), --gu IQ3_XXS --dn IQ4_NL
  • for the ctab=on decode path, each expert_parity --ab run (both ctab on and ctab off, IQ2_XS/IQ3_XXS/IQ2_XXS/IQ2_S) passed its in-process memcmp of plan/h/ye/y. Trick: quant.metal and units.metal are empty stubs in the kernel dir, and moe.metal / the --ab file are each version's quant+units+moe concatenation, so A and B carry different quant.metal.
  • moe_mm_parity --v3 --gu {IQ2_XS,IQ3_XXS} with QUANT_CTAB=1 (mm_gu3 ctab on): same hashes too
  • dense_mm_parity --dq-ab with DQ_AB_KERNELS: memcmp=0 for every shape tried (IQ4_XS, IQ4_NL, IQ3_S, Q8_0)

Measurements (paired, M5) #

mm_gu3 (prompt, moe_mm_parity --v3 --tile 16,32,32,4 --dtile 64,32,2 --gu-bench --ab KB, kernel dir = base, T=2048) #

diff = (cand - base) / base, med / min, %

format A/A control final candidate
IQ3_XXS -0.04/-0.07, -0.02/-0.08, +0.03/0.00 -1.73/-2.88, -2.92/-2.24, -3.66/-1.58, -3.91/-2.29, -3.79/-2.95
IQ2_XS +0.03/+0.04, +0.09/-0.03, +0.08/+0.12 +0.06/+0.14, +0.14/+0.06, +0.03/+0.17, +0.07/0.00, +0.02/+0.15 (same code path)
IQ2_XXS -0.06/+0.07, -0.43/+0.12, +0.16/+1.34 +0.03/+0.64, -0.26/+0.51, +0.22/+0.51, -0.02/+0.63, +0.13/+0.57 (same code path)
all-popcount (rejected), IQ2_XS +1.53..+1.95 med, +1.65..+2.11 min

IQ2_S: removing its dead sign copy / shrinking TAB came out +1..3% but within that run's noise (A/A spread up to 11%). Left unchanged.

moe_gu (decode, expert_parity --bench --ab, cand/base gate/up ratio over NE=1,2,3,4 same + ntok=1,2,4,8 sh=0.3) #

  • IQ2_XS ctab: 0.953 0.995 0.965 0.975 0.945 0.954 0.971 0.968 (earlier run 0.947..0.993)
  • IQ2_XS tab (unchanged path): 0.992..1.007
  • IQ3_XXS ctab (unchanged path): 0.998..1.006, plus one 1.064 at ntok=8
  • IQ3_XXS tab: 1.047..1.076 (non-default path, see trade-off above)
  • IQ2_XXS ctab / tab (unchanged path): 0.993..1.006 / 0.996..1.018

dq_f16 (DQ_AB_KERNELS=base dense_mm_parity --dq-ab TYPE, CHUNK totals) #

  • A/A: IQ4_XS 60.20 -> 60.39 ms, IQ3_S 55.14 -> 55.09
  • IQ4 4-register LUT: IQ4_XS 60.06 -> 60.61 (2560x128: 0.78x, 2560x48: 0.85x), IQ4_NL 65.44 -> 65.45 (320x10240p4: 0.82x)
  • IQ4 constant kvalues_iq4nl: 56.53 -> 57.15
  • perm guard: 56.55 -> 56.51
  • 16-byte stores: IQ4_XS 56.75 -> 56.82, IQ3_S 55.18 -> 55.12, Q8_0 80.27 -> 76.45 (one 2560x10240 outlier; the other 12 shapes 1.00x)

gemv unroll pragmas (gemv_parity --one, 3 alternating runs per shape, min per tile) #

best tile base -> cand: IQ4_XS 2560x10240 104.6 -> 114.1 (+9.2%; first sweep +24%); Q6_K 2560x10240 160.0 -> 164.5 (+2.8%); IQ4_NL 640x2560 +3.3%; Q2_0 2560x640 +0.6%; Q3_K -2.2%; Q4_K -0.9%; Q5_K -0.7%; IQ3_S -0.3%; Q4_0 -0.1%; BF16 -0.1%. Rejected.

New test mode (in patch) #

DQ_AB_KERNELS=<dir> dense_mm_parity --dq-ab [type] [shape] compiles <dir>'s quant/units/mma/dense_mm next to this build's. It runs both at dq_tg's threadgroup size, memcmps the fp16 output and times them in paired rounds. The warm-up now covers both sides (before, the first shape was biased against A by ~14%). Without the env var the mode is unchanged (64 threads vs dq_tg's).

Notes / caveats #

  • ksg picks by at (the grid's word count: 256 IQ3_XXS, 512 IQ2_XXS, 1024 IQ2_XS). It is a compile-time constant at every call site (quant.metal and moe_mm.metal), so the branch folds away. moe_mm.metal's call sites didn't need changing.
  • The moe_mm.metal:691 ks pointer (tab + 256) is unused, so the IQ3_XXS shrink is safe. The only other readers of tab + G go through ksg.
  • About the heartbeat's note: my moe_mm runs used kernel dir = base and --ab = candidate, so A = base and B = candidate. Negative means the candidate is faster, as reported here (the A/A controls confirm the direction).
  • Not tested end to end (no full-model runs, per the rules). Expected effect: decode routed gate/up for the IQ2_XS K512 pack about -4%; prompt mm_gu3 for IQ3_XXS packs about -3%.