opus_r9_qsa — Round 9 Track A: qsa_flash threadgroup memory 19,904 B -> 7,616 B #

patch: /tmp/opus_r9_qsa/opus_r9_qsa.patch (git -C vendor/strata apply --check clean against the current dirty working tree). touches only metal/kernels/qsa_prompt.metal (qsa_flash only; qsa_attn_mm untouched) and metal/tests/qsa_prompt_parity.cpp. scratch: /tmp/opus_r9_qsa/strata, build /tmp/opus_r9_qsa/build, baseline kernel /tmp/opus_r9_qsa/baseline_qsa_prompt.metal, experiment kernel dirs /tmp/opus_r9_qsa/v*/.

result (Apple M5, paired alternating, 12 dispatches per command buffer = 12 QSA layers, pos0 0) #

(second repeat: 24.76->21.17, 90.79->75.62, 330.84->272.44.) single-dispatch, pos0 30000: T512 17.19->13.75, T1024 34.61->27.32, T1984 67.57->52.95 ms. T=96 pos0 0 is ~+5% (0.545->0.570 ms, tiny grid, two extra barriers per chunk); T=96 pos0 30000 is -17%.

bitwise: memcmp PASS for T = 96, 512, 1024, 1984 x pos0 = 0, 30000, plus (96,1984), (96,2050), (1,0), (64,1987), (200,1900). full qsa_prompt_parity (default mode, incl. qsa_prep fused memcmp through qsa_flash): ALL PASS. every process < 3 s.

what changed #

• sp[4][16*QCH] (16 KiB) -> sp[2][HPG*QCH] (6 KiB); S partials stored with the matmul2d tensor store into a 12-row (dextents(QCH, HPG)) view, in two rounds: sg2/sg3 store, the softmax's own lanes (8 lanes x 8 cells per head) read s2+s3 into registers, barrier, sg0/sg1 store, lanes compute ((s0+s1) + s23) * 0.0625f -> same IEEE order as before, kept in registers for both the max and the exp loops (no re-read / re-add).
• p (fp16 P, 2 KiB) aliases sp[1], which is dead after the combine (one extra barrier before the P writes); sg3 still zeroes P's 4 padding rows each chunk (S stores clobber them), fac[12..15] zeroed once before the chunk loop.
• the serial lsum reads float4 at a time, same sequential add order.
• total 7,616 B -> 4 TGs/core by the 32 KiB rule (was 1).
• test: --flash-compare T pos0 baseline [layers pairs] now times layers back-to-back dispatches per sample and prints median+min (old 5-arg form unchanged: layers 1, pairs 3).

what the brief got wrong (measured, worth knowing) #

• Variant A as written is ~40% slower (T1984: 27.4 -> 38.5 ms/layer) while bitwise-identical. the cost is not the barrier and not occupancy: element-wise access of the destination cooperative tensor (s[i] + get_multidimensional_index) is expensive. isolation: original code + one extra barrier = no change (27.43); original code with manual scalar stores instead of s.store = 37.5; float4 manual stores (dest layout is runs of 4 cells, coop_layout 16_64_32) = 35.1. the tensor store into a 12-row view instead of 16 is ~2.5% faster by itself (26.76).
• variant B (single 3 KiB buffer, 4 serialized store rounds, 6,592 B) is bitwise but slower than the 2-buffer version (T1984 x12: 285.9 vs 279.4) - the serialized stores cost more than the extra occupancy buys.
• attribution of the 2-buffer register-combine version at 9,664 B (before p aliasing): same code padded back to 19,904 B gives 302.6 ms x12, so of the 330 -> 279 win about half is the code (no 4-way re-read, 12-row store) and half is occupancy; p aliasing (9,664 -> 7,616 B, 3 -> 4 TGs/core) took 279 -> 272.