Round 6 Track C Audit Report: Register-Cached rv in gr_norm_h* + 1-Simdgroup-Per-Head Barrier-Free gdn_pre & gdn_post #

• Agent: argon-gr-regs-and-gdn-post
• Device: Apple M5 (10-core GPU, 24 GB unified memory)
• Patch: /tmp/argon_r6_gr_gdn/argon_r6_gr_gdn.patch (git -C vendor/strata apply --check clean)
• Modified files (within owned boundaries):
  • metal/kernels/prompt.metal (lines 171-319)
  • metal/kernels/gdn_prompt.metal (gdn_pre and gdn_post)
  • metal/src/engine/window.cpp (lines 3407 and 3412 only, inside 3402-3414 boundary)
  • metal/tests/comb_fold_parity.cpp
  • metal/tests/gdn_prompt_parity.cpp

1. Target 1: Register-Cached rv in gr_norm_h, gr_norm_h4, gr_norm_h_moe, gr_norm_h_moec (metal/kernels/prompt.metal) #

Design & Exact Arithmetic Preservation #

• In gr_norm_h (1 stream per threadgroup of 256 threads, n = 2560 = 10 * 256), Pass 1 now caches the updated residual stream values rv[i] = v (float rv[10], i = 0..9) in registers and Pass 2 computes o[d] = half(rv[i] * w[d] * rs) directly from registers without re-reading r[d] from device memory.
• In gr_norm_h4, gr_norm_h_moe, and gr_norm_h_moec (hc = 4 streams per threadgroup of 256 threads), Pass 1 caches rv[c][i] = v (float rv[4][10]) in registers and passes rv to gr_norm_h_out, which writes o[d] = half(rv[c][i] * w[d] * rs) directly from registers with zero device re-reads of R (-83.89 MB per call at T=2048, -7.89 GB across 96 calls at T=1983).

Bitwise Parity & M5 GPU Benchmark (comb_fold_parity + ple_moe_parity + prompt_parity) #

• Bitwise Parity (memcmp == 0): 100% PASS across T in {1, 7, 64, 1024, 2048} on both R' and xnh for gr_norm_h_moec, gr_norm_h_moe, gr_norm_h4 (apply 0), and gr_norm_h4 (apply 1).
• M5 GPU Benchmark (48 reps per command buffer, best of 3):

(Note: In standalone microbenchmarks where a single threadgroup's 40 KB slice of R[t] remains hot in L1/SLC between Pass 1 and Pass 2, the savings reflect eliminated load instructions + L1/SLC traffic; in full-window execution with concurrent dispatches and cache pressure, eliminating the second R read avoids SLC eviction re-reads from DRAM.)

2. Target 2: 1-Simdgroup-Per-Head Barrier-Free gdn_pre & gdn_post (metal/kernels/gdn_prompt.metal + window.cpp) #

Design & Exact Arithmetic Preservation #

1. Barrier-Free 1-Simdgroup-Per-Head gdn_post (grid = (12, T), 128 threads = 4 heads/TG):

   • Each simdgroup (32 lanes) handles one head h = int(tg.x) * 4 + int(sg) (with automatic fallback to h = sg == 0 ? int(tg.x) : hv when launched with legacy grid.x == 48).
   • Lane lane (0..31) loads ov0 = o[base], ov1 = o[base + 32], ov2 = o[base + 64], ov3 = o[base + 96] and computes: p0 = simd_sum(ov0 * ov0), p1 = simd_sum(ov1 * ov1), p2 = simd_sum(ov2 * ov2), p3 = simd_sum(ov3 * ov3), tot = (p0 + p1) + (p2 + p3).
   • Because each simd_sum reduces the exact same 32 elements (0..31, 32..63, 64..95, 96..127) as the 4 simdgroups in the original kernel and sums them with identical (p0 + p1) + (p2 + p3) parenthesization, the result is 100% bitwise identical (memcmp == 0) while using 0 bytes of threadgroup memory, 0 threadgroup_barriers, and 4x fewer threadgroups (24,576 vs 98,304 at T=2048).

2. Barrier-Free 1-Simdgroup-Per-Head gdn_pre (grid = (8, T), 256 threads = 8 heads/TG):

   • When launched with grid.x == (h_k + h_v) / 8 = 8 (2 key-head threadgroups + 6 value-head threadgroups per token, vs 40 threadgroups in the legacy grid; legacy grid.x == 40 is also supported):
     • For x < h_k / 8 (2 TGs = 16 simdgroups for 16 key heads): each simdgroup (32 lanes) owns one key head kh = x * 8 + sg and computes both cv_q[0..3] and cv_k[0..3] (4 channels per lane) in registers.
     • L2 norms tot_q = (sq0 + sq1) + (sq2 + sq3) and tot_k = (sk0 + sk1) + (sk2 + sk3) and the dot product kq = (p[0] + p[1]) + (p[2] + p[3]) are reduced within the single simdgroup in the exact same floating-point association order as part[0..3] and part[4..7] in the original kernel.
     • For x >= h_k / 8 (6 TGs = 48 simdgroups for 48 value heads): each simdgroup (32 lanes) owns one value head h = (x - 2) * 8 + sg (4 channels per lane).
   • Eliminates threadgroup float part[8] and threadgroup float prod[128] (544 B threadgroup memory -> 0 B), eliminates all 3 threadgroup_barriers (3 -> 0), and cuts threadgroups by 5x (8 * T vs 40 * T) while remaining 100% bitwise identical (memcmp == 0) on all 6 outputs (qn, kn, vc, gb, kq, hist_out).

Bitwise Parity & M5 GPU Benchmark (gdn_prompt_parity) #

• Bitwise Parity (memcmp == 0): Verified 100% bitwise identical (PASS) between packed 1-simdgroup-per-head (8, T) / (12, T) and legacy (40, T) / (48, T) across T in {1, 3, 37, 300}, and all 28 gdn_step recurrence parity checks pass (history exact).
• M5 GPU Benchmark (T = 2048 and T = 1024):