Round 29 Track C: Decode moe_route Unrolls + moe_gu / deq<T_IQ2_S / T_IQ2_XXS / T_IQ3_XXS / T_IQ2_XS> Optimization #

1. Summary of Changes #

• metal/kernels/quant.metal (lines 145–201):
  • deq<T_IQ2_S> (34 of 48 routed-expert layers in k512-iq2_xs-metal.gguf): Replace the four per-iteration 8-bit uchar loads and variable shift (qh << (8 - 2 * l)) with two 2-byte-aligned 16-bit ushort loads (q2 = *(device const ushort*)(b + 2 + 2 * ui) and s2 = *(device const ushort*)(b + 34 + 2 * ui)), pre-shift qh = uint(b[66 + ib]) >> (4 * h) once outside the loop, and add #pragma clang loop unroll(full) on l2 = 0, 1.
  • deq<T_IQ2_XXS> (12 of 48 routed-expert layers): Eliminate a0 = rd32(b + 2 + 8 * ib) (two ushort loads + shift + OR + variable shifts by 8 * l) in favor of direct byte loads qs = b + 2 + 4 * ib + 2 * ui (qs[0], qs[1]), pre-shift s0 = a1 >> (14 * h) once before the loop, and add #pragma clang loop unroll(full).
  • deq<T_IQ3_XXS>: Pre-shift s0 = aux >> (14 * h) once before the loop and add #pragma clang loop unroll(full).
  • deq<T_IQ2_XS>: Add #pragma clang loop unroll(full) on l2 = 0, 1.
  • deq<T_IQ3_S>: Kept baseline (unrolling l2 across two g_iq3s table lookups per iteration was +0.9% slower on moe_gu 32,1,4 due to register pressure).
• metal/kernels/moe.metal (moe_route lines 163–334, gu_body lines 381–410):
  • Add e512 = (a.n_expert == 512u) fast path and #pragma clang loop unroll(full) across all 16-element per-lane loops and SIMD xor reduction loops in moe_route.
  • Hoist has_r1, has_moce, and min_scaled_rank outside the 10-iteration rank loop, split the !has_moce vs has_moce rank loops, and preload the 16-bit residency mask res_mask from smap once before the rank loop when has_moce is active (removing 16 conditional device loads from inside the 10-iteration rank loop).
  • Replace the dynamically-indexed array write if ((ex & 31) == lane) v[ex / 32] = -INFINITY; with a branchless unrolled select #pragma clang loop unroll(full) for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; outside the divergent if ((ex & 31) == lane) block, keeping v[0..15] in pure scalar registers without dynamic indexing or divergent predication.
  • In gu_body, hoist etok[0] and advance x and xs16 in place before the u loop when NE == 1 (zero additional live registers; hoisting int tok[NE] or preloading xs[2][16] for NE >= 2 was tested and rejected due to +1..+3% occupancy regression on 32,1,4).
• metal/tests/expert_parity.cpp (lines 220–233):
  • Extend --ab to accept a baseline kernel directory path (compiling quant.metal, units.metal, and moe.metal from that directory) as well as a single .metal file path.

2. Verification & Paired A/B Benchmarks (Apple M5) #

Bitwise Parity #

• route_parity: 200 windows K=512 + 200 windows K=256 + MoCE=0.85: 0 mismatches (PASS)
• route_parity --ab: 200 trials (moe_route + moe_plan, split = 0, 1): 0 fails (ALL PASS, 100% bitwise identical)
• expert_parity --ab: IQ2_S, IQ2_XXS, IQ3_S, IQ2_XS, IQ3_XXS (NE=1..4,8 + random windows): 0 fails (ALL PASS, 100% bitwise identical)
• gemv_parity: All formats, shapes 2560/640/10240, ntok = 1, 3, 8: ALL PASS

Paired A/B Timing: moe_route (route_parity --ab --bench, 192 reps x 9 interleaved rounds) #

Paired A/B Timing: moe_gu (expert_parity --ab --bench, QUANT_CTAB=1) #

Across the 48 routed-expert layers of k512-iq2_xs-metal.gguf (34 IQ2_S + 12 IQ2_XXS + 2 IQ3_S), this saves ~110 us/tok at ntok=1, ~230 us/tok at ntok=2, and ~470 us/tok at ntok=4 in moe_gu alone, plus ~75 us/tok in moe_route (split=1).