Round 22 Track C Audit Report: moe_down (T_Q2_0R/T_Q2_0) Compressed-Register Weight Reuse & moe_route / plan_body Optimization #

• Agent: argon-moe-down-q2r-and-route (gemini-4-argon)
• Scratch workspace: /tmp/argon_r22_moe/strata
• Unified patch: /tmp/argon_r22_moe/argon_r22_moe.patch (git -C vendor/strata apply --check clean, exit code 0)
• Files modified: metal/kernels/moe.metal, metal/tests/expert_parity.cpp

1. Key Findings & Design #

1.1 dn_body & gu_body: Compressed-Register Weight Reuse (q2_combine<F>) Across NR & NE #

Previously, in dn_body (metal/kernels/moe.metal):

• Holding float q[NR][16]; float4 sc[NR]; (40 registers for NR = 2) across r = 0..NR-1 spilled registers on Apple M5 when NE >= 2, which forced dn_body for F == T_Q2_0 || F == T_Q2_0R to place r as the outer loop and reload xs[16] from h NR = 2 separate times per token t.
• In gu_body, F == T_Q2_0 at NE >= 2 reloaded xs[16] 2 * NR = 4 times per token t (g0 and u0 for each r).

By introducing q2_combine<F>(ushort2 pk, float d, thread const float* xs, float xsum0):

1. Compressed register staging (4 registers total for NR = 2 in dn_body, 8 in gu_body):
   • In dn_body (F == T_Q2_0 || F == T_Q2_0R, across all NE = 1..4), we pre-load ushort2 pk[NR]; float d[NR]; for all r = 0..NR-1 first (sharing (u >> 2) * 18 and 2 + 4 * (u & 3) across rows), then loop for (int t = 0; t < NE; ++t), loading X<F>::load(h + (c + t) * int(a.n_ff), ...) once per token t (-50% h device loads at NR = 2, NE >= 2) and accumulating ac[r][t] += q2_combine<F>(pk[r], d[r], xs, xsum[0]) without ever materializing float q[16].
   • In gu_body (F == T_Q2_0 || F == T_Q2_0R, NE >= 2), we pre-load pkg[NR], pku[NR], dg[NR], du[NR] (8 registers for NR = 2), load X<F>::load once per token t (-75% x loads at NR = 2, NE >= 2), and accumulate ag[r][t] and au[r][t] via q2_combine<F>.
   • Added #pragma clang loop unroll(full) to the r < NR and t < NE loops in dn_body and gu_body (keeping r < NR un-pragmad in gu_body when UNIT == 32 to avoid qg[32]/qu[32] register pressure in Q4_K).
2. Exact IEEE-754 safe-math bitwise parity:
   • q2_combine<F> preserves the exact 4-way FMA accumulation order (s0, s1, s2, s3) and scale/offset expression of Dec<F>::run + unit_combine<16>, achieving std::memcmp == 0 across every buffer (plan, h, ye, y).

1.2 moe_route & plan_body: v[16] Indexing Hypothesis Falsified, sg == 0 Parallel SIMD Prefix Sum Shipped #

1. Falsified hypothesis (v[ex / 32] = -INFINITY -> 16-way if (i == wi) loop):
   • Tested in isolation on Apple M5 (route_parity --ab): replacing v[ex / 32] = -INFINITY; with #pragma clang loop unroll(full) for (int i = 0; i < 16; ++i) if (i == wi) v[i] = -INFINITY; regressed moe_route by +22% to +54% (5.12 us -> 7.89 us at ntok=1 split=0, 6.95 us -> 9.76 us at ntok=1 split=1).
   • Root cause: Apple Silicon GPU hardware supports native register-indirect indexing on v[ex / 32] = -INFINITY; (1 instruction in the winning lane), whereas the 16-way if (i == wi) loop emits 16 compare/select instructions per rank (160 extra instructions across 10 ranks) and increases live scalar register pressure in the 512-thread (16-simdgroup) threadgroup. Rejected per rule ("Only ship changes that are proven net wins on M5").
2. Proven win in plan_body (used by moe_route, moe_plan, and router_route):
   • Previously, after the first threadgroup_barrier, thread tid == 0 executed a serial 2 x 16 = 32-iteration loop reading and writing s_grp and s_ent in threadgroup memory while the other 511 threads waited at the barrier, and unconditionally ran simd_prefix_exclusive_sum(miss) / c1 even when !split.
   • Replaced the tid == 0 serial loop with a single-instruction parallel simd_prefix_exclusive_sum across sg == 0 (lane < nsg) and gated the split second-half prefix sums on split, cutting moe_route time by -8.8% to -22.7% and moe_plan time by -14.6% to -35.4% across ntok = 1..8.

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

2.1 Bitwise Parity (std::memcmp == 0) #

• route_parity: K=512 (200 windows) 0 mismatches, K=256+MoCE=0.85 (200 windows) 0 mismatches (PASS).
• route_parity --ab /tmp/argon_r22_moe/moe_ref.metal: 200 trials (moe_route + moe_plan, split = 0, 1, ntok = 1..8): 0 fail(s) ALL PASS.
• expert_parity --ab /tmp/argon_r22_moe/moe_ref.metal: Q2_0/Q2_0(142), IQ2_XS/Q2_0(prep), IQ3_XXS/Q5_1, Q4_K/IQ4_NL across NE = 1, 2, 3, 4, 8 (same) and ntok = 1, 2, 3, 4, 8 (random): 0 fail(s) ALL PASS (memcmp == 0 on plan, h, ye, y).

2.2 Paired Interleaved A/B Benchmarks (expert_parity --ab ... --bench) #

A. Q2_0 Gate/Up (32,2,4) + T_Q2_0R (142) Down (moe_down_142_8_2_4) #

B. IQ2_XS Gate/Up (32,2,4) + T_Q2_0 (prep) Down (moe_down_134_8_2_4) #

C. moe_route & moe_plan Paired A/B (route_parity --ab ... --bench) #