Round 36 Track C: qsa_prep Single-Barrier norm_rope & w == 128 tg_sum + qsa_score_mm Loop Unrolling #

Summary & Decision #

• Owner: argon-r36-qsa-prep-score-and-gdn
• Owned files: metal/kernels/qsa.metal, metal/kernels/qsa2.metal, metal/kernels/qsa_prompt.metal, metal/tests/qsa_parity.cpp, metal/tests/qsa_prompt_parity.cpp
• Kernel decision: No-ship on kernel changes (qsa.metal, qsa2.metal, qsa_prompt.metal unchanged); ship paired A/B harness improvements (--ab in qsa_parity.cpp and qsa_prompt_parity.cpp).
  • All tested qsa_prep and qsa_score_mm kernel variants achieved 100% IEEE-754 bitwise parity (memcmp == 0) across every output buffer (qcur, qidx, kq, vq, ksc, vsc, tail, dead, pooled, scores, ids, part, xp), including QSA_MROPE=1 and QSA_Q_F16_FC=true.
  • However, on quiet Apple M5 (10-core GPU) in warmed 15-pair alternating A/B benchmarks, the current working-tree qsa_prep and qsa_score_mm kernels are already at their M5 optimum:
    1. Single-barrier norm_rope (if (rope && tid < QROT) row[tid] = x; before tg_sum256 + partner scale (rs * row[tid ^ 32]) * g[tid ^ 32]): Bitwise identical (memcmp == 0), but slightly slower across decode and prompt shapes (ratio 1.0006–1.0136, +0.5% to +1.4% at T = 1, 1984, 2048) because holding x live across tg_sum256 + issuing a second uncoalesced/divergent device const float* g load (g[tid ^ 32]) for tid < 64 outweighs the threadgroup barrier savings on M5.
    2. tg_sum256 #pragma clang loop unroll(full) + nsg = (w <= 128 ? 4 : 8): Bitwise identical (memcmp == 0), neutral within ±0.2% at T = 1984, 2048 (961.05 -> 962.95 us, ratio 1.0010 at P=0, T=1984; 1000.15 -> 1005.97 us, ratio 0.9994 at P=3000, T=1984; 1039.76 -> 1044.30 us, ratio 1.0004 at P=3000, T=2048), because Metal Clang already unrolls nsg = 8 when constant-folded and qsa_prep is memory-bandwidth bound on qfull / kcur / vcur / kq / vq.
    3. qsa_score_mm #pragma clang loop unroll(full) + qp pre-offsetting (qsa_prompt.metal): Bitwise identical (memcmp == 0), neutral (0.1788 -> 0.1770 ms at pos0=0, 0.4654 -> 0.4699 ms at pos0=2051, 1.3466 -> 1.3326 ms at pos0=8192, 4.4671 -> 4.4764 ms at pos0=30000 for T = 1984), confirming Round 26 (audits/sol_r26_qsa.md) that qsa_score_mm is dominated by the 64x64x128 matmul2d cooperative tensor op and DRAM reads/writes.

Measured Quiet-GPU A/B Results #

1. qsa_prep Variants (qsa_parity --paired / --ab, 15 alternating pairs of 8x/32x dispatches) #

2. qsa_score_mm Loop Unrolling + qp Pre-offsetting (qsa_prompt_parity --mm-compare score, 15 alternating pairs) #

Harness Additions Shipped in Patch #

• metal/tests/qsa_parity.cpp:
  • Added --ab REF_DIR [P] mode that runs full-buffer bitwise parity (memcmp == 0) and warmed alternating paired timings across decode nt = 1, 2, 4, 8 and prompt nt = 1984, 2048 in a single < 6s invocation.
  • Switched --paired timing loops to use cx.run({ds[v][i]}, repeats) with cx.warm(...) and removes per-turn CPU memset of GPU buffers inside the timing loop.
• metal/tests/qsa_prompt_parity.cpp:
  • Added --ab T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs mode testing qsa_prep, qsa_score_mm, and combined (qsa_prep + qsa_score_mm) with full 10-buffer bitwise memcmp verification (supporting QSA_MROPE=1 and QSA_QF16=1) and 12-layer-per-command-buffer paired A/B timings.
  • Added cx.warm(...) before --mm-compare paired timing loops.

Verification Artifacts #

• Unified patch: /tmp/argon_r36_qsa/argon_r36_qsa.patch
• git -C vendor/strata apply --check /tmp/argon_r36_qsa/argon_r36_qsa.patch: exit 0 (clean).
• Baseline kernel snapshots preserved in /tmp/argon_r36_qsa/base_kernels/.