round 21 track b — qsa2 #

experiment contract #

owner: qsa2 attention and merge kernels; standalone parity evaluator. scratch-only implementation; no server or full-model runs.

hypothesis: fixed-trip unrolling and unconditional full eight-cell v loads reduce attention cost; shared merge weights and the existing prepared-plane writer reduce redundant work. keep chunk layout, sequential fma/reduction order, selection, interfaces and defaults unchanged. reject on any part/xp bit mismatch. compare baseline/candidate on identical seeded inputs in paired rounds; >3% consistent speedup is useful, smaller differences are noise. 40 position/window cases include partial chunks and the 2051-cell budget boundary. short runs only, no port 8080. measure six workloads, two confirmation runs; stop on correctness failure or after two setup repairs.

baseline: snapshot of current dirty vendor tree, saved in /tmp/sol_r21_qsa2; patch relative to those exact starting files, not vendor HEAD.

implementation and review #

  • unroll query staging, head/dimension dot loops, accumulator init and stores; hoist query lane pointer.
  • split v accumulation into unconditional unrolled eight-cell blocks and scalar tail, preserving sequential fma order.
  • merge caches maxima/weights in mk[128]/ek[128], computes denominator once, then uses qsa_write_prepared from preceding qsa.metal. all barriers remain uniform. threadgroup storage is 1028 bytes versus the old 1024, not a claimed memory reduction; prepared-plane stage itself becomes barrier-free.
  • QSA_ATTN=2 qsa_parity --bitwise REF-qsa2.metal selects matching 65-chunk/four-group dispatch and allocation; absent that environment variable preserves original v1 mode. baseline/candidate share the unchanged qsa.metal dependency. separate part/xp failure diagnostics; alternate ab/ba timing order.
  • reviewed full diff with view_diff, then test setup and merge context with view_code. bytes from qsa_topk2 onward are unchanged. no tracked vendor source was edited.

build and reproduction #

vendor base HEAD: 296b59d2f73c8f04a805df29c2fa97655adace78, with pre-existing dirty changes. baseline file sha256:

  • qsa2.metal: f8587c2fa8016cc695e101a8c38cb52b535766491c57748e6e6feea9b385b7e1
  • qsa_parity.cpp: ea7a66b94683790e582afed0f47752bc2d2167973e7b9e84e0939cde11b75c3d

scratch: /tmp/sol_r21_qsa2/strata; originals: /tmp/sol_r21_qsa2/qsa2.ref.metal and qsa.ref.metal. scratch git index holds starting files; working files hold candidate. no code modifications during any measured run.

cmake -S metal -B build-metal -G Ninja -DCMAKE_BUILD_TYPE=Release   -DMETAL_CPP_DIR=/tmp/sol_r21_qsa2/strata/third_party/metal-cpp   -DSTRATA_GGML_DIR=/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/llama.cpp
cmake --build build-metal --target qsa_parity strata_window_cli -j 4
QSA_ATTN=2 /tmp/sol_r21_qsa2/strata/build-metal/qsa_parity --bitwise /tmp/sol_r21_qsa2/qsa2.ref.metal
/tmp/sol_r21_qsa2/strata/build-metal/qsa_parity --bitwise /tmp/sol_r21_qsa2/qsa.ref.metal

both build targets passed; cli was built only, never invoked. logs: /tmp/sol_r21_qsa2/build.log, rebuild.log. device: apple m5. seeded random generator 41, 32768 kv cells, ntok 1/2/4/8; pos0 0,1,63,64,512,2047,2048,2050,8192,30000. compiler runtime default fast_math=false. live server was not touched; external server contention remains a timing caveat.

measured results #

v2 runs 1 and 2: 40/40 cases have exact memcmp == 0 on both full part and full xp buffers, including zeroed inactive regions. final formatting-only evaluator rebuild reran v2: 40/40 again. original v1 mode also passed 40/40. each invocation exited 0 in 2.00s, 1.39s, 1.61s, and 1.21s respectively, all under 5s.

paired kernel times below are run 2 medians, microseconds, 9 alternating paired rounds x 32 dispatches after 50ms warm-up. run 1 independently showed the same improvements.

pos0 ntok attn baseline / candidate attn reduction merge baseline / candidate merge reduction
512 1 37.14 / 22.07 40.6% 6.61 / 6.06 8.3%
512 4 82.67 / 50.49 38.9% 10.83 / 9.25 14.6%
2048 1 82.11 / 50.65 38.3% 14.27 / 12.40 13.1%
2048 4 309.16 / 192.27 37.8% 55.30 / 50.73 8.3%
8192 1 82.21 / 51.04 37.9% 14.30 / 12.45 12.9%
8192 4 309.24 / 195.33 36.8% 55.28 / 50.70 8.3%

logs: /tmp/sol_r21_qsa2/bitwise-v2-run1.log, bitwise-v2-run2.log, bitwise-v2-final.log, bitwise-v1.log. final run had a startup outlier at 512/1 attention (68.18/34.13us); remaining workloads agree closely with earlier runs. do not interpret that outlier as additional speedup. no formal confidence intervals were collected.

decision #

accept as a validated local kernel improvement: attention time drops 36.8–40.6%, merge 8.3–14.6% in the confirmation table, with exact outputs on all requested cases. no end-to-end speedup claimed, and v2 is still slower than the existing v1 attention path in these standalone workloads; no default change recommended.

patch: /tmp/sol_r21_qsa2/sol_r21_qsa2.patch. git -C vendor/strata apply --check passed, exit 0. patch changes only the two owned files and only qsa_attn2/qsa_merge2 in the kernel file (original lines 1–120; added lines shift subsequent numbering).