Round 19 Track C: gemv.metal (gemv / gemv_cat / gemv_catk) K-Split & Inner-Loop Unrolling + router.metal Unrolling #

• Agent: argon-gemv-and-router-unroll (opencode-go/gemini-4-argon-contributor)
• Target files: metal/kernels/gemv.metal, metal/kernels/router.metal
• Unified patch: /tmp/argon_r19_gemv/argon_r19_gemv.patch (verified clean via git -C vendor/strata apply --check)
• Standalone A/B harness: /tmp/argon_r19_gemv/ab_bench.cpp (< 0.25s minimal shader subset compile per library)

1. Key Findings & Root-Cause Analysis #

1.1 gemv, gemv_cat, gemv_catk, gemv_pair, and gemv_x (metal/kernels/gemv.metal) #

1. Why for (int t = 0; t < NTOK; ++t) and for (int r = 0; r < NR; ++r) Were Not Unrolled by Default:
   • NTOK ([[function_constant(0)]]) is a Metal function constant (not a C++ template parameter). Without #pragma clang loop unroll(full), the Apple Metal compiler does not unroll for (int t = 0; t < NTOK; ++t) when the body contains X<F>::load and unit_combine<UNIT>, nor does it unroll for (int r = 0; r < NR; ++r) when NR > 1 (acc[r][t], q[r], sc[r]).
   • Consequently, acc[r][t] was indexed by a runtime loop variable t (and r), preventing scalar replacement of aggregates (SROA) into registers and causing thread-local stack indexing and loop-branch overhead on every unit u and every K-split reduction.
2. K-Split (KS > 1) Reduction & Bounds Check Hoisting (gemv and gemv_catk):
   • Adding #pragma clang loop unroll(full) to for (int r = 0; r < NR; ++r), for (int t = 0; t < NTOK; ++t), and for (int k = 0; k < KS; ++k) in the KS > 1 reduction emits KS immediate-offset threadgroup loads (red[(k * RPT + lrow0 + r) * 8 + t]) directly into registers.
   • Hoisting if (row0 + r < int(a.n_out)) outside the t and k reduction loops in gemv and gemv_catk avoids redundant threadgroup loads on tail rows.
3. gemv_pair Init Bound (t < NTOK vs t < 8):
   • Changed for (int t = 0; t < 8; ++t) ag[t] = au[t] = 0.0f; to #pragma clang loop unroll(full) for (int t = 0; t < NTOK; ++t) ag[t] = au[t] = 0.0f; and unrolled all t < NTOK loops in gemv_pair and gemv_x.
4. Tested & Rejected Sub-Variant (if (NR == 1) branch inside u loop):
   • Explicitly branching on if (NR == 1) inside u with scalar q0[UNIT] / sc0 produced identical codegen once #pragma clang loop unroll(full) was present on for (int r = 0; r < NR; ++r), except for a +3.62% scheduling wobble on IQ3_S 6144->2560 (16_1_4_1, ntok=2). Kept the clean #pragma clang loop unroll(full) loop without the redundant if (NR == 1) branch.

1.2 router_route (metal/kernels/router.metal) #

1. Why Guarding threadgroup_barrier with if (fmt<F>::TAB > 1) Regresses router_route:
   • We A/B tested if (fmt<F>::TAB > 1) { load_tab<F>(tab, tid, nth); threadgroup_barrier(mem_flags::mem_threadgroup); } in router_route across 21 interleaved pairs x 128 reps.
   • Removing the initial threadgroup_barrier(mem_flags::mem_threadgroup) at line 131 consistently regressed router_route_30 (BF16) by +1.22% .. +1.57% (24.06 -> 24.35 us at ntok=1, 39.23 -> 39.71 us at ntok=4) and router_route_0 (F32) by +3.84% .. +4.35% (42.04 -> 43.79 us at ntok=1).
   • Root cause: router_route launches a 512-thread (16-simdgroup) threadgroup where all 16 simdgroups read the same xp / s16 vector across u = lane; u < U; u += 32 and then synchronize at the device-fence / atomic counter tail (ngx = 32 threadgroups). The initial threadgroup_barrier(mem_flags::mem_threadgroup) aligns all 16 simdgroups within each threadgroup at u = 0 so simdgroup 0..15 hit the L1 cache in lockstep for xp[u] and finish together at the tail barrier. Without it, simdgroups drift out of phase on xp loads.
   • Therefore, per Rule 2.3 ("Only keep changes that are 100% memcmp == 0 and neutral-to-faster (zero regressions)"), we kept the initial synchronization in router_route and kept #pragma clang loop unroll(full) on the three for (int t = 0; t < NTOK; ++t) loops in router_route (which is 100% bitwise identical and neutral-to-faster: -0.28% on BF16 ntok=2, -0.33% on F32 ntok=1).

2. Standalone A/B Benchmark Results (Apple M5) #

All measurements use interleaved A/B pairs (48..128 reps per command buffer, 7..15 pairs, both pipelines pre-warmed) and verify 100% bitwise memcmp == 0 across all ntok = 1..8.

2.1 gemv Production Shapes & Tiles (T_Q8_0, T_IQ3_S, T_Q2_0) #

2.2 gemv_cat, gemv_catk, gemv_pair, and router_route #

3. Parity Verification #

• gemv_parity: ALL PASS across all 13 formats (F32, BF16, Q8_0, Q4_0, Q5_0, IQ4_NL, IQ4_XS, Q2_0, IQ3_S, Q3_K, Q4_K, Q5_K, Q6_K), n_in in {2560, 640, 10240}, ntok in {1, 3, 8}, all 20 tiles (window == one token: bitwise).
• fuse_parity: ALL PASS (gemv_cat ntok=1..8 bitwise, router_route BF16 & F32 ntok=1..8 2 runs each bitwise, gr_norm_comb & gr_norm_comb_4s ntok=1..8 bitwise).
• route_parity: ALL PASS (K=512 200 windows 0 mismatches, K=256 + MoCE=0.85 200 windows 0 mismatches).