round 17 track b: prepared-sum shuffles #

delivered #

  • patch: /tmp/sol_r17_shuf/sol_r17_shuf.patch (only metal/kernels/gr.metal, gr2.metal, and the original qsa.metal:399–444).
  • git -C vendor/strata apply --check /tmp/sol_r17_shuf/sol_r17_shuf.patch: exit 0.
  • tracked vendor files were never edited. scratch: /tmp/sol_r17_shuf/strata.
  • baseline repo head: 296b59d2f73c8f04a805df29c2fa97655adace78; the actual working-file snapshot, not head alone, is preserved in /tmp/sol_r17_shuf/baseline. manifest.json hashes all tested snapshots.

changes and correctness argument #

gr_up_f, gr_up3, and qsa_write_prepared now compute even/odd sequential partial sums with simd_shuffle_down, matching the existing gdn.metal pattern. only lanes 0 and 16 publish the sum. their source lanes 0–15 and 16–31 respectively remain inside their half-simdgroup. the token guard is uniform within those halves, including odd token counts 1 and 3. the k iteration and s0/s1 addition order are unchanged. this is not a tree reduction.

qsa_write_prepared has one caller, so its threadgroup-pointer parameter is removed. the kernels keep their host signatures and layouts. gr_norm has a full-unrolled 10-element fast path for n=2560, preserving fma accumulation order and the original generic path. no changes to weight formats, host dispatch, or persistent state.

pipeline-reported static threadgroup bytes (bf16 and q8):

kernel old new barriers per threadgroup old → new
gr_up_f 2560 2048 2 → 1
gr_up3 12800 12288 3 → 2
qsa_merge 1024 0 1 → 0
gr_norm 32 32 1 → 1

for the supplied schedule of 97 gr_up_f + 12 qsa_merge dispatches, this removes 109 dispatch-level barrier sites per window, not merely 109 physical threadgroup executions. the 97 up dispatches alone remove 15,520 threadgroup barrier executions at 160 groups each. gr_up3 is an alternate path and is not added to that schedule count. no full-model trace was collected.

verification #

standalone /tmp/sol_r17_shuf/ab uses the existing mtl runtime, compiling frozen old and new sources separately in safe math mode. deterministic random inputs (seed 19+size; weights seed 31), n=2560, hc=4, lr=320. all output bytes are compared, including all prepared planes and sums, gr_down inject, and gr_norm's mutated residual. gr_up3 receives independent finite synthetic D/Rn/wn/rs, so this isolates that kernel rather than testing a down3→up3 chain.

  • ntok 1,2,3,4,8, both bf16 and q8: all memcmp comparisons pass for gr_norm (apply 0 and 1), gr_down_8_8, gr_up, gr_up3, and qsa_merge.
  • baseline and candidate qsa_parity 3000 with QSA_TOPK=1: all pass.
  • baseline and candidate qsa_prompt_parity: all pass.
  • baseline and candidate prompt_parity: all pass.
  • final rerun of all three parity programs: exit 0; final benchmark matrix: ten exit-0 invocations.
  • every standalone test/benchmark invocation completed below 5 seconds (slowest exploratory benchmark 3.70s, final parity 1.73s). no :8080 access or full-model benchmarks.

an initial qsa_parity 64 run failed identically in baseline and candidate because that context intentionally skips indexer q preparation while this test still compares it. the corrected active-selection command above passes; the original failed logs remain parity0.log and parity3.log.

not covered: other gr_down SG/KS variants, q4/q5 bytewise comparisons, generic n!=2560 norm runtime tests, nonfinite activation inputs, full-model output or throughput. templates compile through the existing library, but compile success is not numerical coverage.

paired microbenchmark evidence #

apple m5; safe math; warm-cache synthetic weights, one untimed correctness dispatch per pipeline then 7 alternating a/b pairs with 12 dispatch repetitions per sample. medians are gpu time per dispatch. there is no server-state manipulation: this is a shared live machine and gpu clocks/contention drift substantially. in particular identical gr_down code and even identical norm code in unroll comparisons sometimes differ, so small gains and outlier regressions must not be read as reliable throughput estimates. norm apply=1 repeatedly mutates its residual within a timed sample; both arms start each sample from the same input. these are microbenchmarks, not a decode result.

final candidate/baseline median-time ratios (lower better; not controlled performance claims):

ntok format norm apply0 norm apply1 down up up3 qsa_merge
1 bf16 0.7004 0.8705 1.0004 0.9692 0.9641 0.9798
1 q8 0.7034 0.8731 1.0007 1.0052 1.1289 0.9917
2 bf16 0.5989 0.9265 1.0191 0.9296 0.9551 0.9936
2 q8 0.7125 0.8819 1.0022 0.9650 0.9882 0.9942
3 bf16 0.6991 0.8673 1.0005 0.9826 0.9469 0.9914
3 q8 0.7101 0.8698 0.8882 0.9556 0.9896 0.9922
4 bf16 0.6973 0.8599 0.9997 1.0046 0.9943 1.0136
4 q8 0.7030 0.8606 0.9984 0.9541 0.9915 0.9919
8 bf16 0.6291 0.7554 1.0027 0.9587 0.9681 0.9196
8 q8 0.6292 0.7581 0.9967 0.9895 1.2389 0.9901

first-pass ab_* logs had norm ratios ~0.61–0.70 (apply0) and ~0.76–0.88 (apply1), consistent with eliminating the dynamic local-array indexing. reduction-only gains were generally small; final q8 up3 includes regressions/outliers, so no dependable gr_up3 speedup is established. the barrier and memory reductions themselves are confirmed by source and pipeline metadata. recommend parent verify latency in its controlled integration run before making any speed claim.

unroll exploration rejected #

unroll_all/ additionally forced full KS and GR_NTOK inner loops and requested cc unroll_count(4); all five bf16 ntok cases remained byte-identical. down ntok8 regressed 11.8% in that exploratory comparison, while other down cases were approximately neutral. unroll_cc/ tested cc alone across all five q8 cases; byte equality passed but timings were inconsistent (including an ntok3 2.10x up outlier). neither variant is in the patch. the generic hc loop stays unchanged rather than hardcoding hc=4. no unsupported claim that loop pragmas help.

reproduction and artifacts #

clang++ -O2 -std=c++20 -Ivendor/strata/metal/src /tmp/sol_r17_shuf/ab.cpp vendor/strata/build-metal/libmtl_rt.a -framework Metal -framework Foundation -o /tmp/sol_r17_shuf/ab
/tmp/sol_r17_shuf/ab 3
/tmp/sol_r17_shuf/ab 3 q8

repeat ntok=1,2,3,4,8; each command has a 4.8s timeout in the tool runner. ab.cpp explicitly references baseline and scratch absolute kernel paths. ab_unroll.cpp compares v1 to the scratch kernel dir; point its latter dir at the preserved unroll_all/ or unroll_cc/ to reproduce those explorations. v1/ and candidate/ equal the final variant. final_* logs preserve the confirmation matrix and parity results; ab_*, unroll_*, cc_*, and parity* logs preserve earlier runs. plan.md records the pre-run contract. final diff and context received a visual review; no extra helper was added because gdn's same arithmetic is inline and the kernel libraries compile independently.