Round 15 Track C Audit: qsa_prep Threadgroup-Memory & Barrier Reduction + route_topk n_exp == 512 Unrolling #

• Agent: argon-qsa-prep-and-route-topk
• Scratch workspace: /tmp/argon_r15_prep/strata
• Patch artifact: /tmp/argon_r15_prep/argon_r15_prep.patch (verified clean with git -C vendor/strata apply --check)
• Target files & line boundaries:
  • metal/kernels/qsa.metal lines 58–187 (tg_sum256, norm_rope, qsa_prep)
  • metal/kernels/prompt.metal lines 16–115 (route_topk)

1. Summary of Changes #

Target 1: qsa_prep / norm_rope Threadgroup Memory (1,088 B -> 320 B) & Barrier Reduction (metal/kernels/qsa.metal) #

1. Shrank threadgroup float row[256] (1,024 B) to threadgroup float row[QROT] (64 floats = 256 B):
   • norm_rope only ever reads row[p] (0..31) and row[p + 32] (32..63) when rope && tid < QROT (QROT = 64).
   • Replaced row[tid] = y; with if (rope) { if (tid < QROT) row[tid] = y; threadgroup_barrier(mem_flags::mem_threadgroup); ... }.
   • Total static threadgroup memory of qsa_prep drops from 1,088 B to 320 B (-768 B, -70.6%) (part[8] = 32 B, row[64] = 256 B, amax[8] = 32 B), eliminating 192 wasted shared-memory stores per threadgroup when rope == true and all 256 shared-memory stores + the RoPE barrier when rope == false (cell == 0 dead key).
2. Eliminated Redundant Trailing Barriers in tg_sum256 and norm_rope (4 -> 2 barriers per single-call threadgroup):
   • Removed the trailing threadgroup_barrier(mem_flags::mem_threadgroup) from tg_sum256 (safe because part[0..7] is distinct from row[0..63] and amax[0..7]).
   • Added bool reuse = false to norm_rope so the trailing threadgroup_barrier(mem_flags::mem_threadgroup) is only executed when reuse == true (the r == 32 && t_self == 0 loop over t < n_pre where part and row can be reused across cell == 0 and (cell & 3) == 3 or subsequent t iterations).
   • Every single-call threadgroup (r < 32 and r == 32, t_self >= 3 — i.e., 67,583 of 67,584 threadgroups at T = 2048) now executes 2 threadgroup_barriers instead of 4 (-50%).
3. Evaluated Per-Token Early-Exit for r >= 28 && r < 32 (qidx heads, line 180):
   • Inspected qsa_prompt.metal (qsa_score_mm lines 190–225) and metal/tests/qsa_parity.cpp: when pos0 + n_tok > 2051 (e.g., a chunk crossing QBUDGET = 2051 such as pos0 = 1984, T = 96), qsa_score_mm executes a matmul2d over all 16 tokens in the tile (tQ sliced at t0 * IH) and scores the tail block dead for every token in the tile before qsa_topk or qsa_flash runs. Leaving qidx[t] uninitialized (or NaN-poisoned) for tokens t with pos0 + t + 1 <= 2051 when pos0 + n_tok > 2051 would feed uninitialized memory into qsa_score_mm and alter scores[t].
   • Kept if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; as-is so all stages and parity tests remain 100% bitwise identical.

Target 2: route_topk n_exp == 512 Unrolled Fast Path (metal/kernels/prompt.metal) #

• Added an if (n_exp == 512) fast path with #pragma clang loop unroll(full) for:
  1. Initial logit load + per-lane max reduction (v[i] = l[lane + 32 * i]; mx = max(mx, v[i]);),
  2. Softmax exp(v[i] - mx) + per-lane sum reduction,
  3. Normalization v[i] *= rcp + isnan(v[i]) guard, and
  4. The MoCE rank == p_prot && moce_scale > 1.0f scaling loop.
• Eliminates all idx < n_exp (lane + 32 * i < 512) branch predicates on the K = 512 path while preserving the exact floating-point operation sequence and fallback else path for n_exp != 512 (K = 256, 320).

2. Verification & Bitwise Parity Results #

All standalone parity suites and a dedicated A/B bitwise (memcmp == 0) & paired-dispatch benchmark (ab_bench) were run on Apple M5:

• qsa_parity: ALL PASS (q rel L2 5.22e-08, indexer q 5.60e-08, pooled keys 4.36e-08, top-k 0/6153 diff, attention + gate 2.44e-07)
• qsa_prompt_parity: ALL PASS (qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS)
• prompt_parity: ALL PASS (route_topk 300 tokens: id mismatches 0, worst weight rel err 2.28e-07)
• comb_fold_parity: ALL PASS (T = 1, 7, 64, 1024 bitwise PASS)
• ab_bench (memcmp == 0 between baseline and patched kernels across all outputs qcur, qidx, kq, vq, ksc, vsc, tail, dead, pooled, ids, w): 100% BITWISE PASS (fails = 0).

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

qsa_prep (12-layer back-to-back command buffers, 9 alternating pairs, static_tg_mem: 1,088 B -> 320 B) #

route_topk (48-layer back-to-back command buffers, 9 alternating pairs, n_exp = 512) #