6.1 kB 70 lines
a6202763
Round 24 blitz: W8<T_Q6_K> (+12.8..24.9%), qsa_flash 32-bit ko/vo (-11.7% @8K), TAB 1024/512 + gr_norm_h_moec k-outer + gemm_hc_down *0.25f
dawn
⋯
audits/argon_r24_moemm.md
sourcerendered
download raw history
at main Markdown

Round 24 Track C Audit: T_IQ2_XS / T_IQ2_XXS Dead Sign-Table Elimination + deq sign4_iq2xs + W8<T_Q6_K> / W4<T_Q6_K> Bitwise Folding #

1. Summary of Changes (metal/kernels/quant.metal, metal/kernels/moe_mm.metal) #

Target 1: Dead Sign-Table (load_signs / ksg) Elimination & TAB Shrinking in quant.metal #

  1. Shared sign4_iq2xs(uint j, int hh) in quant.metal:
    • Moved static inline uint sign4_iq2xs(uint j, int hh) { return hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j; } to quant.metal so both quant.metal (deq) and moe_mm.metal (W4, W8) share it.
  2. Zero-Lookup / Half-popcount deq<T_IQ3_XXS>, deq<T_IQ2_XXS>, deq<T_IQ2_XS> + fma Scale Math:
    • Because put_signed4(q, g, s) only inspects bits 0..3 of s, the first 4 values (q + 8 * l2) pass raw 7-bit s directly (zero table lookup, zero popcount), and the second 4 values (q + 8 * l2 + 4) pass sign4_iq2xs(s, 1).
    • Folded qs = b + 2 + 4 * ui in deq<T_IQ3_XXS> and qs = (device const ushort*)(b + 2) + 2 * ui in deq<T_IQ2_XS>, and replaced (0.5f + float(s)) * 0.25f / 0.5f with fma across deq<T_IQ3_XXS>, deq<T_IQ2_XXS>, deq<T_IQ2_XS>, and deq<T_IQ2_S>.
  3. Eliminated ksg and load_signs, and Shrunk TAB:
    • Removed ksg and load_signs completely from quant.metal (including the unused load_signs(tab, 2048, tid, nth) in load_tab<T_IQ2_S>).
    • Shrunk fmt<T_IQ2_XXS>::TAB from 512 + 32 -> 512 and fmt<T_IQ2_XS>::TAB from 1024 + 32 -> 1024.
    • Kept fmt<T_IQ2_S>::TAB = 2048 + 32 so mm_gu3<T_IQ2_S> stays above the 16 KiB threadgroup occupancy boundary (ot[4][512] + tab = 16,512 B, capping at 2 TGs/core; shrinking T_IQ2_S to 2048 regressed -0.55%, whereas keeping 2048 + 32 without load_signs gains +0.19% to +0.33%).

Target 2: W8<T_Q6_K> and W4<T_Q6_K> Bitwise Address & Scalar Unpack Folding in moe_mm.metal #

  1. Bitwise Address & Shift Folding (j = uint(k & 255), j4 = j >> 4):
    • ql offset 64 * n + ((qd & 1) ? 32 : 0) + l folded to ((j >> 1) & 64u) | (j & 63u) (since ((qd & 1) << 5) | l == j & 63u).
    • qh offset 128 + 32 * n + l folded to 128u + (((j >> 2) & 32u) | (j & 31u)).
    • sc offset 192 + 8 * n + 2 * qd + (l / 16) folded to 192u + j4 (sc[0] and sc[1] in W8<T_Q6_K> where l < 16).
    • Shift counts sl = (qd >> 1) ? 4u : 0u and shh = uint(2 * qd) folded to sl = j4 & 4u and shh = j4 & 6u.
  2. Scalar uint 6-Bit Code Assembly Before as_type<uchar4>:
    • Replaced the two 4-wide uint4 promotions + vector shift + vector OR (uint4(as_type<uchar4>((la >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type<uchar4>((ha >> shh) & 0x03030303u)) << 4)) with a single scalar uint OR qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4) followed by float4(as_type<uchar4>(qa)) in both W8<T_Q6_K> and W4<T_Q6_K>.
    • Note: we also tested scalar uint assembly in W8<T_Q5_K> and char2 scale loads in W8<T_Q6_K>; both caused compiler load hoisting that reduced threadgroup occupancy at T=2048 (-31% on Q5_K, -20% on Q6_K) and were rejected/reverted.

2. Bitwise Parity Verification (memcmp == 0) #

All standalone parity suites pass with 100% bitwise parity (memcmp == 0 where checked):

  • moe_mm_parity --v3 --ab /tmp/argon_r24_moemm/base_kernels across all 11 gate/up and down format pairs (Q2_0/Q2_0, IQ3_XXS/IQ4_NL, IQ2_XXS/Q5_1, IQ2_XS/Q8_0, IQ2_S/Q2_0, IQ3_S/IQ4_NL, IQ4_XS/Q5_1, Q4_K/Q8_0, Q5_K/Q2_0, Q6_K/IQ4_NL, Q8_0/Q5_1, plus Q6_K/Q2_0 at TK=16 exercising W4<T_Q6_K>):
    • ab bitwise (T=64): HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS on all 12 tests.
  • moe_mm_parity --v2 (IQ3_XXS, IQ2_XXS, IQ2_XS, IQ2_S, Q6_K): ALL PASS.
  • expert_parity (Q2_0, IQ3_XXS, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_S, Q6_K): ALL PASS.
  • gemv_parity: ALL PASS (bitwise window == one token across all formats).
  • dense_mm_parity: ALL PASS (memcmp(wh)=0, memcmp(xh)=0 across all formats).

3. Paired A/B Benchmark Results (A = Candidate /tmp/argon_r24_moemm/strata/metal/kernels, B = Baseline /tmp/argon_r24_moemm/base_kernels) #

Command:

STRATA_METAL_KERNELS=/tmp/argon_r24_moemm/strata/metal/kernels \
  /tmp/argon_r24_moemm/strata/build-metal/moe_mm_parity \
  --v3 --gu <FMT> --dn Q2_0 --tile 16,32,32,4 --dtile 64,32,2 --T <T> \
  --ab /tmp/argon_r24_moemm/base_kernels --gu-bench
Format T Candidate A med (min) us Baseline B med (min) us diff = (B - A)/A med (min)
Q6_K 128 18616.0 (18609.9) 21185.7 (20958.0) +13.80% (+12.62%)
Q6_K 512 20437.1 (20328.2) 23052.0 (22870.8) +12.79% (+12.51%)
Q6_K 1024 26174.7 (22409.5) 29390.0 (26684.4) +12.28% (+19.08%)
Q6_K 2048 35518.8 (35495.2) 36823.4 (36618.8) +3.67% (+3.17%)
IQ2_XS 512 7678.7 (7670.4) 7681.9 (7677.9) +0.04% (+0.10%)
IQ2_XS 1024 12820.3 (12792.4) 12830.2 (12826.5) +0.08% (+0.27%)
IQ2_XS 2048 21932.5 (21892.8) 21967.1 (21918.3) +0.16% (+0.12%)
IQ2_XXS 512 8104.6 (8076.7) 8102.3 (8080.1) -0.03% (+0.04%)
IQ2_XXS 1024 13477.9 (13449.4) 13497.4 (13463.2) +0.14% (+0.10%)
IQ2_XXS 2048 22948.2 (22883.5) 22987.7 (22912.5) +0.17% (+0.13%)
IQ2_S 512 10102.7 (10073.4) 10122.2 (10111.2) +0.19% (+0.37%)
IQ2_S 1024 16946.7 (16882.8) 16982.6 (16942.4) +0.21% (+0.35%)
IQ2_S 2048 28478.5 (28464.5) 28571.3 (28554.5) +0.33% (+0.32%)

In expert_parity --bench (ntok=1, QUANT_CTAB=1, production decode path in window.cpp), IQ2_XS moe_gu also improves from 93.4 us (101.4 GB/s) to 92.5 us (102.4 GB/s, +1.0%).

4. Artifact Paths #

  • Clean unified diff (git -C vendor/strata apply --check verified): /tmp/argon_r24_moemm/argon_r24_moemm.patch
  • Audit report: audits/argon_r24_moemm.md