7.2 kB 51 lines
aa551030
Round 32 blitz: qsa_flash sp[4] 3->1 barrier & ko[j] offset + unroll (-4.0..-5.5% / -6.1 ms in kprof), mm_gu3 T_IQ2_S half4 HS store (-9.0 ms in kprof -> 3175.3 ms / 624.5 tok/s GPU prefill!), deq uint2 grid loads + deq_iq_packed (-> 408.75 tok/s prefill & 42.16 tok/s decode!)
dawn
⋯
audits/argon_r32_moe.md
sourcerendered
download raw history
at main Markdown

Round 32 Track C: Decode deq<T_IQ2_S / T_IQ2_XXS / T_IQ2_XS / T_IQ3_S> 64-Bit uint2 Grid Loads (quant.metal) + gu_body Low-Register Packed iq_combine (moe.metal) #

1. Summary of Changes #

  • Patch: /tmp/argon_r32_moe/argon_r32_moe.patch (verified clean via git -C vendor/strata apply --check, exit 0).
  • Files Modified (within ownership boundaries):
    1. metal/kernels/quant.metal (lines 160–224):
      • deq<T_IQ2_XXS>, deq<T_IQ2_XS>, deq<T_IQ2_S>: Replaced the two separate 32-bit grid lookups per l2 iteration (g_iq2*(tab, idx, 0) and g_iq2*(tab, idx, 1)) with a single 64-bit uint2 load (const uint2 g = ctab ? as_type<uint2>(iq2*_grid[idx]) : ((threadgroup const uint2*)tab)[idx]), cutting threadgroup/constant grid loads in half (4 -> 2 loads per 16-weight unit).
      • deq<T_IQ3_S>: Replaced the per-iteration uchar sign loads b[74 + 4 * ib + l] with one hoisted 16-bit ushort load s2 = *(device const ushort*)(b + 74 + 2 * ui) before the l2 loop, added #pragma clang loop unroll(full), and replaced float(1 + 2 * ((sc >> (4 * (ib & 1))) & 15)) with fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f).
    2. metal/kernels/moe.metal (lines 376–549):
      • Added iq_combine(uint4 g, uint s16, float sc, thread const float* xs) and deq_iq_packed<F>(b, ui, g, s16, tab) for the 16-weight i-quant formats (T_IQ2_S, T_IQ2_XXS, T_IQ2_XS, T_IQ3_S, T_IQ3_XXS), keeping a decoded 16-weight unit in 6 registers (uint4 g, uint s16, float sc) instead of 17 registers (float q[16] + float sc).
      • In gu_body<F, LPR, NR, NE>, enabled fused gate+up decoding and single X<F>::load reuse per token at NR == 1 && NE == 4 with loop-invariant ui = (LPR % 16 == 0) ? (li & 15) : (u & 15), holding both gate and up packed units (12 registers total vs 34 floats) across the t = 0..3 token loop and cutting X<F>::load calls in half (8 -> 4 per unit step) without register spilling.

2. Verification & Bitwise Parity #

All standalone parity and A/B suites pass with 0 fails (ALL PASS, 100% bitwise identical):

  • expert_parity --ab /tmp/argon_r32_moe/base_kernels: 0 fail(s) ALL PASS across IQ2_S, IQ2_XXS, IQ2_XS, IQ3_S, IQ3_XXS and tiles 32,2,4, 32,1,4, 16,1,4 (NE = 1..4, 8 + random windows).
  • gemv_parity --ab /tmp/argon_r32_moe/base_kernels IQ3_S 2560 2560 {1,2,4,8} gemv: bitwise (ALL PASS).
  • Full gemv_parity, dense_mm_parity, expert_parity: ALL PASS.

3. Paired A/B Benchmarks on Quiet Apple M5 (expert_parity --ab --bench --layers 1 --reps 16) #

A. Threadgroup-Table Mode (QUANT_CTAB=0, used by IQ3_S in window.cpp and default expert_parity) #

Format Tile (L,R,S) NE=1 Ratio NE=2 Ratio NE=3 Ratio NE=4 Ratio ntok=1 Ratio ntok=2 Ratio ntok=4 Ratio ntok=8 Ratio
IQ3_S 32,2,4 0.9894 (-1.1%) 0.9713 (-2.9%) 0.9290 (-7.1%) 0.9254 (-7.5%) 0.9817 (-1.8%) 0.9836 (-1.6%) 0.9763 (-2.4%) 0.9722 (-2.8%)
IQ3_S 32,1,4 0.9757 (-2.4%) 0.9857 (-1.4%) 0.9688 (-3.1%) 0.8557 (-14.4%) 0.9842 (-1.6%) 0.9756 (-2.4%) 0.9711 (-2.9%) 0.9779 (-2.2%)
IQ3_S 16,1,4 1.0009 0.9753 (-2.5%) 0.9593 (-4.1%) 0.8355 (-16.5%) 0.9960 (-0.4%) 0.9893 (-1.1%) 0.9719 (-2.8%) 0.9641 (-3.6%)
IQ2_S 32,2,4 0.9890 (-1.1%) 0.9917 (-0.8%) 1.0003 0.9965 (-0.4%) 0.9915 (-0.9%) 0.9918 (-0.8%) 0.9931 (-0.7%) 0.9915 (-0.9%)
IQ2_S 32,1,4 0.9993 (-0.1%) 0.9963 (-0.4%) 0.9713 (-2.9%) 0.8686 (-13.1%) 0.9996 0.9976 (-0.2%) 0.9962 (-0.4%) 0.9949 (-0.5%)
IQ2_S 16,1,4 1.0000 0.9973 (-0.3%) 0.9593 (-4.1%) 0.8416 (-15.8%) 1.0064 1.0008 0.9958 (-0.4%) 0.9930 (-0.7%)
IQ2_XXS 32,2,4 0.9928 (-0.7%) 0.9976 (-0.2%) 0.9943 (-0.6%) 0.9961 (-0.4%) 0.9926 (-0.7%) 0.9906 (-0.9%) 0.9921 (-0.8%) 0.9915 (-0.9%)
IQ2_XS 32,2,4 0.9777 (-2.2%) 0.9947 (-0.5%) 0.9960 (-0.4%) 0.9934 (-0.7%) 0.9734 (-2.7%) 0.9819 (-1.8%) 0.9815 (-1.9%) 0.9830 (-1.7%)
IQ2_XS 32,1,4 0.9925 (-0.8%) 0.9998 1.0180 0.9048 (-9.5%) 0.9941 (-0.6%) 0.9952 (-0.5%) 0.9907 (-0.9%) 0.9966 (-0.3%)
IQ3_XXS 16,1,4 0.9995 1.0050 1.0009 0.8900 (-11.0%) 1.0015 0.9987 (-0.1%) 0.9986 (-0.1%) 0.9964 (-0.4%)

B. Constant-Table Mode (QUANT_CTAB=1, used by IQ2_S, IQ2_XXS, IQ2_XS, IQ3_XXS in window.cpp) #

Format Tile (L,R,S) NE=1 Ratio NE=2 Ratio NE=3 Ratio NE=4 Ratio ntok=1 Ratio ntok=2 Ratio ntok=4 Ratio ntok=8 Ratio
IQ2_S 32,1,4 1.0023 1.0000 1.0005 0.8580 (-14.2%, 237.8 -> 204.0 us) 1.0016 1.0028 1.0012 0.9969 (-0.3%)
IQ2_S 16,1,4 1.0002 1.0003 0.9998 0.8533 (-14.7%, 225.0 -> 192.0 us) 0.9958 (-0.4%) 1.0015 0.9994 0.9989 (-0.1%)
IQ2_XXS 16,1,4 1.0021 1.0005 1.0005 0.8899 (-11.0%, 220.1 -> 195.8 us) 1.0010 1.0022 0.9991 0.9969 (-0.3%)
IQ2_XS 16,1,4 1.0005 1.0009 1.0000 0.9057 (-9.4%, 223.2 -> 202.2 us) 1.0017 1.0020 0.9979 (-0.2%) 0.9974 (-0.3%)
IQ3_XXS 16,1,4 1.0009 1.0032 1.0003 0.8866 (-11.3%, 230.1 -> 204.0 us) 1.0009 1.0037 1.0007 1.0000

C. Dense gemv_parity --ab (IQ3_S 2560 -> 2560, tile 32,2,4,1) #

  • ntok=1: 24.56 us -> 24.04 us (-2.1%), bitwise identical.
  • ntok=2: 30.87 us -> 30.61 us (-0.8%), bitwise identical.
  • ntok=8: 77.12 us -> 76.77 us (-0.5%), bitwise identical.

4. Tried & Rejected Ideas #

  • Preloading 32-bit q4 = rd32(b + 2 + 4 * ui) in deq<T_IQ3_S>: Tested against hoisted s2 + #pragma clang loop unroll(full) + fma; adding q4 = rd32(...) increased live register pressure across the four divergent g_iq3s lookups in deq<T_IQ3_S> and lost the -2.9%..-7.5% gain. Keeping the byte index loads b[2 + 8 * ib + 2 * l] while hoisting the 16-bit sign load s2 + unroll(full) + fma preserved the full -1.6%..-7.5% speedup on IQ3_S.
  • Fused deq_iq_packed + iq_combine at NR == 2 or NE == 2, 3: At NR == 2 (32,2,4), holding two packed units or adding per-element sign selection inside the token loop was +1.2%..+3.7% slower than unit_combine<16>. At NR == 1, NE == 3, the extra bitwise sign extractions across 3 tokens cost +2.9%..+5.0% more than saving 3 X<F>::load calls; at NR == 1, NE == 4, where the baseline float q[16] path suffered occupancy pressure from 8 X<F>::load calls and 17 live floats alongside ag[1][4] + au[1][4], deq_iq_packed + iq_combine delivered a clean -9.4% to -16.5% speedup across all five 16-weight i-quant formats.