round 6 track b — qsa flash #

experiment contract #

owner: prompt qsa kernels and the prompt_chunk dispatch block. preserve output bytes, selected-index semantics above budget, scalar fallback, and existing runtime interfaces. scratch only: /tmp/sol_r6_qsa/strata.

hypothesis: cached half scales plus eight-lane-per-head softmax lowers standalone flash time without changing any fp16 output bit. identity selection avoids topk below budget. scalar qsa_attn still needs ids, so retain its topk dispatch.

evaluator: same deterministic random int8/scales/query/gate input against frozen baseline kernel in /tmp/sol_r6_qsa/baseline; bitwise memcmp is mandatory. boundary cases include 2051/2052 and mixed chunks. standalone flash pairs at t=512,1024,1984; three paired samples per process, no warm duration loop; each process capped below 5 seconds. no full-model execution or port 8080. reject on byte mismatch; performance improvement must recur across sizes, otherwise report inconclusive. maximum three implementation repairs.

delivered change #

• qsa_flash / qsa_attn_mm derive identity ids per token when n_kv <= 2051; longer tokens retain explicit selection, including chunks that cross the threshold.
• flash stages both half4 scales once per selected cell. removes the dead vs store. static threadgroup storage rises from 18,880 to 19,904 bytes; max threads remains 1024 on m5.
• eight lanes per head compute max and exp; shuffle reduction stays inside each eight-lane partition. reuse dead sp[0] for float exponentials; lane zero sums them in the original order. padding heads stay zero.
• host dispatch skips qsa_topk for under-budget flash prompts. deliberate safety deviation: keeps it for !p.flash, because this branch dispatches scalar qsa_attn, which unconditionally reads ids and is outside this track's ownership. unconditional under-budget skipping would break that fallback. only original window.cpp line 3454 changes.
• test adds --flash-compare T pos0 baseline_qsa_prompt.metal and poisons implicit ids. ordinary parity now also exercises mixed-budget windows at pos0=1984 and 2050.

actual runs #

machine: apple m5, safe metal math (runtime default), rng seed 5. baseline is the frozen working tree, not clean upstream: repository head 296b59d2f73c8f04a805df29c2fa97655adace78; exact owned-file hashes in /tmp/sol_r6_qsa/sha256.json, initial status in baseline-status.txt. no tracked vendor file was edited. source snapshots and test executable are under /tmp/sol_r6_qsa/.

compiled test with:

clang++ -std=c++20 -O2 /tmp/sol_r6_qsa/strata/metal/tests/qsa_prompt_parity.cpp -I/tmp/sol_r6_qsa/strata/metal/src -I/tmp/sol_r6_qsa/strata/metal/tests vendor/strata/build-metal/libmtl_rt.a -framework Metal -framework Foundation -o /tmp/sol_r6_qsa/qsa_prompt_parity

the existing runtime static library was reused; candidate msl is runtime-compiled from scratch, selected with STRATA_METAL_KERNELS=/tmp/sol_r6_qsa/strata/metal/kernels. exact invocation per pair:

/tmp/sol_r6_qsa/qsa_prompt_parity --flash-compare T POS /tmp/sol_r6_qsa/baseline/metal/kernels/qsa_prompt.metal

three interleaved paired gpu timings per process after one correctness dispatch each; no sustained warmup. all processes exit 0, all under 1.5 seconds wall, each launch timeout=4.9s.

every listed comparison passed full-output memcmp == 0 against frozen flash baseline, with implicit ids poisoned to -1 in candidate correctness dispatch. timing-only candidate ids were restored for fair baseline pairing.

ordinary parity command (same environment, no args): exit 0, ALL PASS, 0.208 seconds final run. qsa_attn_mm relative l2 vs scalar reference ranged 3.09e-4–3.62e-4; flash ranged 3.72e-4–4.28e-4 across pos0 0/1984/2050/30000. score worst relative difference 2.76e-6. existing fused prep qcurh and flash memcmp both pass. logs: /tmp/sol_r6_qsa/parity-final.log, flash-*.log, confirm-flash-*.log.

decision / limitations #

recommend applying the combined patch: reproducible ~20% standalone under-budget flash reduction, tested byte parity, topk dispatch removed on the flash short-prompt path. topk host savings are not timed here. no end-to-end prefill speed claim; full model and :8080 were not run/touched. m4 / non-k4 execution not available (the branch compiles on m5); scalar fallback host behavior is preserved by retaining topk. no full engine rebuild was run; host edit is a single conditional guard.

patch: /tmp/sol_r6_qsa/sol_r6_qsa.patch; git -C vendor/strata apply --check passes. reviewed complete scratch diff with view_diff and kernel/test context with view_code. host edit verified restricted to original line 3454, inside authorized 3440–3475.