Round 5 Track C Report: Direct FP16 qcurh in qsa_prep + PromptGemm Dead Buffer Removal & Layer-Local Aliasing #

• Agent: argon-qsa-prep-f16-and-scratch (opencode-go/gemini-4-argon-contributor)
• Patch Artifact: /tmp/argon_r5_scratch/argon_r5_scratch.patch (verified clean via git -C vendor/strata apply --check /tmp/argon_r5_scratch/argon_r5_scratch.patch and clang++ -fsyntax-only on window.cpp)
• Test Binaries: /tmp/argon_r5_scratch/{qsa_prompt_parity,qsa_parity,arena_budget_test,qsa_prep_bench,need_bytes_check}

1. Summary of Changes #

Target 1: Fuse qsa_q_f16 into qsa_prep when p.flash is enabled (metal/kernels/qsa.metal + metal/src/engine/window.cpp) #

• Added constant bool QSA_Q_F16_FC [[function_constant(25)]] (qsa_qf16) to metal/kernels/qsa.metal:
  • When qsa_qf16 is true, qsa_prep writes ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); directly to buffer 10 instead of float (qcur[(t * QH + r) * QD + tid] = o;).
• In Window::prompt_chunk (metal/src/engine/window.cpp lines 3424–3444):
  • When p.flash is true, qsa_prep is dispatched with qsa_c().b(25, true) (pipeline key qsa_k() + "#qf16") and p.qcurh bound directly at buffer 10 (p.flash ? p.qcurh : p.qcur), eliminating the separate qsa_q_f16 kernel dispatch.
  • Across the 12 QSA layers per 2048-token chunk, this eliminates 12 qsa_q_f16 dispatches, 906 MB of DRAM reads/writes (12 * (50.33 MB FP32 write + 50.33 MB FP32 read - 25.17 MB FP16 write)), and makes p.qcur (50.33 MB) completely unused when p.flash == true.

Target 2: Remove Dead Prompt Scratch Allocations (p.XQ, p.yg, p.qcur) + Alias Disjoint GDN vs QSA Prompt Scratch Buffers (metal/src/engine/window.cpp lines 3080–3189) #

• Dead buffer removal (-323.96 MB at MT=2048, -171.39 MB at MT=1024):
  1. p.XQ: Changed from unconditional F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384) (223.35 MB at MT=2048) to p.flash ? F32(16) : F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384).
  2. p.yg: Changed from F32(MT * 6144) (50.33 MB at MT=2048, never written because gdn_post runs with ga.f32 = 0) to F32(16).
  3. p.qcur: Changed from F32(MT * 24 * 256) (50.33 MB at MT=2048) to p.flash ? F32(16) : p.z.
• Disjoint GDN (l % 4 != 3) vs QSA (l % 4 == 3) layer-local scratch aliasing (-250.59 MB additional at MT=2048, -165.15 MB additional at MT=1024):
  • Because every GDN and QSA scratch buffer is produced and consumed strictly inside its own if (l % 4 != 3) ... else ... branch before gemm(..., p.Wout, p.bo, ...) writes p.bo, QSA's layer-local buffers cleanly alias GDN's allocations:
    • p.qkv sized to F32(MT * 12288) (+16.78 MB over 10240) and shared with p.qfull = p.qkv (-83.89 MB net at MT=2048).
    • p.kcur = p.gqn (512 <= 2048 floats/token, -4.19 MB) and p.vcur = p.gkn (512 <= 2048 floats/token, -4.19 MB).
    • p.qcurh = p.ygh (both F16(MT * 6144), -25.17 MB).
    • p.gvc sized to F32(MT * std::max<size_t>(6144, (size_t) MB)) and shared with p.scores = p.gvc (-50.33 MB).
    • p.qids = p.go (2051 <= 6144 words/token, -16.81 MB).
    • p.Wqkv sized to F16((size_t) 12288 * N) (+10.49 MB over 10240 * N) and shared with p.Wq = p.Wqkv (-52.43 MB net).
    • p.Wk = p.Wz (512 * N <= 6144 * N, -2.62 MB), p.Wv = p.Wal (Wal sized to F16(512 * N), -0.25 MB net), p.Wiq = p.Wbe (Wbe sized to F16(512 * N), -0.25 MB net).
  • Both make() in Window::prompt_gemm_init and Window::prompt_gemm_need_bytes() are updated in lockstep.

2. Verification & Benchmark Results on Apple M5 #

2.1 Bitwise Parity (qsa_prompt_parity, qsa_parity, arena_budget_test) #

• qsa_prompt_parity:
  • qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS (100% bitwise identical qcurh and qsa_flash output).
  • All existing qsa_attn_mm, qsa_flash, and qsa_score_mm checks PASS.
• qsa_parity: ALL PASS (prep, score, top-k, attention + gate).
• arena_budget_test: all passed (46/46 checks).
• need_bytes_check: Verified across 480 combinations of C in {64, 512, 1024, 2048, 4096}, flash in {false, true}, stream in {false, true}, pred2 in {false, true}, bo_alias in {false, true}, ple_on in {false, true}, max_cells in {4096, 32768, 131072} that need == prompt_gemm_need_bytes() and every aliased buffer's capacity is >= its required QSA byte size.

2.2 Microbenchmark (qsa_prep_bench on Apple M5) #

2.3 prompt_gemm_need_bytes() Memory Reduction (max_cells = 32768, stream = true, ple_on = true) #