7.1 kB 70 lines
35f84cbe
Round 18 blitz: per-stream gr_norm_comb 5.09x (-80%), xprep float4 (-46.8%), moe_combine_sh (-49%), Q6_K moe_gu NE>=3 (-25%), qsa_topk2/qsa_score (-5..9%) -> 40.03 tok/s decode & 445.79 tok/s 10.18K prefill
dawn
⋯
audits/argon_r18_qsa2.md
sourcerendered
download raw history
at main Markdown

Round 18 Track C: qsa_topk2 Register-Cached Keys & Barrier Reduction + qsa_score Query Hoisting & Unrolling #

1. Summary of Changes #

Target 1: metal/kernels/qsa2.metal (qsa_topk2, lines 121–238) #

  1. Redundant threadgroup_barrier Removal (20 -> 13 radix-pass barriers, -7 barriers per qsa_topk2 dispatch = -84 barriers per 12-layer decode window):
    • Hoisted the initial if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); once before the pass = 0..3 loop.
    • Inside each radix pass, immediately after cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed);, cleared the bin for the next pass via if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed);. Because two threadgroup_barriers (the sgs barrier and sel barrier) execute between that store and the next pass's atomic_fetch_add_explicit, hist is guaranteed zeroed before pass + 1 accumulates into hist.
    • Removed the redundant per-pass start barrier and trailing barrier after need = sel[1]; (sel[0] and sel[1] are only read at the end of the pass and are not written again until after the histogram and sgs barriers of pass + 1).
  2. Register-Cached Order Keys (nb <= 4096, n_kv <= 16,384) with Scoped Live Range & Uniform Early break:
    • Scoped uint keys[4]; inside if (nb <= 4096) (covering n_kv <= 16,384) so the nb > 4096 (pos0 = 30000) path keeps zero keys register live-range across its barriers.
    • Preloaded keys[0..3] via #pragma clang loop unroll(full) before pass = 0..3 and reused keys[i] across all 4 radix passes with an early if (b >= nb) break; so nb <= 1024 (pos0 = 3000) exits after i = 0 (1 iteration instead of 4) and nb = 2049 (pos0 = 8192) exits after i = 1..2.
    • Evaluated ekeys[4] in the emit pass (E4 vs M1_break): Adding a second uint ekeys[4] register array across the two emit-pass threadgroup_barriers increased peak register pressure in the 1024-thread threadgroup and measured slower than M1_break at pos0 = 3000 (6.52 us in M1_break vs 6.61 us in E4) and pos0 = 8192 (7.81 us in M1_break vs 8.07 us in E4). Retaining M1_break (without ekeys[4]) delivers the fastest qsa_topk2 across both pos0 = 3000 (-3.17% to -4.95%) and pos0 = 8192 (-5.60% to -5.87%) while also improving pos0 = 30000 (-2.19% to -2.34%).

Target 2: metal/kernels/qsa.metal (qsa_score, lines 199–222) #

  • Hoisted device const float* qp = qidx + (t * IH) * ID + lane;, preloaded float q[IH][4] and float k[4] with #pragma clang loop unroll(full), and unrolled the h < IH and i < 4 FMA reduction loops so all 16 query loads (qp[h * ID + 32 * i]) and 4 key loads (key[lane + 32 * i]) use compile-time immediate offsets.
  • Evaluated threadgroup float q_tg[512] (vC) vs Register Hoisting (vB): Staging qidx (2 KiB) in threadgroup float q_tg[512] required adding a 256-thread cooperative load and threadgroup_barrier(mem_flags::mem_threadgroup) to a barrier-free kernel and increased static_tg_mem from 0 B to 2,048 B, which was slower than register hoisting (vB) across every tested configuration (-0.79%..-5.19% for vC vs -5.47%..-8.94% for vB). Shipped vB (0 B shared memory, 0 barriers).

2. Verification & Bitwise Parity #

  1. Standalone Parity Suites (/tmp/argon_r18_qsa2/strata/build-metal):
    • qsa_parity (P = 3000 and P = 8192): ALL PASS (0 of 6153 ids differ, score rel L2 5.88e-08, < 0.6s).
    • qsa_prompt_parity: ALL PASS (pos0 = 0, 1984, 2050, 30000, < 1.1s).
  2. Paired A/B Bitwise Parity (/tmp/argon_r18_qsa2/ab_bench):
    • 100% bitwise memcmp == 0 on both scores and ids (including 0xa5-poisoned padding/unused entries) across all 56 boundary & context combinations pos0 in {0, 1, 3, 2048, 2050, 2051, 2052, 3000, 4092, 8192, 16380, 16384, 20000, 30000} x ntok in {1, 2, 3, 4}, plus adversarial tie-breaking and signed-zero (+0.0f / -0.0f / all-equal / small-integer ties) tests across pos0 in {2052, 3000, 8192, 16380, 16384, 30000}.

3. Paired Interleaved A/B Benchmark Results (Apple M5) #

Measured on a quiet Apple M5 GPU (11 interleaved pairs x 32 dispatches per command buffer after 0.30s P-state warmup; median per-dispatch time in microseconds):

pos0 ntok Stage Baseline (us) Candidate (us) Speedup Delta Bitwise memcmp
3000 1 qsa_score 2.50 2.36 1.06x -5.57% PASS (0)
3000 1 qsa_topk2 6.86 6.52 1.05x -4.95% PASS (0)
3000 1 combined 9.52 9.05 1.05x -4.87% PASS (0)
3000 2 qsa_score 3.64 3.43 1.06x -5.90% PASS (0)
3000 2 qsa_topk2 6.80 6.52 1.04x -4.08% PASS (0)
3000 2 combined 10.66 10.14 1.05x -4.90% PASS (0)
3000 4 qsa_score 6.23 5.76 1.08x -7.48% PASS (0)
3000 4 qsa_topk2 6.77 6.56 1.03x -3.17% PASS (0)
3000 4 combined 13.27 12.48 1.06x -5.99% PASS (0)
8192 1 qsa_score 4.71 4.42 1.07x -6.33% PASS (0)
8192 1 qsa_topk2 8.30 7.84 1.06x -5.60% PASS (0)
8192 1 combined 13.16 12.40 1.06x -5.76% PASS (0)
8192 2 qsa_score 8.02 7.48 1.07x -6.80% PASS (0)
8192 2 qsa_topk2 8.29 7.81 1.06x -5.72% PASS (0)
8192 2 combined 16.47 15.43 1.07x -6.32% PASS (0)
8192 4 qsa_score 14.92 13.78 1.08x -7.61% PASS (0)
8192 4 qsa_topk2 8.30 7.81 1.06x -5.87% PASS (0)
8192 4 combined 23.39 21.77 1.07x -6.92% PASS (0)
30000 1 qsa_score 25.48 25.45 1.00x -0.12% PASS (0)
30000 1 qsa_topk2 14.01 13.71 1.02x -2.19% PASS (0)
30000 1 combined 39.96 39.65 1.01x -0.78% PASS (0)
30000 2 qsa_score 26.86 25.27 1.06x -5.94% PASS (0)
30000 2 qsa_topk2 14.06 13.73 1.02x -2.32% PASS (0)
30000 2 combined 41.60 39.51 1.05x -5.04% PASS (0)
30000 4 qsa_score 52.52 47.82 1.10x -8.94% PASS (0)
30000 4 qsa_topk2 14.12 13.79 1.02x -2.34% PASS (0)
30000 4 combined 67.35 62.22 1.08x -7.61% PASS (0)

4. Artifacts #

  • Unified patch: /tmp/argon_r18_qsa2/argon_r18_qsa2.patch (verified clean with git -C vendor/strata apply --check /tmp/argon_r18_qsa2/argon_r18_qsa2.patch)
  • Paired A/B benchmark source & binary: /tmp/argon_r18_qsa2/ab_bench.cpp, /tmp/argon_r18_qsa2/ab_bench
  • Scratch workspace & build tree: /tmp/argon_r18_qsa2/strata/build-metal