Round 14 Track C: gdn_scanr_4 hk-Grouped Head Locality & gr_norm_h* / moe_comb Optimization Report #

1. Summary & Deliverables #

• Scratch Workspace: /tmp/argon_r14_gdn_gr/strata
• Unified Patch: /tmp/argon_r14_gdn_gr/argon_r14_gdn_gr.patch (verified clean via git -C vendor/strata apply --check)
• Files Modified (within Track C ownership boundaries):
  • metal/kernels/gdn_prompt.metal (gdn_scan_r lines 225–227)
  • metal/kernels/prompt.metal (gr_norm_h*, moe_comb_slots, moe_comb, gr_apply_moec lines 170–440)
  • metal/tests/comb_fold_parity.cpp (grvec paired A/B 5-sample median reporting)
• Bitwise Parity: 100% bitwise identical (memcmp == 0) across all gdn_prompt_parity (T = 1, 3, 16, 17, 32, 37, 300, 512, 1024, 2048, f32 = 0, 1) and comb_fold_parity (T = 1, 7, 64, 512, 1024, 2048) checks.

2. Target 1: gdn_scan_r (gdn_scanr_4) hk-Grouped Head Scheduling (metal/kernels/gdn_prompt.metal) #

2.1 What Worked (Kept) #

• hk-Grouped Head Remapping (hv == 3 * hkn): In Qwen3.8-Flash-Next, hv = 48 value heads share hkn = 16 query/key heads (hk = h % 16). Previously, h = int(tg.x) / 4 scheduled tg.x = 0..63 across hk = 0..15, separating the 3 value heads (12 threadgroups: h = hk, hk + 16, hk + 32) that read the exact same qn and kn rows (2 MB per hk at T = 2048) by 64 threadgroups. Remapping:

  const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k);
  const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn);
  const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx;
  

  schedules all 12 threadgroups sharing a given hk consecutively (tg.x = 0..11 share hk = 0, tg.x = 12..23 share hk = 1, etc.), keeping qn and kn hot in L1/L2 cache across the 3 value heads. Because each (h, j) threadgroup writes exclusively to its own slice of state and o, this is 100% bitwise identical (memcmp == 0) by construction.

2.2 What Was Ruled Out (Discarded) #

• sg == 0-Only qa (s.y) Reduction in gdn_scan_r (TB == 0): Adding an if (sg == 0) branch after the per-token threadgroup_barrier so that sg = 1..3 only sum red[slot][0..3][lane].x while sg == 0 sums float2 (s.x, s.y) regressed gdn_scanr_4 by +1.0% to +1.4% (49.45 ms -> 50.03 ms at T = 512, 98.50 ms -> 99.50 ms at T = 1024). On Apple M5, uniform float2 shared-memory loads (ld.tg.v2.f32) across all 4 simdgroups are faster than divergent simdgroup control flow inside the tight T-token recurrence loop.

2.3 Paired A/B Timings (gdn_prompt_parity --ab, 36 GDN layers, f32 = 0, median of 3 alternating pairs) #

3. Target 2: gr_norm_h_moec, gr_apply_moec, and gr_norm_h4 (metal/kernels/prompt.metal lines 170–440) #

3.1 What Worked (Kept) #

1. SIMD-Compacted 32-bit Row Offsets in moe_comb_slots / moe_comb: Previously, moe_comb looped over all k = 0..topk-1, checking if (rk[k] >= 0) and computing 64-bit ulong(rk[k]) * n + d 100 times per thread (10 d-iterations $ imes$ 10 topk). In moe_comb_slots, simdgroup 0 now compacts the valid routed slots (r >= 0) in order using simd_prefix_exclusive_sum(ushort(valid)) and precomputes 32-bit row offsets roff[prefix] = uint(r) * n (max_rows * 2560 <= 52.4M < 2^32) and weights wk[prefix] = w, storing the valid count cnt = simd_sum(ushort(valid)) in roff[16]. moe_comb then runs a branchless loop for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); with 32-bit indexing, preserving the exact summation order (100% memcmp == 0).
2. Hoisted !a.apply Fast Path + n == 2560 Full Loop Unrolling in gr_norm_h4, gr_norm_h_moec, and gr_norm_h_stream_out:
   • Separating if (!a.apply) in gr_norm_h4 eliminates zero-initializing bv[10], evaluating exp(-inj) per stream, and branching on if (a.apply) inside the inner loop, speeding up gr_norm_h4 (apply 0) by -2.4% to -3.5%.
   • Specializing if (n == 2560) with #pragma clang loop unroll(full) on the i = 0..9 (d = tid + i * 256) stream loops in gr_norm_h4, gr_norm_h_moec, and gr_norm_h_stream_out (plus #pragma clang loop unroll(full) on the 8-way part[i] reductions) keeps bv[0..9] and rv[0..9] in registers rather than dynamically indexed thread stack memory.

3.2 What Was Ruled Out (Discarded) #

• Swapping the gr_norm_h_moec moe_comb Loop Order (k outer, i = 0..9 inner) or Full Unrolling of the moe_comb Outer Loop: Accumulating bv[0..9] across k = 0..cnt-1 in registers (or #pragma clang loop unroll(full) on the i = 0..9 loop calling moe_comb) keeps all 10 accumulators live simultaneously across the YE global loads, increasing register pressure and hurting occupancy (+1.5% to +3.1% slower at T = 1024, 2048). Keeping the outer for (uint d = tid, i = 0; d < n; d += 256, ++i) loop over moe_comb un-unrolled while unrolling the stream loops (rc, w_norm, xnh) gave the best performance at all T.

3.3 Paired A/B Timings (comb_fold_parity T grvec, 4 reps per sample, 5 alternating samples) #