7.7 kB 66 lines
5dc99052
Round 27 blitz: gdn_step_t 5->4 barriers + dt/ssm_a hoist + red2[IB-1] (-3.0..-7.9%), qsa_merge2/qsa_merge float2(m,l) + simd_max (-8..-20%), gemv_tile n_out=512 32,2,4,1 (-3.2..-25.5%) + Q2_0 640->2560 8,2,4,1
dawn
⋯
audits/argon_r27_gdn.md
sourcerendered
download raw history
at main Markdown

Round 27 Track C Audit Report: gdn_step_t (gdn.metal) Barrier Elimination (5 -> 4 Barriers/Token) + dt/ssm_a Hoist + part[8]/red2[IB-1] Shrink #

  • Agent: argon-r27-gdn-and-gr-opt (gemini-4-argon)
  • Patch: /tmp/argon_r27_gdn/argon_r27_gdn.patch (git -C vendor/strata apply --check /tmp/argon_r27_gdn/argon_r27_gdn.patch passes cleanly)
  • Owned files: metal/kernels/gdn.metal, metal/tests/gdn_parity.cpp (vendor/strata/ untouched)

1. Summary of Changes #

  1. Barrier Elimination (5 -> 4 threadgroup_barriers per Token) + part[32] -> part[8] Shrink (metal/kernels/gdn.metal):

    • Before the first barrier of each token, threads 0..127 (sg = 0..3) and 128..255 (sg = 4..7) already hold raw cv in registers: they now write raw qs[tid] = cv and ks[tid - 128] = cv directly to shared memory alongside their simd_sum(cv * cv) write to part[sg], while threads 256..383 (sg = 8..11) write vs[tid - 256] = cv (skipping redundant simd_sum(cv * cv) on tid >= 256 and shrinking threadgroup float part[32] to part[8]) and idle lanes 384 / 416 compute gb[0] / gb[1].
    • After that single barrier, each thread computes rsq = rsqrt(tot_q / float(S) + a.eps / float(S)) and rsk = rsqrt(tot_k / float(S) + a.eps / float(S)) from part[0..7] and normalizes kl = (rsk * ks[qk_idx]) * scale and ql = (rsq * qs[qk_idx]) * scale on the fly before broadcasting across the simdgroup.
    • Because (rsk * ks[qk_idx]) * scale and (rsq * qs[qk_idx]) * scale execute the exact same two IEEE-754 roundings (rs * cv) * scale in the exact same order on the exact same cv value, the result is 100% bitwise identical while eliminating GDN_NTOK 512-thread barriers per dispatch (144 barriers per 4-token verify window across 36 GDN layers).
    • Write-after-read safety across tokens t -> t + 1 holds with three subsequent barriers in token t separating all reads of qs, ks, vs, part, and gb from their next-token writes.

  2. Token-Invariant dt[h] & ssm_a[h] Hoist + red2[IB][S] -> red2[IB - 1][S] Shrink (metal/kernels/gdn.metal):

    • Preloaded dth = tid == 384 ? dt[h] : 0.0f and ssm_ah = tid == 384 ? ssm_a[h] : 0.0f before the state load loop and token loop so their device loads issue in parallel with qkv, hist, conv_w, gamma, and state.
    • Shrunk threadgroup float red2[IB][S] to red2[IB - 1][S] (saving 512 B of threadgroup memory; total threadgroup memory drops 5784 B -> 5176 B for IB = 4): ib == 0 keeps its partial readout o in its register across the barrier and accumulates for (int k = 0; k < IB - 1; ++k) o += red2[k][j], preserving the exact IEEE-754 addition order ((o_0 + o_1) + o_2) + o_3 while eliminating 128 shared stores and 128 shared loads + additions per token.
    • Tested #pragma clang loop unroll(full) on for (int t = 0; t < GDN_NTOK; ++t) and ruled it out: full unrolling of the outer token loop degraded ntok >= 4 performance by +6% .. +16% due to register pressure under [[max_total_threads_per_threadgroup(512)]].

  3. Extended gdn_parity --ab Coverage to ntok = 1..8 (metal/tests/gdn_parity.cpp):

    • Updated --ab bitwise parity and paired A/B benchmark loops in metal/tests/gdn_parity.cpp from ntok = 1..4 to 1..MAXT (1..8), covering all 88 (ntok, commit, plane) combinations (commit = 0, 1, 0x101 .. 0x107; plane = 0, 49152) in < 4.2s total runtime.

2. Verification & Paired A/B Benchmark Results on Apple M5 #

  • Pipeline threadgroup occupancy:
    • gdn_step (IB = 4): maxTotalThreadsPerThreadgroup = 512 (needs 512) for all GDN_NTOK = 1..8.
    • gdn_step8 (IB = 8): maxTotalThreadsPerThreadgroup = 1024 (needs 1024) for all GDN_NTOK = 1..8.
  • Bitwise parity (gdn_parity --ab /tmp/argon_r27_gdn/base_gdn.metal):
    • gdn_step (IB = 4): 88 / 88 PASS (0 fail(s) ALL PASS, memcmp == 0 on y, state, hist across all ntok = 1..8, commit = 0/1/partial, plane = 0/49152).
    • gdn_step8 (IB = 8): 88 / 88 PASS (0 fail(s) ALL PASS).
  • Paired A/B Benchmark (./build-metal/gdn_parity --ab /tmp/argon_r27_gdn/base_gdn.metal --bench, 36 layers rotated = 113 MB state):
ab bench plane     0 ntok 1 commit  base   59.03 us  cand   58.54 us  ratio 0.9885 (x36: 2.125 -> 2.107 ms)
ab bench plane     0 ntok 1 verify  base   35.23 us  cand   34.51 us  ratio 1.0120 (x36: 1.268 -> 1.242 ms)
ab bench plane     0 ntok 2 commit  base   67.51 us  cand   66.34 us  ratio 0.9924 (x36: 2.430 -> 2.388 ms)
ab bench plane     0 ntok 2 verify  base   41.03 us  cand   40.69 us  ratio 0.9915 (x36: 1.477 -> 1.465 ms)
ab bench plane     0 ntok 3 commit  base   67.96 us  cand   73.18 us  ratio 1.0698 (x36: 2.447 -> 2.635 ms)
ab bench plane     0 ntok 3 verify  base   53.70 us  cand   49.92 us  ratio 0.9458 (x36: 1.933 -> 1.797 ms)
ab bench plane     0 ntok 4 commit  base   85.32 us  cand   82.23 us  ratio 0.9813 (x36: 3.071 -> 2.960 ms)
ab bench plane     0 ntok 4 verify  base   66.33 us  cand   61.44 us  ratio 0.9449 (x36: 2.388 -> 2.212 ms)
ab bench plane     0 ntok 5 commit  base   97.72 us  cand   97.38 us  ratio 0.9693 (x36: 3.518 -> 3.506 ms)
ab bench plane     0 ntok 5 verify  base   76.71 us  cand   71.98 us  ratio 0.9729 (x36: 2.762 -> 2.591 ms)
ab bench plane     0 ntok 6 commit  base  105.55 us  cand  107.56 us  ratio 1.0220 (x36: 3.800 -> 3.872 ms)
ab bench plane     0 ntok 6 verify  base   85.64 us  cand   80.99 us  ratio 0.9284 (x36: 3.083 -> 2.916 ms)
ab bench plane     0 ntok 7 commit  base  115.73 us  cand  109.36 us  ratio 0.9450 (x36: 4.166 -> 3.937 ms)
ab bench plane     0 ntok 7 verify  base   92.26 us  cand   83.80 us  ratio 0.9412 (x36: 3.321 -> 3.017 ms)
ab bench plane     0 ntok 8 commit  base  121.96 us  cand  113.98 us  ratio 0.9386 (x36: 4.390 -> 4.103 ms)
ab bench plane     0 ntok 8 verify  base   98.90 us  cand   93.92 us  ratio 0.9496 (x36: 3.561 -> 3.381 ms)
ab bench plane 49152 ntok 1 commit  base   61.16 us  cand   60.90 us  ratio 0.9698 (x36: 2.202 -> 2.193 ms)
ab bench plane 49152 ntok 1 verify  base   39.40 us  cand   38.15 us  ratio 0.9683 (x36: 1.418 -> 1.373 ms)
ab bench plane 49152 ntok 2 commit  base   69.72 us  cand   67.33 us  ratio 0.9847 (x36: 2.510 -> 2.424 ms)
ab bench plane 49152 ntok 2 verify  base   47.07 us  cand   46.84 us  ratio 0.9995 (x36: 1.695 -> 1.686 ms)
ab bench plane 49152 ntok 3 commit  base   77.60 us  cand   75.86 us  ratio 0.9621 (x36: 2.794 -> 2.731 ms)
ab bench plane 49152 ntok 3 verify  base   57.29 us  cand   56.05 us  ratio 0.9837 (x36: 2.062 -> 2.018 ms)
ab bench plane 49152 ntok 4 commit  base   94.69 us  cand   86.28 us  ratio 0.9218 (x36: 3.409 -> 3.106 ms)
ab bench plane 49152 ntok 4 verify  base   67.10 us  cand   64.39 us  ratio 0.9207 (x36: 2.415 -> 2.318 ms)
ab bench plane 49152 ntok 5 commit  base  106.59 us  cand  101.82 us  ratio 0.9709 (x36: 3.837 -> 3.666 ms)
ab bench plane 49152 ntok 5 verify  base   79.16 us  cand   75.86 us  ratio 0.9398 (x36: 2.850 -> 2.731 ms)
ab bench plane 49152 ntok 6 commit  base  115.77 us  cand  109.96 us  ratio 0.9570 (x36: 4.168 -> 3.959 ms)
ab bench plane 49152 ntok 6 verify  base   95.05 us  cand   88.38 us  ratio 0.9473 (x36: 3.422 -> 3.182 ms)
ab bench plane 49152 ntok 7 commit  base  126.66 us  cand  118.77 us  ratio 0.9567 (x36: 4.560 -> 4.276 ms)
ab bench plane 49152 ntok 7 verify  base  102.09 us  cand   96.49 us  ratio 0.9493 (x36: 3.675 -> 3.474 ms)
ab bench plane 49152 ntok 8 commit  base  127.43 us  cand  120.34 us  ratio 0.9517 (x36: 4.587 -> 4.332 ms)
ab bench plane 49152 ntok 8 verify  base  109.12 us  cand  102.96 us  ratio 0.9830 (x36: 3.928 -> 3.707 ms)