Round 11 Track C Audit: mm_gu3 & mm_down3 Row-Pointer / Block-Offset Hoisting & Branchless W8<T_IQ3_S> / W4<T_IQ3_S> #

1. Summary of Changes (metal/kernels/moe_mm.metal & metal/tests/moe_mm_parity.cpp) #

• Patch file: /tmp/argon_r11_moemm/argon_r11_moemm.patch (verified clean via git -C vendor/strata apply --check).
• Scratch checkout: /tmp/argon_r11_moemm/strata (tracked vendor/strata/ untouched).
• Files touched: Only metal/kernels/moe_mm.metal and metal/tests/moe_mm_parity.cpp.

Change 1: Hoist Row Base Pointer (g0_l, u0_l, w0_l), Row Stride Step (step8), kl, and Per-k0 Block Offset (boff) in mm_gu3 and mm_down3 (moe_mm.metal lines 852–885 & 983–1007) #

• On Apple M5 (coop_k4 && TK == 32), each lane's right-input cooperative tensor holds half_cap / 4 runs (4 runs in mm_gu3 16x32x32, 8 runs in mm_down3 16x64x32) at consecutive row offsets n_r = n_0 + 8 * r (r = 0 .. half_cap/4 - 1) and a fixed intra-tile K offset kl = ix0[0] = (lane & 3) * 4.
• Previously, inside every k0 iteration (80 times in mm_gu3, 20 times in mm_down3):
  1. mm_gu3 recomputed g0 + ulong(n) * a.gu_row 4 times for gate and 4 times for up (640 64-bit multiplies per lane), and W8<F>::run recomputed row + (k >> 8) * BB 8 times per k0 (640 block-offset multiplies per lane).
  2. mm_down3 recomputed w0 + ulong(ix[1]) * a.d_row 8 times per k0 (160 64-bit multiplies per lane), and W8<F>::run recomputed row + (k >> 6) * BB 16 times per k0 (320 block-offset multiplies per lane).
• Hoisted kl = ix0[0], g0_l = g0 + ulong(ix0[1]) * a.gu_row, u0_l = u0 + ulong(ix0[1]) * a.gu_row, w0_l = w0 + ulong(ix0[1]) * a.d_row, and step8 = 8u * row_bytes outside the k0 loop, and computed boff and the masked intra-block index j_in ((k0 + kl) & 255 for QK == 256, (k0 + kl) & 63 for QK == 64, kl for QK == 32) once per k0 iteration, advancing a single running pointer rgw += step8 / ruw += step8 / rdw += step8 across the unrolled i loop.
• Because j_in is masked to < QK before calling W8<F>::run(rgw, j_in, ...), LLVM KnownBits eliminates (j_in >> 8) * BB (or (j_in >> 6) * BB) and j_in & 255 inside every W8<F>::run call without duplicating W8<F> specializations.
• Excluded F == T_Q6_K from the outer pointer-hoisting branch in mm_gu3 (if (coop_k4 && TK == 32 && F != T_Q6_K && mm_probe != 1)) because W8<T_Q6_K> sits at a VGPR occupancy threshold where keeping g0_l, u0_l, step8 live across op.run regressed T = 1024.

Change 2: Eliminate Ternary hh ? ... : ... Branches in W8<T_IQ3_S> and W4<T_IQ3_S> (moe_mm.metal lines 531 & 727–731) #

• In W8<T_IQ3_S>::run, replaced:

  const int l1 = l + 2;
  const uint i0 = hh ? (uint(bl[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256))
                     : (uint(bl[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256));
  const uint i1 = hh ? (uint(bl[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256))
                     : (uint(bl[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256));
  

  with branchless pointer indexing and shift arithmetic (hh is (k & 7) >> 2 in {0, 1}):

  const P qs = bl + 2 + 8 * ib + 2 * l + hh;
  const int sh0 = 8 - 2 * l - hh;
  const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u);
  const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u);
  const uint s0 = bl[74 + 4 * ib + l], s1 = bl[76 + 4 * ib + l];
  

  and similarly in W4<T_IQ3_S>::run:

  const uint i = uint(b[2 + 8 * ib + 2 * l + hh]) | ((qh << (8 - 2 * l - hh)) & 256u);
  

• Eliminates two divergent byte-load pairs and select instructions per W8<T_IQ3_S> call (8 per k0 iteration = 640 per lane), shrinking VGPR usage across the unrolled decode loop and speeding up mm_gu3<T_IQ3_S> by -25.0% at T = 2048, -50.8% (2.03x faster!) at T = 1024, and -32.7% at T = 512.

Change 3: In-Process Interleaved --ab <kernel_dir> Mode in moe_mm_parity.cpp #

• Added --ab <dir> to moe_mm_parity for --gu-bench and --dn-bench so baseline and candidate libraries are compiled in the same Metal context and timed in 5 alternating pairs (AB, BA, AB, BA, AB) on the exact same weight and activation buffers.

Ruled-Out Candidates #

1. Explicit 4- / 8-element pointer arrays rgw[4], ruw[4], rdw[8] hoisted before k0 (k_hoist1): Forced the Metal compiler to keep 8 64-bit pointers (16 32-bit VGPRs) live across op.run in mm_gu3 and mm_down3, reducing simdgroup occupancy and regressing IQ2_XS --gu-bench by +5.1% .. +6.6% and Q2_0 --dn-bench by +4.2% .. +5.3%. Using a single base pointer g0_l + boff advanced by + step8 inside the unrolled loop keeps only one pointer live during the unrolled loop and zero across op.run.
2. Outer kb (256-element super-block) + inner ks loop in mm_gu3 (k_hoist3b): Added a nested loop induction variable live across op.run, regressing IQ2_XS at T = 1024 by +2.90% (13,848 -> 14,250 us) and T = 2048 by +0.66% compared to the single k0 loop with boff = uint(k0 >> 8) * BB (k_hoist3a, -2.09% at T = 1024 and -2.24% at T = 2048).
3. W8<T_Q6_K> uchar4 bitwise pack (as_type<uchar4>(((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4))): Bitwise identical (memcmp == 0), but caused Apple's Metal backend to scalarize the uchar4 -> float4 conversion differently from float4(uint4), increasing VGPR pressure and regressing Q6_K --gu-bench at T = 1024 by +46.9% (26.2 -> 38.4 ms) and T = 2048 by +12.1% (45.0 -> 50.5 ms). Rejected.

2. Bitwise Parity Verification (100% Byte-for-Byte memcmp == 0) #

Verified --dump output (HS + y, 6,512,640 bytes per format at T = 256, 64 experts, --v3 --tile 16,32,32,4 --dtile 64,32,2 --dn Q2_0) between baseline (k_base) and patched (k_final) across all 11 supported mm_gu3 formats (including all 5 required formats):

3. Paired A/B Benchmark Results on Apple M5 (T = 512, 1024, 2048) #

In-process alternating 5-pair A/B timings (--v3 --tile 16,32,32,4 --dtile 64,32,2 --ab, 512 experts, quiet GPU):

Projected Full-Model Prefill Savings at T = 1983 (48 Layers) #

• mm_gu3 T_IQ2_XS (x34 layers, 807.06 ms baseline): -2.10% -> -16.9 ms
• mm_gu3 T_IQ3_S (x2 layers, 67.78 ms baseline): -25.01% -> -17.0 ms (and -50.8% on T = 1024 chunks!)
• mm_down3 T_Q2_0 (x48 layers, 334.82 ms baseline): -4.58% -> -15.3 ms
• Total Track C savings at T = 1983: ~49.2 ms per 2k chunk (~270 ms across an 11k prompt!) with 100% bitwise parity.