6.4 kB 62 lines
e1737f7d
Round 23 blitz: ksg popcount + gemv_direct_x/gemv_q2_raw_combine + mm_gu3 sign4_iq2xs/ptr-folding -> 42.01 tok/s decode & 602.1 tok/s GPU prefill
dawn
⋯
audits/argon_r23_moemm.md
sourcerendered
download raw history
at main Markdown

Round 23 Track C: moe_mm.metal Prompt Prefill mm_gu3 & mm_down3 Optimization (argon-moe-mm-prompt-opt) #

1. Summary of Shipped Changes #

  • Deliverable patch: /tmp/argon_r23_moemm/argon_r23_moemm.patch (git -C vendor/strata apply --check verified clean, exit 0).
  • Touched files: metal/kernels/moe_mm.metal (lines 507–574, 688–765) and metal/tests/moe_mm_parity.cpp (lines 306–336).

Key Optimizations in metal/kernels/moe_mm.metal #

  1. 4-bit Register Sign-Parity Helper sign4_iq2xs(uint j, int hh) (W4 & W8 for T_IQ2_XS, T_IQ2_XXS, T_IQ3_XXS):
    • w4_signed(g, s, sc) only reads the low 4 bits (s & 1, s & 2, s & 4, s & 8) of its second argument.
    • When hh == 0 (50% of SIMD lanes), the low 4 bits of ksigns_iq2xs[j] are simply j itself (v >> 9) — bit 7 (popcount(j) & 1u) is not needed at all.
    • When hh == 1, (ksigns_iq2xs[j] >> 4) & 15 is (j >> 4) | ((popcount(j) & 1u) << 3).
    • sign4_iq2xs(j, hh) = hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j eliminates divergent threadgroup byte loads ksg(tab, ..., j) in W4<T_IQ2_XS>, W8<T_IQ2_XS>, W4<T_IQ2_XXS>, W8<T_IQ2_XXS>, W4<T_IQ3_XXS>, and W8<T_IQ3_XXS>, while avoiding popcount on hh == 0 lanes.
    • Note on moe_mm_parity --ab sign convention: moe_mm_parity.cpp:424 reports diff = (mb - ma) / ma * 100.0 where A is the candidate and B is --ab, so positive +X% means A is faster than B (ma < mb). Direct head-to-head A/B comparison of A = ksg vs B = sign4_iq2xs (with all other changes identical) confirmed B (sign4_iq2xs) is 3.47%–10.63% faster than A (ksg).
  2. Constant-Offset Pointer Folding ((k & 255) >> 3 and (k & 255) >> 2):
    • In W8<T_IQ2_XS>, 4 * ib + l (with ib = (k & 255) >> 5, l = (k & 31) >> 3, l < 2) is identical to (k & 255) >> 3. Computing device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3) turns q[l], q[l + 2] into constant-offset loads q[0], q[2] and eliminates l.
    • In W8<T_IQ3_S>, 8 * ib + 2 * l + hh is (k & 255) >> 2, 4 * ib + l is (k & 255) >> 3, and 2 * l + hh is (k & 31) >> 2. Folding these into qs = bl + 2 + ((k & 255) >> 2) (qs[0], qs[4]) and sp = bl + 74 + ((k & 255) >> 3) (sp[0], sp[2]) eliminates l in the hot path and reduces VGPR pressure across the 4-way unrolled mm_gu3 loop.
    • Applied the same constant-offset pointer folding to W8<T_IQ2_S>, W8<T_IQ3_XXS>, W4<T_IQ2_XS>, W4<T_IQ2_S>, W4<T_IQ3_S>, and W4<T_IQ3_XXS>.
  3. Bitwise-Identical fma Scale Math under MTLMathModeSafe:
    • Because 0.25f ($2^{-2}$) and 0.5f ($2^{-1}$) are exact powers of two and rdh(bl) is an fp16 value promoted to fp32 (exponent $\ge -24$), d * fma(float(nib), 0.25f, 0.125f) is 100% bitwise identical to d * (0.5f + float(nib)) * 0.25f across all 65,536 fp16 bit patterns and all 16 nibble values 0..15, while replacing FADD + FMUL + FMUL (which MTLMathModeSafe cannot reassociate) with FMA + FMUL and removing rdh(bl) from the first stage of the dependency chain.
  4. Ruled-Out Candidates (Not Shipped):
    • Custom W8<T_Q2_0> in mm_down3: tested both shared-d and split-d W8<T_Q2_0> specializations; both regressed mm_down3 by 0.7%–3.0% at T=1024, 2048 because CSEing the block pointer across the two W4 halves allows LLVM to hoist all 16 loads across the 8-way unrolled loop, increasing VGPR pressure near the occupancy boundary. Kept the generic two-W4<T_Q2_0> path untouched.
    • W8<T_Q6_K> scalar byte packing: regressed Q6_K by 2.4%–4.5% vs uint4(as_type<uchar4>(...)). Reverted so W8<T_Q6_K> is untouched.

2. Paired A/B Benchmark Results on M5 (moe_mm_parity --v3 --tile 16,32,32,4 --dtile 64,32,2 --gu-bench) #

Note: In moe_mm_parity --ab, A is the candidate kernel, B is the baseline reference kernel, and diff = (B - A) / A * 100% (+ = candidate A is faster than baseline B).

Format T Candidate A Med (Min) us Baseline B Med (Min) us Speedup diff Med (Min)
IQ2_XS 512 7,682.3 (7,669.9) 7,982.3 (7,978.3) +3.91% (+4.02%)
IQ2_XS 1024 16,603.6 (12,901.3) 17,684.0 (13,439.3) +6.51% (+4.17%)
IQ2_XS 2048 25,535.6 (25,486.3) 26,712.9 (26,492.7) +4.61% (+3.95%)
IQ3_S 512 17,946.8 (17,777.0) 19,823.0 (19,711.9) +10.45% (+10.88%)
IQ3_S 1024 28,841.0 (28,506.7) 31,010.2 (30,957.8) +7.52% (+8.60%)
IQ3_S 2048 48,501.1 (46,976.4) 52,371.5 (51,207.6) +7.98% (+9.01%)
IQ2_S 512 35,342.2 (34,052.6) 36,953.1 (35,455.6) +4.56% (+4.12%)
IQ2_S 1024 44,448.1 (42,382.8) 45,269.8 (44,114.7) +1.85% (+4.09%)
IQ2_XXS 1024 32,583.5 (30,793.0) 33,326.4 (32,109.3) +2.28% (+4.27%)
IQ2_XXS 2048 56,247.3 (54,307.5) 58,809.2 (56,789.6) +4.55% (+4.57%)
IQ3_XXS 512 28,919.5 (22,779.1) 33,020.1 (24,348.8) +14.18% (+6.89%)
IQ3_XXS 1024 33,955.7 (29,351.1) 35,193.0 (32,111.7) +3.64% (+9.41%)
IQ3_XXS 2048 50,338.7 (45,779.3) 53,857.8 (50,270.8) +6.99% (+9.81%)

3. Bitwise Parity Verification (HS, YE, Y memcmp == 0) #

Verified 100% bitwise identity (HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS) between candidate and unmodified baseline across all formats (IQ2_XS, IQ3_S, Q6_K, Q2_0, IQ2_S, IQ2_XXS, IQ3_XXS) at T = 256, 512, 1024, 2048 on both TK=32 (W8) and TK=16 (W4):

[W8 TK=32] IQ2_XS   T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W8 TK=32] IQ3_S    T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W8 TK=32] Q6_K     T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W8 TK=32] Q2_0     T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W8 TK=32] IQ2_S    T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W8 TK=32] IQ2_XXS  T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W8 TK=32] IQ3_XXS  T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS
[W4 TK=16] IQ2_XS..IQ3_XXS T=256: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS