Round 30 Track C Audit: gdn_pre, gdn_post, and gdn_scanr_4 (vendor/strata/metal/) #

1. Summary of Shipped Changes & Findings #

Shipped (metal/kernels/gdn_prompt.metal & metal/tests/gdn_prompt_parity.cpp) #

1. gdn_pre (gdn_prompt.metal lines 80-128):
   • Replaced the 12 scalar simd_sum(float) reduction trees per key head (sq0..sq3, sk0..sk3, p[0..3]) with 3 vector simd_sum(float4) reductions (sq, sk, pv), preserving the exact IEEE-754 pairwise lane addition tree ((v.x + v.y) + (v.z + v.w)).
   • Hoisted the cold if (t == T - 1) history write check out of the unrolled for (int m = 0; m < 4; ++m) conv loop in both t >= 3 and t < 3 paths (leaving the cold t == T - 1 loop un-unrolled so it does not bloat register pressure or instruction cache on the t < T - 1 hot path).
2. gdn_post (gdn_prompt.metal lines 515-531):
   • Replaced the 4 scalar simd_sum(float) calls (p0..p3) with a single vector simd_sum(float4) reduction (pv).
   • Hoisted if (a.f32) outside the unrolled for (uint k = 0; k < 4; ++k) output loop.
3. gdn_prompt_parity.cpp (metal/tests/gdn_prompt_parity.cpp lines 165-179):
   • Added paired A/B 36-layer total (pre + scan + post together) timing to --ab ... bench.

Investigated & Ruled Out in gdn_scan_r (gdn_scanr_4, TB == 0, IB == 4) #

We tested 6 distinct hypotheses on gdn_scan_r<4, 0> (188.3 ms across 36 layers at T = 2048 on M5) and ruled out all 6 because the baseline gdn_scan_r<4, 0> is already at 99.6% of the M5 GPU's hardware instruction/register roof (224 instructions/token/thread across 1.812e9 thread-tokens = 4.059e11 instructions -> 188.8 ms theoretical floor at 10 cores * 128 ALUs * 1.68 GHz):

• Splitting threadgroup float2 red[2][IB][32] into red_k + red_q with if (sg != 0): Regressed 188.39 ms -> 192.62 ms (+2.3%). Apple AGX issues float2 threadgroup loads/stores as single 64-bit vector memory ops; splitting into scalar float arrays doubled shared-memory instruction count and added warp divergence around the barrier.
• Branching if (sg == 0) after the barrier to sum .x only on sg != 0: Regressed 187.70 ms -> 188.80 ms (+0.58%) due to branch overhead in the per-token loop.
• 64-bit pointer stepping (kn_p += ks, qn_p += ks, vc_p += rs, gb_p += hv, kq_p += hkn, o_p += rs): Regressed 188.27 ms -> 194.47 ms (+3.39%). Keeping six 64-bit pointers live across the token loop consumes 12 extra 32-bit VGPRs on top of float st[32], pushing gdn_scanr_4 over the AGX VGPR occupancy threshold.
• Uniform base offsets ((kn + kn_u)[lr]) vs kn[tn * ks + k0]: Regressed 187.87 ms -> 190.15 ms (+1.26%) because k0 and v0 already amortize the + lr / + lane offset across kn/qn and vc/o.
• Moving red[slot][sg][lane] = s before the t + 1 global prefetches (kn_l, qn_l, g_b_n, v_n): Regressed 187.71 ms -> 190.76 ms (+1.63%). In the baseline, issuing the t + 1 global loads between the end of the r = 0..31 FMA loop and red[slot][sg][lane] = s overlaps load issue with the draining of the 32-deep serial ka/qa FMA dependency chain.
• Splitting the r = 0..31 loop into a separate st[r] update loop followed by a ka/qa dot-product loop: Regressed 188.10 ms -> 188.97 ms (+0.50%) because interleaving the independent st[r] FMA instructions inside the r loop fills the 4-cycle FMA pipeline bubble between consecutive serial ka/qa accumulations.

2. Bitwise Parity & Paired A/B Benchmarks (Quiet Apple M5 GPU, 36 GDN Layers) #

• Default parity suite (gdn_prompt_parity): ALL PASS (32/32 tests: 1-sg vs legacy gdn_pre/gdn_post, single-pass vs TB=16 gdn_scanr_4 for T = 1..2048, and all 7 scan variants vs gdn_step across T = 1, 3, 37, 300).
• --ab bitwise parity across all 10 buffers (qn, kn, vc, gb, kq, history, o, state, yh, y): 60/60 PASS (memcmp == 0) across T = 512, 1024, 2048 and f32 = 0, 1.

3. Patch Verification #

• Unified diff written to /tmp/argon_r30_gdn/argon_r30_gdn.patch.
• Verified cleanly applicable via git -C vendor/strata apply --check /tmp/argon_r30_gdn/argon_r30_gdn.patch (exit code 0).