Round 25 Track C (argon-r25-moemm-down-and-plan): mm_down3 W4<T_Q2_0> 64-Block Stepping + mm_gu3 Coalesced half4 HS Store #

1. Summary of Shipped Changes #

1. W4<T_Q2_0>::run_q<OFF>(device const uchar* b, uint kl4) & mm_down3 64-Element Block Stepping (metal/kernels/moe_mm.metal):

   • In mm_down3 (F == T_Q2_0 && TK == 32, coop_k4), two consecutive TK = 32 steps (k0 and k0 + 32) sit inside the same 64-element, 18-byte Q2_0 block (QK = 64, BB = 18).
   • Previously, each k0 += 32 iteration computed boff = uint(k0 >> 6) * 18 and j_in = (k0 + kl) & 63, and W8<T_Q2_0> called W4<T_Q2_0>::run(rdw, j_in) and W4<T_Q2_0>::run(rdw, j_in + 16), where the second call had to recompute + ((k >> 6) * 18) and 2 + ((k & 63) >> 2) because the compiler could not prove (j_in + 16) < 64.
   • By stepping k0 += 64, w0_l += 18 and introducing W4<T_Q2_0>::run_q<OFF>(device const uchar* b, uint kl4) reading (b + kl4)[OFF] with compile-time constant immediate byte offsets OFF in {2, 6, 10, 14} and unsigned kl4 = uint(kl) >> 2 (eliminating 32-to-64-bit sign extension sext, (k >> 6) * 18, and (k & 63) while keeping the minimal 1-W4 VGPR live range across the 8-row unroll), mm_down3 speeds up by +2.48% (-121.4 us/layer, -5.83 ms across 48 layers) at T = 1024 and +1.85% (-146.1 us/layer, -7.01 ms across 48 layers) at T = 2048, with +0.15% at T = 512.

2. mm_gu3 half4 Coalesced HS Store Pointer Hoisting + Format Expansion (metal/kernels/moe_mm.metal):

   • Hoisted threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]); and device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f; with stride hs_step = uint(RP) * a.n_ff out of the unrolled m store loop in mm_gu3.
   • Enabled the coalesced half4 HS store path for T_IQ3_XXS, T_IQ2_XXS, and T_Q2_0 in addition to T_IQ2_XS (+0.09% on IQ2_XS, +0.33% on IQ3_XXS, +0.91% on IQ2_XXS, +0.60% on Q2_0 at T = 1024).

3. moe_mm_parity Paired A/B Harness (metal/tests/moe_mm_parity.cpp):

   • Added candidate warmup and alternating (i & 1) A/B ordering to gather_bench plan timing (7 paired iterations).
   • Updated --ab isolated kernel benchmarking (--gu-bench / --dn-bench) to run 19 interleaved pairs with external-contention spike filtering (<= 1.10 * best_pair).

2. Measured A/B Results on Apple M5 #

All benchmarks run with STRATA_METAL_KERNELS=/tmp/argon_r25_moemm/strata/metal/kernels (A = candidate) and --ab /tmp/argon_r25_moemm/base_kernels (B = baseline):

3. Ruled-Out Hypotheses (Measured Regressions / Neutral) #

• Sharing d and loading q0, q1 upfront in W8<T_Q2_0>: -0.84% at T=1024, -1.39% at T=2048 (increases VGPR live ranges across the 8-row unrolled loop).
• SWAR 0x40100401u multiply in W4<T_Q2_0>: -0.65% at T=1024, -2.58% at T=2048 (adds 4 fp32 fadd -1.0f ops on the FPU pipe).
• SWAR 0x00204081u char4 sign negation in w4_signed: -10.11% at T=1024 (serializes float4 conversion behind sign4_iq2xs + 32-bit integer multiplies).
• Unrolling k0 by 2 (k0 += 64) in mm_gu3: -1.22% at T=1024, -3.42% at T=512 (mm_gu3 already has 2 i loops and 2 op.run calls per k0 step; doubling it to 4 spills VGPRs).
• Adding #pragma clang loop unroll(full) on acc / accg zero-init loops: -0.10% to -0.40% vs leaving MPP's cooperative tensor loop default.

4. Artifact Paths #

• Clean unified diff (git -C vendor/strata apply --check verified): /tmp/argon_r25_moemm/argon_r25_moemm.patch
• Audit report: audits/argon_r25_moemm.md