Round 26 Track C Audit Report: DW<T_IQ4_XS> / DW<T_IQ4_NL> / DW<T_Q6_K> Specializations & mm_gu3 (W8<T_IQ2_XS>) Investigation #

• Agent: argon-r26-moemm-gu-and-dqf16 (gemini-4-argon)
• Scratch Repo: /tmp/argon_r26_moemm/strata
• Unified Patch: /tmp/argon_r26_moemm/argon_r26_moemm.patch (git -C vendor/strata apply --check verified 0)
• Touched Files:
  • metal/kernels/mma.metal
  • metal/kernels/dense_mm.metal
  • metal/kernels/moe_mm.metal
  • metal/tests/dense_mm_parity.cpp
  • metal/tests/moe_mm_parity.cpp

1. Target 1: DW<T_IQ4_XS>, DW<T_IQ4_NL>, DW<T_Q6_K> Direct Specializations & Full Loop Unrolling (mma.metal, dense_mm.metal, moe_mm.metal) #

Findings & Implementation #

1. Direct DW<T_IQ4_NL> and DW<T_IQ4_XS> specializations (metal/kernels/mma.metal):
   • Previously, T_IQ4_XS and T_IQ4_NL fell back to the generic DW<F>, which allocated a separate float q[32] array in deq<F> and then multiplied w[j] = q[j] * s in a second loop before dq_f16 / gemv_mma / stage_w copied w[j] out.
   • Added dedicated DW<T_IQ4_NL> and DW<T_IQ4_XS> specializations that fuse the scale multiply float(w0.x) * d directly into the unrolled 8-iteration ushort table lookup loop (#pragma clang loop unroll(full)), eliminating the intermediate float q[32] array and second loop while preserving exact fp32 multiplication order (memcmp == 0).
   • Also tested replacing the threadgroup tab lookup in dq_f16<T_IQ4_XS> / dq_f16<T_IQ4_NL> with direct constant-memory kvalues_iq4nl[...] nibble lookups (skipping load_tab + barrier); this regressed chunk time by ~1% (58.88 -> 59.47 ms on IQ4_XS, 65.61 -> 65.98 ms on IQ4_NL) due to constant-cache divergent lookup serialization on Apple GPUs, so the threadgroup tab lookup is retained.
2. 4-Byte rd32 + Scalar uint Bitwise Unpack in DW<T_Q6_K> (metal/kernels/mma.metal):
   • Replaced the 16-iteration byte-by-byte loop in DW<T_Q6_K> with a 4-iteration unrolled loop using rd32(ql + j) and rd32(qh + j) plus SWAR mask/shift (la & 0x0F0F0F0Fu) | (((ha >> sh) & 0x03030303u) << 4).
   • On small-row dq_f16 dispatches (2560x128, 2560x48, 10240x4, 2560x1), this speeds up dq_f16<T_Q6_K> by 1.04x–1.07x (3.7 -> 3.6 us, 2.2 -> 2.1 us, 1.8 -> 1.7 us, 1.5 -> 1.4 us) and 6144x2560 by 1.02x (322.0 -> 316.2 us), while large matrices remain DRAM-write bandwidth bound at ~100 GB/s.
3. Full Unrolling Across DW<F>, dq_f16, gemv_mma, and stage_w:
   • Added #pragma clang loop unroll(full) to all DW<*> specializations (DW<F>, DW<T_IQ4_NL>, DW<T_IQ4_XS>, DW<T_Q2_0>, DW<T_Q5_1>, DW<T_F32>, DW<T_BF16>, DW<T_Q4_K>, DW<T_Q5_K>, DW<T_Q6_K>, DW<T_Q3_K>), gemv_mma's threadgroup store loop, dq_f16's half4 store loop, and moe_mm.metal's stage_w loop.
4. Paired Alternating A/B Timing in dense_mm_parity.cpp:
   • Updated dq_ab_bench in metal/tests/dense_mm_parity.cpp to alternate execution order (round & 1), filter thermal/contention outliers (<= 1.15 * best_pair), and accept STRATA_DQ_AB_KERNELS in addition to DQ_AB_KERNELS.

Verification (dense_mm_parity & mma_parity) #

• dense_mm_parity: ALL PASS across all 14 formats (F32, BF16, Q8_0, Q4_0, Q5_0, Q5_1, IQ4_NL, IQ4_XS, Q2_0, IQ3_S, Q3_K, Q4_K, Q5_K, Q6_K), memcmp(wh)=0, memcmp(xh)=0 on every shape.
• mma_parity: ALL PASS across all 12 formats (ntok = 1, 3, 8), window == one token: bitwise.
• Quiet-GPU dense_mm_parity --dq-ab paired results (memcmp=0 on all 11–13 shapes per format):
  • IQ4_XS: 60.91 -> 60.95 ms (1.00x, 2560x6144: 293.5 -> 291.3 us 1.01x, memcmp=0)
  • Q6_K: 66.13 -> 65.98 ms (1.00x chunk; 6144x2560: 322.0 -> 316.2 us 1.02x, 2560x48: 2.2 -> 2.1 us 1.06x, 10240x4: 1.8 -> 1.7 us 1.06x, 2560x1: 1.5 -> 1.4 us 1.07x, memcmp=0)
  • IQ4_NL: 66.13 -> 65.94 ms (1.00x, 2560x6144: 286.7 -> 283.7 us 1.01x, memcmp=0)
  • IQ3_S: 59.56 -> 59.17 ms (1.01x, 2560x6144: 278.9 -> 274.4 us 1.02x, 6144x2560: 286.4 -> 282.7 us 1.01x, memcmp=0)

2. Target 2: mm_gu3 (W8<T_IQ2_XS>) Investigation & Ablations (moe_mm.metal, moe_mm_parity.cpp) #

Hardware Profiling (MM_PROBE Breakdown at T=512, tile=16,32,32,4, 512 Experts) #

• MM_PROBE=0 (full mm_gu3 decode + matmul2d): 7723.7 us (4.34 TFLOP/s useful)
• MM_PROBE=3 (no tab lookup, v0 * 0x01010101u, with w4_signed + sign4_iq2xs + device loads): 7320.4 us (-403 us = 5.2% of total time)
• MM_PROBE=5 (device loads v0, v1, scb, d + sc0, sc1 only, no tab lookup, no w4_signed/sign4_iq2xs): 6604.0 us (-716 us = 9.3% of total time)
• MM_PROBE=1 (no weight decode, cheap constant fill + matmul2d): 3680.6 us (-2923 us = 37.8% of total time)
• QUANT_CTAB=1 (constant-memory grid instead of threadgroup tab): 8550.8 us (+10.7% slower)

Controlled Ablations Tested on M5 (T=512 / 1024 / 2048, all HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS) #

1. Pre-added 3-pointer stepping (rgw, rgw_q, rgw_sc += step8):
   • Maintaining three 64-bit pointers across the unrolled i loop added pointer-increment instructions and live registers across op.run(tA, rw, accg), regressing T=512 min time by +0.93% (15269 -> 15411 us under contention).
2. Outer block loop (kb += 256, ib = 0..7) + run_q(rgw, qoff, ib, hh, ...):
   • Keeping qoff = kl_q + 8u * ib, ib, and hh live across op.run(tA, rw, accg) increased register pressure around the cooperative-tensor matmul2d op, regressing T=512 by 0.55% (7748 -> 7791 us), T=1024 by 1.35% (12896 -> 13073 us), and T=2048 by 1.43% (22193 -> 22515 us).
3. Full compile-time template unrolling of IB = 0..7 (W8<T_IQ2_XS>::run_q<IB>):
   • Unlike mm_down3 (Q2_0, QK=64 = 2 * TK), IQ2_XS has QK=256 = 8 * TK. Unrolling all 8 IB steps (16 op.run calls per k0 += 256 iteration) blew up kernel code size and register allocation, regressing T=512 from 7708 us to 33176 us (-76.8%).
4. Vectorized packed_ushort4 load (q4 = *(device const packed_ushort4*)(bl + qoff)):
   • Because 74 % 4 == 2, odd blocks are only 2-byte aligned, so packed_ushort4 loads all 8 bytes (q[0..3]) as 16-bit loads including the unused q[1] and q[3] words for the other l half, matching baseline within 0.00%–0.28% (12925.1 vs 12925.0 us at T=1024).
5. Cross-lane simd_shuffle_xor(vm, 1) sharing between adjacent hh=0 and hh=1 lanes:
   • Both lanes already hit the same L1 cache line in the same cycle; adding simd_shuffle_xor + select instructions inside the unrolled loop regressed T=512 by -40.6% (7744 -> 13028 us) and T=1024 by -16.0% (12914 -> 15367 us).
6. Dual rw + ru cooperative tensors in a single i loop:
   • Allocating two right_input_cooperative_tensor instances simultaneously doubled the half register footprint (32 -> 64 halfs per lane) across op.run, reducing occupancy and regressing T=512 by -9.54% (7735 -> 8550 us) and T=2048 by -4.03% (22217 -> 23149 us).
7. Bitwise float32 sign-XOR (as_type<uint4>(p) ^ m) and w4_iq2xs direct shift:
   • The Apple M5 shader compiler fuses (s & bit) ? -float(v[k]) : float(v[k]) into hardware conditional-negate/select instructions, whereas integer sign-bit XOR required float <-> uint bitwise reinterpretation (-17.4% at T=512), and splitting s3 parity select (w4_iq2xs) added an extra branch/select per w4 (-3.75% at T=512).

Conclusion on mm_gu3 (W8<T_IQ2_XS>) + Fast-Compile Guard (MM_V3_ONLY) #

• Baseline W8<T_IQ2_XS>::run + mm_gu3's single j_in = (k0 + kl) & 255 index is already register- and instruction-optimal on M5: keeping only k0 and kl live across op.run(tA, rw, accg) minimizes live scalar registers during the matmul2d cooperative-tensor op while LLVM CSE shares (j_in & 255) >> 5, (j_in & 255) >> 3, and (j_in & 7) >> 2 across the 4 unrolled i iterations.
• Kept W8<T_IQ2_XS> and mm_gu3 at exact baseline (100% bitwise identical, HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS), and added #ifndef MM_V3_ONLY guards around non-v3 kernels (mm_gu, mm_down, mm_gu2, mm_down2, mm_gu4, mm_gu5) in metal/kernels/moe_mm.metal + #define MM_V3_ONLY 1 in metal/tests/moe_mm_parity.cpp when running --v3, cutting moe_mm_parity --v3 Metal shader compilation time from ~15s to ~1.5s while leaving normal engine compilation 100% unchanged.