8.6 kB 95 lines
86255f22
Round 12 blitz: qsa_topk 3.85x (cached keys + simdgroup scan + paired emit), gemm_hc_up_reg_32_128_4 + unrolled epilogue (-6.9%), Q6_K hoisted mm_gu3 (-4.3%), 3,375.4 ms GPU kprof (587.5 tok/s GPU)
dawn
⋯
audits/sol_r12_qsa.md
sourcerendered
download raw history
at main Markdown

round 12 track b: qsa top-k and partial flash tiles #

decision #

retain the top-k change: bitwise-identical ids and 2.55–3.85x lower median kernel time across the nine requested random-score workloads. retain the one-line flash partial-tile bound as a small, bitwise-checked improvement: clear early-context benefit, roughly 0–1.4% on the sampled longer contexts, not a general long-context speedup claim. no full-model throughput measurement was run.

patch: /tmp/sol_r12_qsa/sol_r12_qsa.patch (passes git -C vendor/strata apply --check). source and all raw logs: /tmp/sol_r12_qsa/. no tracked vendor files edited; only the three owned paths occur in the patch. no server access or full-model benchmark.

design and hypothesis #

owner: the existing metal top-k/flash kernels, on the per-token prompt/decode path. preserve ascending ids, weighted partial blocks, lowest-cell tie ordering, the no-selection fast path, the >4096-block fallback, and exactly the existing flash arithmetic order for every nonempty tile. falsifier: any baseline/candidate memcmp difference or absent measured benefit. bounds remain fixed: four cached keys per lane, 256 histogram buckets, 32 simdgroup prefix records. no new public engine interface or persistent state.

  • cache four order keys per thread when nb <= 4096; retain device loads otherwise.
  • first simdgroup scans eight descending histogram buckets per lane; its exclusive lane prefix identifies the unique lane crossing need.
  • emit scans the pair (equal, greater) once, then clips the equal prefix at the remaining tie budget. this removes the second scan of already-clipped counts and reduces four threadgroup barriers to two per 1024 blocks. a simdgroup barrier protects each group's prefix reads before its lane 31 reuses the slot on the next iteration; the first threadgroup barrier protects totals before simdgroup 0 overwrites them.
  • stop the flash PV loop at nc, skipping its second all-zero tile for nc <= 32.

address-hoisting follow-up was inspected but not included: rk uses the cell coordinate on its second axis, whereas rv uses its first axis and a different tile extent. copying the lane's rk indices/pointers directly into rv is not justified by the existing coop_k4 probe (which checks contiguous k runs, not cross-tensor cell equivalence). the existing cid/scale tables already share the device-id loads. extra shared row offsets or cross-lane remapping need an independently measured variant; no unmeasured locality rewrite is bundled here.

reproducibility #

machine: Apple M5. build: AppleClang 21.0.0.21000327, release -O2, runtime Metal compilation via the existing mtl runtime. scratch CMake project configured without engine dependencies:

cmake -S metal -B build -G Ninja
cmake --build build --target qsa_prompt_parity qsa_parity

commands below run from any directory; default kernel path is the scratch tree baked into its mtl runtime:

/tmp/sol_r12_qsa/strata/build/qsa_prompt_parity --topk-compare T pos0 /tmp/sol_r12_qsa/baseline_qsa.metal 9
/tmp/sol_r12_qsa/strata/build/qsa_prompt_parity --flash-compare T pos0 /tmp/sol_r12_qsa/baseline_qsa_prompt.metal 1 15
QSA_TOPK=1 /tmp/sol_r12_qsa/strata/build/qsa_parity
/tmp/sol_r12_qsa/strata/build/qsa_prompt_parity

all individual invocations had a five-second hard timeout and completed well under five seconds. timings are alternating-order paired GPU samples, one dispatch per sample, after correctness dispatches. seed 5. top-k default scores are positive random values with the real partial-block tail bias; QSA_TOPK_EDGE=1 selects all-equal, small integer ties, signed zeros and random mixtures. flash uses the existing random int8 caches, half scales, queries, gates and sorted ids. samples were chosen before measurement: top-k T=512/1024/1984 × pos0=2048/8192/30000; flash early/budget/crossing/long cases plus those larger T samples.

initial exploratory logs (topk_*.log, flash_*.log, *_first.log) include clock drift and competing GPU work and are not the confirmation timing set. confirmation logs start confirm_; both active sibling benchmarkers explicitly paused for this window. frequency ramp / nonbenchmark system work still appears in early confirmation samples, so absolute time must not be compared across rows. raw paired observations are retained rather than trimmed. the top-k effect is much larger than this noise; the smallest flash effects are not established improvements.

top-k confirmation #

all nine full-buffer ids comparisons: memcmp == 0. each row is the median of nine samples (ms).

T pos0 baseline candidate speedup
512 2048 2.078375 0.539500 3.85x
512 8192 3.127958 1.026333 3.05x
512 30000 3.450500 1.322958 2.61x
1024 2048 0.969208 0.251750 3.85x
1024 8192 1.897208 0.526208 3.61x
1024 30000 2.870375 1.126083 2.55x
1984 2048 1.861250 0.491125 3.79x
1984 8192 2.652417 0.917042 2.89x
1984 30000 5.543250 2.157125 2.57x

flash confirmation #

all seven full-buffer half-output comparisons: memcmp == 0. fifteen pairs each; columns show median times, median paired percentage reduction, and number of candidate wins. larger-context changes are tiny; the T=1024/pos0=8192 row is neutral.

T pos0 baseline ms candidate ms paired reduction wins
96 0 0.297 0.278 6.14% 15/15
96 1984 1.368 1.351 0.90% 12/15
96 2050 1.151 1.136 1.40% 12/15
96 30000 1.592 1.580 0.57% 11/15
512 2048 5.626 5.594 0.52% 15/15
1024 8192 13.785 13.775 -0.01% 7/15
1984 30000 32.730 32.656 0.25% 11/15

correctness and signed zero #

  • qsa_parity_v1_final.log: QSA_TOPK=1, all stages PASS; 0/6153 ids differ from CPU. the ordinary qsa_parity default selects qsa_topk2, so it was also run but is not evidence for this change.
  • qsa_prompt_parity_final.log: ALL PASS at pos0=0,1984,2050,30000, including fused prep qcurh/output memcmp checks.
  • existing CPU/fp32 reference tests use their established numerical tolerances, not bitwise floating-point parity. exact before/after evidence comes from the new top-k compare and existing flash compare modes; this distinction is intentional.
  • final top-k tie/signed-zero tests at pos0=0,16380,30000 all memcmp PASS, including transition across the four-key cache limit. exploratory mixed-pattern tests covered the nine requested sizes too.
  • flash_edge1_*.log: signed-zero queries and K/V scales, zero int8 caches, four position cases, all memcmp PASS.
  • flash_edge2_final_*.log: zero/signed-zero queries/K scales, cell-alternating +/-1 values at minimum-normal-half scale (0x0400), four positions, all memcmp PASS.
  • flash_generic_*.log: forced COOP_K4=0 generic fill path at pos0=0,2050,30000, all memcmp PASS. this is fallback coverage on M5, not M4 hardware validation.

why skipping the zero tile should not change zero signs for finite model inputs: acc is reset to +0 each chunk; all PV operands are half, so every nonzero exact product is a multiple of 2^-48, far from fp32 underflow. fp32 accumulation therefore cannot manufacture a negative underflow zero; exact cancellation rounds to +0, and signed-zero products added to the initial +0 also give +0. the skipped second tile only adds zero. tot still executes the same fma with fac and acc on both paths. the edge tests check actual output bytes rather than assuming numeric equality treats signed zeros correctly. no claim is made about NaN payload propagation or arbitrary invalid cache scales.

source identity and review #

baseline qsa.metal SHA256: f695672bdff60e73cf4165b50377111d8baed8644fb44a4a9a02c94240870118

baseline qsa_prompt.metal SHA256: 45a3505cba897271e76cb9fca23deb8420b9feba8171d86b19e494bba8f3c9f2

final source SHA256:

  • metal/kernels/qsa.metal: b589ba90b4d8ead2148d8c0846ecda1076f59fd8572e42cca66c5e7058b125de
  • metal/kernels/qsa_prompt.metal: 19f364f8de572ade4f56bafb188a1a87bc28febde23fbb99d34ad3e0bade8c6b
  • metal/tests/qsa_prompt_parity.cpp: e65e20c5467a78b64a07715d6d806c2ddfde6ee1e0c24df8c9a64447bec650a2

visual final review used view_diff for every change and view_code for the selection/emit implementation and compare harness. final edits after timing were comment/line-wrap only in kernels and a corrected edge-case generator in the harness; that generator was rebuilt and rerun. no production arithmetic changed after confirmation.