[jschmied]

What this changes

Before, each column threadgroup of simd_qmm's prompt tiles redid X's 64-input group sums (for the affine bias) and, on the half path, X's conversion to half. Both depend only on X, not on the column. Now a prepass, tf_sq_prep, writes them once per call. Each row's group sums use the tiles' own tree (((c0+c1)+(c2+c3))+((c4+c5)+(c6+c7)), each c a left-to-right fma sum of 8), so they're the same floats. With the half variant, it also writes X as half, rows interleaved in pairs, so a tile reads two uint4 per row and step instead of eight floats.

• tf_sq_mmh (the half tiles, n past 6,144) reads both. The prepass runs once after each norm, beside tf_sq_fits. The fp32 fallback, for tiles whose inputs don't fit half, still reads X itself.

• tf_sq_mmx is tf_sq_mmp (the paired-chunk tiles) reading the sums. The paired body is one inline function shared by both kernels. Prompt rows of 64 or more take it after a per-shape check (summedExact, run once per shape on the chip) has shown its bits equal the fp32 tiles'.

• The bits are main's. Prompt prefill is faster.

• Issue: Refs Qwen3.8-27B on M1-M4: prompt GEMM behind MLX's rate, and plain decode behind 0.6.6 #2089 (the bit-equal prepass: v72 for mmh, v8 for mmp, in the bench series).

• Stacked on: main (223da76). Commit: b3692aa. CHANGELOG line, for whoever merges: - simd_qmm on Metal: a prepass writes X's group sums and half operands once per call; prompt prefill on the 27B is about 4 % faster on chips without tensor units, with the same bits (#2089).

• Overlap (checked with gh on all open PRs, then git merge-tree against this head): 12 open PRs touch zig/src/families/qwen27/quant_gpu.zig (train 1113121437: #2851 #2765 #2883 #2889 #2891 #2892 #2878 (CI only) #3202, qwen27 Metal (M5): fused norm and RoPE write K and V into the cache and Q into qa directly (#2940 step 2b) #2960, qwen27 Metal: one q27_qk_norm_rope dispatch for Q and K head norm and RoPE (#2940 step 2a) #2955, qwen27 Metal: K and V read in place from the mixed projection, no unstack copies (#2940 step 1) #2945, qwen27 Metal (M5): out_proj group sums written by the GDN post and attention merge kernels (#2931) #2944, qwen27 Metal (M1-M4): projections no longer barrier internally, and XS sums compile out before tensor units (#2900) #2913, Metal M5: prompt lane projections at every width take a chunk's 16-row blocks in one dispatch #2851, qwen27 metal: a prompt chunk's last layer runs q, o, the MLP and the final norm over the read row's 128-row block #2826, qwen27 metal: qk projections, flat glue and Metal 4 verify passes, each bit-identical and switchable #2342, Metal: serve the dense Qwen3.8-27B straight from a GGUF file (Q4_K_M, Q8_0) (#1920) #2078, Implement GLP control vectors with a fused Metal kernel for Qwen3.8-27B #1835) or also quant_rows_test.zig (Qwen3.5-9B prompt rows in one product in Zig engine #1897). train 1113121437: #2851 #2765 #2883 #2889 #2891 #2892 #2878 (CI only) #3202, Metal M5: prompt lane projections at every width take a chunk's 16-row blocks in one dispatch #2851, qwen27 metal: a prompt chunk's last layer runs q, o, the MLP and the final norm over the read row's 128-row block #2826 and Implement GLP control vectors with a fused Metal kernel for Qwen3.8-27B #1835 merge with this head cleanly; the others conflict, most in quant_gpu.zig. No other open PR touches simd_qmm.metal, simd_qmm.zig or the two contract record files. Whichever lands second rebases.

Checks on this head

• Head commit: b3692aa
• Lean checker, python3 tools/zig/lean_check.py, on main 223da76 and on this head: base 341, head 341. lean_new.py: 0 new. The 341 are inherited from main.
• Tests run, with the command and pass, fail and skip counts:
  • zig build test --summary all: 844/846 steps succeeded, 13,204/13,308 tests passed, 103 skipped, 1 failed. The failure is dashboard_test ("dashboard page polls the public stats snapshot without private details"). It fails the same way on main 223da76, and nothing in it is touched here.
  • zig build test-qwen27: 78/78 steps, 185/185 tests passed.
  • zig build native -Dcpu=apple_m1: 285/285 steps.
  • zig fmt --check on the changed files: clean. git diff --check 223da7659 HEAD: clean. Every changed file is 535 lines or fewer.
  • New device test (quant_rows_test.zig): "chunk pairs on the prepass's sums keep simd_qmm's bits at an odd row count; without the prepass the bits turn". It uses n 512, k 256 and 65 rows (odd, so the prepass's pad row is exercised). The output is bit-equal to the fp32 tiles. The control leaves the sums at zero, and then the bits differ, so the test fails if the prepass is skipped. The half test now runs tf_sq_prep_h into its own buffer and stays bit-equal.
• CI: not yet run on the fork's head.
• Not run, and why: CUDA and HIP (the change is Metal-only); chips other than the M3 Pro (see Receipt).

Receipt

• Environment:

  • Builds: main 223da76 against this branch b3692aa. Timings ran on the branch's served code before its last amend, which changed only source comments (lean).
  • Toolchain: Zig 0.17.0 and Apple metal 32023.921.
  • Machine: Apple M3 Pro (14-core GPU, no tensor units), 36 GB, macOS 27.2 (26B123), iogpu.wired_limit_mb 29517.
  • Checkpoint: TensorFold/Qwen3.8-27B-MLX-4bit. Drafts: the mtp-linked checkpoint (copies) and DFlash2.
  • Built with zig build native -Dcpu=apple_m1; served with tensorfold-native serve.
  • The Mac was in daily use, with 15-18 % memory free while the 27B was served.

• Paths and platforms the change reaches, which I ran and which I could not:

  • Reached: qwen27's prompt chunks of 64 rows or more on chips without tensor units. That covers the 8-chunk half tiles (gate_up, q, q | k | v, z | b | a, and the head over all rows) and the paired-chunk tiles (down, o). Decode rounds, which have fewer than 64 rows, and every other family keep main's path.
  • M5: not reached. qwen27 runs simd_qmm's raw tiles only where !device.tensorUnits().
  • Ran: M3 Pro only. Not run: M1, M2, M4 (reached there, unrun) and M5 (not reached).

• Kernel time, tf-sq-bench at gate_up's shape, 128 rows (fresh process each, md5 of the output equal to main's):

  kernel	main	branch
  half tiles (mmh)	12.20 ms	11.53 ms
  down, paired	6.610 ms	6.476 ms
  o, paired	2.356 ms	2.327 ms

  The prepass itself takes about 0.02 ms. Main's fp32 tiles (mma) run 12.69 ms on the same shape.

• Exactness:

  • Served greedy tokens: one token_sha256 per prompt across every A-B-B-A run below, on main and on the branch (a9b54923… at 3,695 tokens, 71545ae9… at 14,805).
  • bench_concurrent.py --prompts long.json --tokens 128 --reps 1 --temperatures 0,0.7 --levels 1,3 --alone --serial, on four distinct ~1.2K-token code prompts, served --parallel 3 --context 2048 --prompt-cache-gib 0 with mtp-linked copy drafts:
    • Alone: 1/1 equal at one stream and 3/3 at three.
    • Serial: 3/3 equal against "draft": false at each level and temperature.
    • The reply hashes are identical between main and the branch.
  • DFlash2 drafts, --drafter qwen38-27b-dflash2, one stream (the drafter at --parallel 2 or more isn't admitted beside the 27B on this Mac):
    • Alone and serial are equal at temperatures 0 and 0.7, with about 3.0 tokens per chunk.
    • The reply hashes are identical between main and the branch.
  • Resume, tf-qwen27-reuse (turn 1 of 1,500/5,000 tokens, turn 2 of 2,600/7,003; drafted and plain):
    • Every resumed prompt state and reply is equal to the fresh one ("state equal" on all 12 lines).
    • The tool still reports 12 failures on main and on the branch alike, with the same verdicts on both. The cause is memory pressure on this Mac, which sheds the kept states, so turn 2 resumes at 1,400/4,900 or 0 instead of the expected point.
    • The tool's 16 GiB prompt cache isn't admitted here, so it ran with 2 GiB on both builds (a local edit, not in the branch).
  • bench_cache.py (1,536-token prompt, 3 repeats): the replies are equal across repeats and between builds (token_sha256 91df979a…). It found no cache reuse on this Mac (cached 0) on either build.

• Main against this branch, prompt speed, same machine and checkpoint. Served prefill at one stream, greedy, --no-drafts --no-thinking --prompt-cache-gib 0 --context 20480, a fresh server per run, cold, A-B-B-A over two rounds:

  prompt	main	branch	gain
  3,695 tokens	51.76 s	49.49 s	+4.6 %
  14,805 tokens	216.34 s	208.00 s	+4.0 %

  Per run, prefill s:

  • 3,695 tokens: main 52.27, 51.75, 51.62, 51.77; branch 49.49, 49.43, 49.53, 49.49.
  • 14,805 tokens: main 218.81, 216.16, 216.31, 216.37; branch 208.13, 207.86, 207.77, 208.81.

  Every branch run is faster than every main run, and each round alone gains 3.9-5.2 %. 32K and 64K weren't run, because this Mac doesn't admit them beside the 27B.

• Decode speed before and after:

  • One stream, --no-drafts, 256 greedy tokens, a fresh server per run, A-B-B-A over two rounds. Main decodes in 30.48, 31.57, 31.74 and 31.68 s; the branch in 30.91, 31.82, 31.75 and 31.67 s. The medians are 31.62 s and 31.71 s, 8.1 tok/s on each. All 8 replies have token_sha256 c157e66f….
  • One and three streams, paired: bench_concurrent.py --prompts long.json --tokens 256 --reps 1 --temperatures 0 --levels 1,3, served --parallel 3 --context 2048 --prompt-cache-gib 0 with mtp-linked copy drafts, a fresh server per run, A-B-B-A over one round. Per-stream decode, main then branch:
    • 1 stream: main 7.79 and 7.93 tok/s; branch 8.20 and 8.36.
    • 3 streams: main 5.17 and 5.06 tok/s; branch 5.22 and 5.21. Three streams' per-stream rate includes the other streams' prompt passes.
    • Earlier unpaired single runs (the --alone --serial exactness runs above) read 7.7 against 8.0 tok/s for the branch. The paired runs don't repeat that.
  • Decode rounds don't reach the change (fewer than 64 rows).

• Licence of any new weights or drafter, or "n/a": n/a, no new weights.

Kernel contract, only when a GPU kernel is new or changed

• Kernels registered with the contract, and the kernel each one replaces:
  • tf_sq_mmx (the "metal simd_qmm pairs prepass sums q4 g64" entry) replaces simd_qmm's fp32 tiles and is gated bit-equal to them.
    • It runs the served shapes {64, 5120, 1024} and {9, 512, 512}, plus the shape matrix.
    • Its tiles are 16 rows (at least 9), 128 rows (step 32, at least 96) and 512 rows (step 64, at least 64).
  • tf_sq_prep_s and tf_sq_prep_h are the prepass. The existing half entry now runs tf_sq_prep_h before tf_sq_mmh and stays bit-equal.
  • Gate 5 records tf_sq_mmx_16_4_2_2 in simd_pairs_records.zig (256 threads, 16,384 threadgroup bytes; M1 and M2 take 224 threads, which it adapts to). It also records tf_sq_prep_s and tf_sq_prep_h (32 threads each, 0 bytes; M1 and M2 1,024).
• A planted control for each new kernel, and each one fails as planted:
  • "control: metal simd_qmm pairs, the prepass's sums left at zero": red at shapes, from byte 0 at 64x5120x1024.
  • The half control ("every tile marked as fitting half") stays red as before.
  • A prepass that writes wrong sums or halves turns one of these red.
• Each backend's SUMMARY line and exit status, tf-kernel-contract --only simd_qmm:
  • This head: SUMMARY 6 kernels: 6 pass gates 1-4, 0 fail, 0 known findings; 2 controls: 2 red as planted, 0 missed; 0 limit problems; gates 45211 ms, limits 64 ms, exit 0.
  • Main 223da76: SUMMARY 5 kernels: 5 pass gates 1-4, 0 fail, 0 known findings; 1 controls: 1 red as planted, 0 missed; 0 limit problems, exit 0.
  • The full zig build kernel-contract refuses before any gate with UnrecordedRotaryCompiler on main as on this head, so that's inherited.
• Chips I could not run: M1, M2, M4, M5 (M5 doesn't reach these kernels in qwen27).

Checklist

• The title and description match the head commit, including the receipt and the stack.
• New tests fail before the change, pass after it, and skip cleanly without their dependency. The device test and the contract entry each carry a control that fails as planted; the device test skips without a Metal device.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.
• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. Drafted and concurrent: --alone --serial above, with copy drafts and DFlash2. Resumed: every resumed state is equal to the fresh one; the tool's resume-point check fails on main and on the branch alike under this Mac's memory pressure (above).
• No precision traded for speed. The bits don't change: the prepass's sums use the tiles' own tree, and the half operands are the same conversions the tiles made.
• Decode and prompt processing are no slower than main at one stream and at the widest admitted (3 streams here, 1 with DFlash2).