Repository navigation
perf: decode PQ2_0 & Q2_0 with its zero point as a min on Apple7/8 - #31
Conversation
…ter tile
PQ2_0 (Prism ML's ternary format, 2-bit codes, zero point 1, one half d
per 128 inputs) ran on Apple7/8's F32 register decode tile as a seeded
zero-point format: each 16-input half of a group closed its chain
against the table's seed, with a seed load and an epilogue per half.
The F32 instances now decode PQ2_0 through ZeroAsMin, which takes the
zero point as a min (m = -zero d): its chains close once per 32 inputs
against the table's sums with b = -129 d, exact. The coefficient unit
stays four spans and the host's splits are unchanged.
M2 Max, ms at one to four lanes (kernel-ab, ABBA):
pq20 17408x5120 gate/up 0.4730 0.8500 1.1952 1.6329 -> 0.4577 0.7930 1.1028 1.4294 (-3.2%, -6.7%, -7.7%, -12.5%)
pq20 5120x17408 down 0.2503 0.4557 0.6270 0.8580 -> 0.2419 0.4192 0.5778 0.7354 (-3.4%, -8.0%, -7.8%, -14.3%)
pq20 5120x6144 residual 0.0958 0.1673 0.2587 0.3554 -> 0.0934 0.1564 0.2381 0.3026 (-2.5%, -6.5%, -8.0%, -14.9%)
pq20 fused (10240+6144+256) 0.2658 0.4855 0.6513 0.8770 -> 0.2582 0.4463 0.6007 0.7667 (-2.9%, -8.1%, -7.8%, -12.6%)
pq20 248320x5120 vocab 3.0269 5.5959 7.6279 10.1152 -> 2.9096 5.1914 7.0320 9.2018 (-3.9%, -7.2%, -7.8%, -9.0%)
Controls (Q4_0, Q4_K) within noise.
Ternary-Bonsai-2-27B (PQ2_0) end to end (M2 Max):
decode-profile cycle ms (B1-B4, ABBA, 2048 prompt tokens, bf16 KV):
90.00 158.40 216.34 294.67 -> 87.32 149.17 201.06 263.57 (-3.0%, -5.8%, -7.1%, -10.6%)
bench-all decode tok/s (1 / 2 / 4 lanes, ANE share 0.35):
71.6 82.2 88.3 -> 73.6 87.4 99.7 (+2.8%, +6.3%, +12.9%)
backend-benchmark: bit-identical tokens across 1-4 lanes (hash 12422117031052473815, acceptance 0.831).
Outputs move within gguf-projection's fp64 bound (accumulation order change).
Extends the FmtPQ20 line to FmtQ20, one specialization: the formats share FmtTwoBit's algebra — value = d * (q - 1), zero point 1 — so Q2_0's chains close once per 32 inputs against the table's sums with m = -d, b = -129 d, exact. Tested on the model the PQ2_0 line lacked: ISTA-DASLab/Qwen3.8-Flash-Next- GSQ-RCO-GGUF:Q2_0, whose dense decode projections run gguf_decode_sgf_q20_* through this kernel (its routed experts are MoE on another kernel, and untouched). M1 Max, quiet machine: - gguf-projection PASS under MTL_SHADER_VALIDATION=1: 43 register formats against the fp64 reference, 1-4 lanes, S 1-16, fused and one-pass checks. - Isolated projections (gguf-projection-benchmark, median of 100): q20 17408x5120 gate/up, register_gu S1: L1 0.757 -> 0.692 ms (-8.6%), L2 1.468 -> 1.324 (-9.8%), L3 2.090 -> 1.829 (-12.5%), L4 2.621 -> 2.357 (-10.1%); q20 5120x6144 residual, register S4: -7% to -15% across lanes; the staged tiles are untouched and flat. - teacher-logits window 2048:256 over the pack (257 teacher-forced rows of the vocabulary, stock release metallib against this change): top-1 agreement 251/257 (97.7%), mean KL 0.0016, max |dlogit| 0.9 — the fp rounding class gguf-projection's fp64 bound allows. In this pack Q2_0 projections are about 1% of the decode cycle, so there is no end-to-end gain to claim from it; a model whose dense projections are Q2_0 would take the same range the PQ2_0 line measures.
|
You left the Q2_0 line untested for the lack of a model — I have one: Change: Testing (mirrors yours):
The staged tiles are flat, as expected — same shape and magnitude as your PQ2_0 table.
Honest caveat: in this pack the Q2_0 projections are about 1% of the decode cycle, so there is no end-to-end gain to claim from it; a model whose dense projections are Q2_0 should take the same range your PQ2_0 line measures. If you want the line, the branch is one commit on top of yours — happy to rebase it wherever it fits. |
|
Environment for the numbers above, which I should have spelled out there:
|
|
@abssoul great, thank you for testing it on a real Q2_0 model. I had a look at your commit, it is the same line as for PQ2_0 and the algebra is the same, so it is exact for Q2_0 too. Your numbers look like my PQ2_0 ones. I would keep it as your PR, it is your work and your measurement on a model I don't have. Maybe open it on top of this one (or after it is merged), it is a one-line follow-up. @paperniuk unless you prefer it in this PR, then I can cherry-pick it here with abssoul as author. |
|
Please cherry-pick it into this PR — one PR for the whole zero-as-min family is easier for everyone. No need to keep my name on it either: the branch is public ( |
|
@abssoul ok, thank you! I'm running gguf-projection and linear-plan with your commit on my M2 Max now (with and without shader validation). If they pass I will add it to this PR as your commit this morning, unless @paperniuk prefers otherwise. |
|
@abssoul thank you! gguf-projection and linear-plan pass with your commit on my M2 Max, with and without shader validation, so I added it to this PR as your commit. |
|
@paperniuk heads up for when you merge upstream. I sent the Apple9 version of these zero-point views upstream (#569). jianc99 merged them in #602 and then replaced them with one rule for all families (870f07d, in 1.5.0), which was faster than my views on M4 and M5 Max. 1.4.1 doesn’t conflict in this file, but 1.5.0 will conflict with the F32 tile. If you want, I can port his rule to the F32 instances and measure it on my M2 Max. If it’s as fast, the fork can drop its own views and the merge gets simpler. This PR can go in as is either way. |
|
Thanks, merged. Ternary Bonsai PQ2_0 on an M1 Max goes from 35.0 to 36.1 tok/s at one lane and 38.6 to 41.0 at two. It touched the same F32Format lines as #35, I kept both. |
Base:
paperniuk/splashapple7-m1-1.4(branchapple7-m1-pq20-zero-as-min, commit 4ee562b).This one takes #15's
ZeroAsMinview to PQ2_0, Prism ML's ternary format (2-bit codes, zero point 1, one halfscale d per 128 inputs). PQ2_0 is almost all of
prism-ml/Ternary-Bonsai-2-27B-gguf(the 27B target in 6.7 GB, 10 GB with itsdraft), so it fits 16 GB M1/M2 Macs. Now that your fix for #22 (6fb9807) lets Bonsai run on Apple7/8, Bonsai decodes
3% faster at one lane and 13% at four on my M2 Max, with the same tokens. As before, I used Claude Code to author it
and tested it by hand.
Before, PQ2_0 ran on the F32 register tile as a seeded zero-point format: each 16-input half closed its chain against
the table's seed, with a seed load and an epilogue per half. Through
ZeroAsMinthe zero point is an exact min(m = -d), so chains close once per 32 inputs against the table's sums with b = -129 d, exact. The coefficient unit
stays four spans, so the host's splits don't change. I did the same for Apple9 upstream in incoai#569 (3-7% on
projections there). Q2_0 shares
FmtTwoBitand could take the same line, but I have no Q2_0 model to test it, so Ileft it.
Update: the second commit (037f451, by @abssoul) adds the same line for Q2_0, which shares
FmtTwoBitwith PQ2_0, sothe view is exact for it too. On his M1 Max with
ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF:Q2_0the Q2_0 registerprojections take 9-15% less time at 1-4 lanes, and teacher-forced logits stay within rounding (details in his commit
message). Q2_0 is about 1% of that model's decode cycle, so there is no end-to-end number for it. With both commits,
gguf-projectionandlinear-planpass on my M2 Max with and without shader validation.What changed
runtime/metal/kernels/decode/linear_gguf_sgmatrix.metal: specializedF32Format<FmtPQ20>toZeroAsMin<FmtPQ20>, and updated the doc comment aboveZeroAsMin.GGUF_SG_FORMATinstantiatesgguf_decode_sgf_pq20_*usingF32Format<FmtPQ20>::type.Shape<ZeroAsMin<FmtPQ20>>evaluates toHasMin = 1,Seeded = 0,CG = 1,UnitSpans = 4(unchanged from seeded PQ2_0's 4 spans).gguf-projection's fp64 bound.Results (M2 Max, 38-core GPU)
1. Fused decode cycle (
decode-profile, 2048 prompt tokens, bf16 KV, ABBA)Repeat runs agree within 0.24%.
2. End-to-end decode throughput (
bench-all, ANE share 0.35)MLX baseline drift check before and after ran within 0.1%. Prefill latency is unchanged (77.88 s 14K cold, 23.78 s hit+4K, 10.64 s 2K).
In
backend-benchmark(--scenario decode --samples 3), output token hashes and draft acceptance rate (0.831) are bit-identical across all lanes between before and after.3. Isolated projection kernels (
kernel-ab.sh, ABBA, from the first round of this change)Controls remain flat within run-to-run noise.
Testing
linear-plan: PASS without validation, PASS withMTL_SHADER_VALIDATION=1.gguf-projection(every register-tile format, PQ2_0 among them, against an fp64 reference at 1-4 lanes): PASSwithout validation, PASS with
MTL_SHADER_VALIDATION=1.make check: CPU tests, the Python suite (1320 tests) and ruff pass. In the Metal tests, everything up togguf-moepasses under validation (this includesgguf-rotationandgguf-projection).gguf-moetimes outunder validation on my M2 Max on
apple7-m1-1.4itself too (the 120 s command limit, same as perf: column tile for the MMA paths one- and two-row expert tiles #29 reports on anM1 Max), and that stops the rest of
test-engine-metal, so the tests after it didn't run in thismake check.