Skip to content

Qwen3.5-9B support in Zig engine - #1826

Merged
ashhart merged 7 commits into
ashhart:mainfrom
grantoverton:qwen35-9b-support
Nov 2, 2026
Merged

ashhart merged 7 commits into
ashhart:mainfrom
grantoverton:qwen35-9b-support

Conversation

@grantoverton

@grantoverton grantoverton commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

Part of #503. This builds the 9B on the 27B engine, as #503 describes: on main a qwen3_5 checkpoint that isn't the 2B opens the 27B's Metal engine, whose supported() admits only the 27B's shape, so the 9B is refused with UnqualifiedQwenShape. With this branch a dense qwen3_5 checkpoint at the 9B shape (hidden 4096, intermediate 12288, 32 layers on the 27B's 3:1 Gated DeltaNet / full-attention schedule, 16 query heads over 4 KV heads of 256, 16 key and 32 value GDN heads of 128) opens the 27B engine, and its attention kernels run at a KV group of 4. It serves the stock mlx-community/Qwen3.5-9B-MLX-4bit @938d891 and a mixed 4/5-bit oQ4e requant of ornith-ai/Ornith-1.5-9B exactly on an M3 Ultra and an M5 Ultra. The 27B keeps its kernels: at Q27_G 6 and Q27_KVH 4 its attention library compiles to the same optimized IR as main's.

#503 is claimed, so this is offered as the qualification run it asks for; I'm happy to fold it into your branch, split it, or rework it to whatever shape you have in mind.

This branch carries #1125's commit 35cb0f5 (here 24dfbf3, author kept: the KV group as the shader constants Q27_G / Q27_KVH) until it merges, so please review the last five commits here (from "16 query heads over 4 KV heads on the attention Geometry" on). Base: main 9e8c368, which has #1822 (b38f8f6cb, the mixed-bit prompt passes over 128 rows that the oQ4e checkpoint needs on M5).

What this changes

Five commits on #1125's attention Geometry:

  1. 16 query heads over 4 KV heads on the attention Geometry, a KV group of 4 (attention_plan.zig). Geometry.check admits 16 over 4 beside 24 over 4 and 32 over 2; the 9B's rank of a two-way split (8 over 2) is not admitted. Two tests used 16 over 4 as a shape no shader builds: the geometry test now expects a group of 4 there and refuses 8 over 2 and 20 over 4 (GQA 5), and the descriptor test in attention.zig refuses 20 query heads instead of 16, so the refusal it checks is unchanged. Test that fails before: the geometry test's 16-over-4 case (TargetAttentionShape); with the change alone the descriptor test fails, which is why its premise moves to GQA 5.
  2. The decode-only attention build takes the KV-group defines too (simd_attention.zig, decode_round.zig). Before M5, a chip that can't build the prompt attention kernels falls back to a library without them, and that fallback compiled without the Q27_G / Q27_KVH defines, so its kernels took the 27B's group of 6 whatever the checkpoint. simd_attention.grouped puts the defines in front of either build and the runner uses it for both (this applies to qwen27: serve Qwen3.5-MoE (Ornith-1.5-122B) on Metal: routed experts, its DFlash2 drafter and the publisher's 4-bit weights #1125's group of 16 as well; feel free to take it into qwen27: serve Qwen3.5-MoE (Ornith-1.5-122B) on Metal: routed experts, its DFlash2 drafter and the publisher's 4-bit weights #1125). Test that fails before: both builds, with and without the prompt kernels, open with the defines and are otherwise library's text.
  3. Read a checkpoint's A_log in f32 as well as bf16 (gpu_weights.zig). mlx-vlm conversions of Qwen3.5, mlx-community/Qwen3.5-9B-MLX-4bit among them, store linear_attn.A_log in f32, and the loader asked for bf16, so that checkpoint refused with BadNativeTensor once its shape was admitted. The GDN kernels already read either through the params' weight flags, which the round takes from the tensor's own dtype, so the loader now takes A_log as stored. A bf16 A_log loads as before. Test that fails before: f32 and bf16 load, f16 and a wrong shape refuse.
  4. A qwen3_5 dense checkpoint at the 9B shape opens the 27B engine (model.zig). supported() takes either size (hidden, intermediate, layers, query heads, GDN value heads) from a two-row table; the vocabulary, KV heads, head dims, GDN key heads, conv kernel, rotary width and layer schedule checks are unchanged. Test that fails before: the 27B and 9B configs open, and the 4B, the 9B's widths at the 27B's depth and the 27B with the 9B's GDN heads refuse.
  5. A drafter for another target is refused by name before either checkpoint loads (dflash/config.zig, dflash/runtime_model.zig, qwen27_served.zig). With a 9B checkpoint, --drafter <the 27B's DFlash2> used to load the 9B's weights and then fail in the drafter load with TargetBinding. open() now reads both folders' config.json first. A stock DFlash2 whose width, vocabulary or tapped depth belongs to another target is refused with DraftForAnotherTarget, and a DFlash2 config that isn't the stock 27B one with NotStockDFlash2. The drafter load keeps the same check. Test that fails before: two folders holding only config.json; the 27B target admits the 27B's DFlash2, and the 9B target refuses it with DraftForAnotherTarget and a 9B-shaped DFlash2 with NotStockDFlash2. Receipt (M3, serve <9B> --drafter <dir>): before, both refusals came after the 9B loaded (peak footprint 5.87 GB); now they come before either loads (0.25 GB).

Exactness

Checkpoints: mlx-community/Qwen3.5-9B-MLX-4bit @938d891 (uniform 4-bit, group 64; model-00001-of-00002.safetensors sha256 a68b87558c6ef43f…) and an oQ4e requant of ornith-ai/Ornith-1.5-9B @489cb97 (4-bit group 64 with 52 modules at 5 bits; not published, model-00001-of-00002.safetensors sha256 753334c603c34f36…, model-00002-of-00002.safetensors 4ac7aca85476e365…). No drafter: the 9B's DFlash drafters are a follow-up, so every token is a plain lane token. Every check below ran at debe0dfc1, before commit 5 and before the move onto main 9e8c368 (tensorfold-native sha256 15bdcc08dc123a3f… built on the M3, ac74d1962bf1b183… on the M5), greedy, thinking off. The move changes no qwen27 source: main's qwen27 sources are #1822's, which debe0dfc1 carried byte for byte. Main's lane changes since then are opt-in, with their defaults unchanged, and qwen27 sets none of them. Commit 5 changes only how --drafter is read at open, which none of these runs uses. The host suite, goldens, native build, dev-tools and lean check ran again at the tip 6f0480d13 (below).

  • Before: main 5acf042 refuses both: the native Qwen3.8-27B engine cannot load <checkpoint> (UnqualifiedQwenShape).
  • Known answers (8 short greedy chat prompts with fixed answers: a capital, 17 × 23, a reversed word, a chemical symbol, days in a leap year, a planet, a square root, a copy line): 8/8 on both checkpoints and both chips.
  • Concurrent and serial (--parallel 4 --context 65536 --prompt-cache-gib 0, bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial): every alone cell (3 at one stream, 12 at four) equal, every serial ("draft": false) cell equal, 0 failed, on both checkpoints and both chips.
  • Resumed equals fresh (--parallel 4 --context 65536, prompt cache on): a 6,027-token prompt sent twice, the second resumed at 6,020 cached tokens, same token hash (uniform c40a2accd7f8 on both chips; oQ4e 45a7e1591588 on the M5, aad22c3e37f1 on the M3).
  • Seams (tf-qwen27-run --run --tokens <2,048-token public prefill array> --max-tokens 64 --chunk 128 --compare-chunk C --seams 256, C = 16 and 37): the state/KV/logits fingerprint is byte-equal at every 256th prompt row (7 seams), at the prompt end and after the decoded reply, for both chunk pairs, on both checkpoints and both chips.
  • Cancellation and reset (--parallel 4, 14 steps): two reference probes (the two bench_openai prompts, greedy, 128 tokens), then a stream dropped mid-decode, a 22k-token prompt dropped mid-prompt, a stream dropped beside a running probe (whose hash is checked too) and a request past the window (refused); after each, both probes equal their reference hashes. Pass on both checkpoints and both chips (uniform a9af571d979c / a4dc1a079732 on the M3, a9af571d979c / 05cf379ac28e on the M5; oQ4e 7a0121a0d461 / 6ea245ec9906 on the M5, 290f3dd3af44 / 361585a76350 on the M3).
  • Rows. Before M5 a prompt pass is 128 rows, so qwen27: chunk mixed-bit prompt passes at 128 rows #1822's chunking runs one step there; on the M5, where prompt passes reach 2,048 rows, the oQ4e checkpoint's 4/5-bit stacks take its 128-row steps, and every prompt above (up to 65,536 tokens) went through them.
  • Precision. Activations stay bf16 and every product keeps its fp32 sums; no reduction order changes. The 9B's projections run the 27B's kernels at its own shapes; A_log is read at the precision the checkpoint stores.
  • The 27B's attention is unchanged. I rebuilt the attention library text the runner compiles (with and without tensor units, with and without the prompt kernels) from main 5acf042 and from this tip with Q27_G 6 / Q27_KVH 4 in front, and compiled each with xcrun metal -std=metal4.0 -O2 -S -emit-llvm. The preprocessed text (with tensor units and the prompt kernels) differs only where qwen27: serve Qwen3.5-MoE (Ornith-1.5-122B) on Metal: routed experts, its DFlash2 drafter and the publisher's 4-bit weights #1125 names the KV-head constant (G * 4 → G * KVH, two merge kernels and their constants); the optimized IR differs only in the module name and source file name lines (sha256 of the rest: da3770108d7ed1fc, 9f7484c6dedfffd4 and 9fc558a47d996820 on both sides).

Accuracy against float64: attention at a KV group of 4

tf-qwen27-attention-precision --heads H --kv-heads 4 [--prompt] (the engine's own attention library build, against a float64 softmax of the same bf16 rows), H = 24 (the 27B) and 16 (the 9B). Max and mean |err|:

Chip Path Queries 24 over 4 16 over 4
M3 Ultra decode, rows below 8192 [-1, 1) 4.0455e-3, 6.9397e-5 4.1016e-3, 6.9009e-5
M3 Ultra decode, rows from 8192 [-1, 1) 2.3892e-4, 1.7391e-5 2.1997e-4, 1.7493e-5
M3 Ultra decode, rows below 8192 [-4, 4) 4.0788e-3, 3.4130e-4 3.9742e-3, 3.3785e-4
M3 Ultra prompt kernel (simdgroup) [-1, 1) 3.9061e-3, 6.1423e-5 3.9028e-3, 6.1695e-5
M3 Ultra prompt kernel (simdgroup) [-4, 4) 3.9064e-3, 3.0982e-4 3.9064e-3, 3.1142e-4
M5 Ultra decode, rows below 8192 [-1, 1) 4.0455e-3, 6.9396e-5 4.1016e-3, 6.9011e-5
M5 Ultra decode, rows from 8192 [-1, 1) 2.3892e-4, 1.7391e-5 2.1997e-4, 1.7493e-5
M5 Ultra prompt kernel (tensor units) [-1, 1) 4.0934e-3, 6.2411e-5 4.1016e-3, 6.2676e-5
M5 Ultra prompt kernel (tensor units) [-4, 4) 4.1441e-3, 3.1074e-4 4.1320e-3, 3.1235e-4

0 values not finite in every run. The group of 4 sits within the 27B's own error at every row range on both chips. The projection classes run the existing kernels at the 9B's shapes, with no change to their arithmetic; tf-qwen27-widths checks the 27B's shapes only, and I didn't run it at the 9B's (not run, below).

Fidelity

Teacher-forced against mlx_lm 0.32.3 (development build 0d4d3d04f, mlx 0.32.5) on the same checkpoint and the same Mac. mlx_lm renders upstream's four public teacher prompts (tools/zig/qwen27_gate.py: copy, arithmetic, story, code; chat template, thinking off), prefills each in one pass and decodes 256 greedy tokens with EOS ignored, keeping its top two logits at every step; that continuation is the fixed history. This branch's tf-qwen27-run --run --teacher --force-length --chunk 128 --compare-chunk 16 feeds every history token and records its top two at all 1,024 positions (prompt state byte-equal between chunk 128 and 16 on every prompt), and upstream's qwen27_teacher_judge.py judges each difference (its tie policy: the same two contenders, both margins within one bf16 spacing).

Chip Checkpoint Positions Top-1 equal Differences Beyond the tie policy Judge
M3 Ultra uniform 4-bit 1,024 1,017 7 0 pass
M5 Ultra uniform 4-bit 1,024 1,021 3 0 pass
M3 Ultra oQ4e 1,024 1,018 6 0 pass
M5 Ultra oQ4e 1,024 1,011 13 5 fail

The differences beyond the policy, all on the oQ4e checkpoint on the M5 (mlx_lm pick and runner-up with logits and margin; this branch's the same):

Chip Prompt Position mlx_lm top 2 logits margin native top 2 logits margin
M5 copy 12 248044, 248068 2.6875 / 2.4375 0.25 248068, 248044 2.171875 / 2.15625 0.015625
M5 arithmetic 166 20003, 760 24.375 / 24.125 0.25 760, 20003 24.75 / 24.5 0.25
M5 story 90 2919, 4723 16.5 / 16.25 0.25 4723, 2919 16.5 / 16.375 0.125
M5 story 122 11, 2218 22.75 / 22.75 0 27090, 11 22.875 / 22.75 0.125
M5 story 123 5983, 8720 16.375 / 16.375 0 264, 5983 16.25 / 16.25 0

The other differences are near-ties (both margins 0 or 0.125, the same two contenders): uniform M3 story 52, 89, 153, 196, 197, 251 and code 9; uniform M5 story 87, 237 and code 9; oQ4e M5 copy 56, 149, 200, 206 and story 83, 91, 179, 180; oQ4e M3 copy 108, 123, story 90, 103, 146 and code 206. Each chip's mlx_lm history is its own, since mlx_lm's kernels differ by chip too. On the oQ4e checkpoint on the M5, three of the five are margins of two bf16 spacings (0.25) in mlx_lm at 16 to 25, and two (story 122 and 123) name a different runner-up. At copy 12 both engines' top logits are near 2 (a low-confidence position after the copied line) and they rank the same two tokens the other way. Both checkpoints pass on the M3, and the uniform one on the M5. The full judge records (with both engines' top two at every difference and the files' sha256) are in the receipts.

The same judge on a mixed-width 27B on the M5, on main 9e8c368 (with #1822). The checkpoint is an oQ4e requant of Blackfrost-AI/Qwen3.8-27B-ABLITERATED-BF16 @9d85770: 4-bit group 64, 96 modules at 5 bits, lm_head at 6, not published. Over 1,024 positions, 1,018 top-1 are equal, with 6 differences and 2 beyond the tie policy. The two are copy 15 and arithmetic 153, both <|im_end|> against a newline, with mlx_lm margins of 0.5 and 0.25. That is a fail, as the 9B oQ4e fails there, while the uniform 4-bit 9B passes on the same Mac, so it's the M5's mixed-width path. A float64 replay of the 9B oQ4e's M5 histories (mlx_lm's model code on the CPU, weights dequantized exactly, every op in float64) ranks this branch's token first at four of its five positions beyond the policy, and mlx_lm's runner-up first at the fifth (story 123). Over all 1,024 positions, float64's top-1 equals mlx_lm's at 1,008 and this branch's at 1,010.

Performance

  • M3 Ultra: Mac Studio (Mac15,14), 80-core GPU, 256 GB, macOS 26.4.1 (25E253), Xcode 26.6, no tensor units.
  • M5 Ultra: Mac Studio (Mac17,15), 64-core GPU, 96 GB, macOS 27.2, Xcode 27.2, tensor units.

Zig 0.17.0, zig build native tf-qwen27-run -Dcpu=apple_m1 -Dversion=1.0.0 on each Mac. Clients are this branch's tools/ on Python 3.14 (M3) and 3.15 (M5). Server serve <checkpoint> --parallel 4 --context 65536 --prompt-cache-gib 0, no drafter; mlx_lm.server --prompt-cache-size 0 (mlx_lm 0.32.3, a development build at 0d4d3d04f, with mlx 0.32.5) on the same checkpoint. One model process of mine at a time; neither Mac was quiet: both also served other models for other work, and the M3 also ran other builds, so I note beside each table what ran with it. Main has no before-run (it refuses both checkpoints), so the 9B's numbers are served decode beside mlx_lm, and the 27B path's "unchanged" rests on the identical attention IR above. Runs alternate this branch, mlx_lm, mlx_lm, this branch. The raw receipt files (decode, concurrency, prefill and cancellation records, the exact prefill arrays and request-body hashes, seam and fidelity records, the attention float64 logs, the IR comparison and the host-suite summary) are available on request.

Decode and concurrency, M5 Ultra. bench_openai.py … --tokens 256 --reps 3 --temperatures 0 (medians; fibonacci-raw code and gpu-chat-no-think chat) and bench_concurrent.py … --levels 1,4 --alone --serial (aggregate tok/s).

Run bench_openai code chat concurrent code ×1 code ×4 chat ×1 chat ×4 Peak footprint
uniform 4-bit, this branch 1 123.09 123.22 123.1 465.5 123.2 463.1 22.10 GB
uniform 4-bit, mlx_lm 1 138.28 144.65 145.7 143.4 145.3 137.2 13.25 GB
uniform 4-bit, mlx_lm 2 145.62 145.03
uniform 4-bit, this branch 2 123.25 123.18 15.61 GB
oQ4e, this branch 1 101.59 101.61 101.4 364.3 101.4 376.0 22.14 GB
oQ4e, mlx_lm 1 140.84 138.05 141.9 140.1 141.7 139.8 13.32 GB
oQ4e, mlx_lm 2 135.66 141.96
oQ4e, this branch 2 101.69 101.25 15.64 GB

The oQ4e rows come from a build of this branch without commit 3 (8744853a9, tensorfold-native a2d669804e4d569d…), the one quiet window I had for that checkpoint; commit 3 changes nothing for its bf16 A_log, and the tip gives 101.36 / 101.80 and ×4 376.9 / 375.2 on it. Two later tip runs on the oQ4e checkpoint (36.78 / 36.43 and 28.76 / 28.22) ran beside another job that held the M5's GPU at 92% utilization, and I don't count them. On the M5 one native stream decodes below mlx_lm (on the code prompt, the mean of the two runs each: 0.87× on the uniform checkpoint, 0.74× on oQ4e), and at four streams the native aggregate is 3.2× and 2.6× mlx_lm's, which serves the streams one after another.

Decode and concurrency, M3 Ultra (beside other served models and other Zig builds; the spread inside each run shows it).

Run bench_openai code chat concurrent code ×1 code ×4 chat ×1 chat ×4 Peak footprint
uniform 4-bit, this branch 1 48.87 47.81 93.6 283.5 93.0 278.4 20.18 GB
uniform 4-bit, mlx_lm 1 54.74 39.78 34.6 96.6 100.5 97.3 16.23 GB
uniform 4-bit, mlx_lm 2 46.25 70.49
uniform 4-bit, this branch 2 48.96 50.48 13.69 GB
oQ4e, this branch 1 39.95 88.14 85.2 125.1 36.3 90.2 20.21 GB
oQ4e, mlx_lm 1 36.23 49.44 72.1 70.5 71.0 69.5 16.30 GB
oQ4e, mlx_lm 2 72.45 72.28
oQ4e, this branch 2 55.49 53.12 13.73 GB

These M3 numbers are not a speed comparison: the same binary gives 48.87 tok/s in bench_openai and 93.6 alone in bench_concurrent minutes later. I'd rerun them on a quiet M3 before anyone quotes them.

Cold prefill. Exact token arrays: upstream's public prefill paragraph ("A lighthouse keeper writes in the log as the boats cross the bay.") repeated, tokenized once without special tokens by the checkpoint's own tokenizer.json, each array the first N ids (each decodes and re-encodes to itself; both checkpoints' tokenizers give the same arrays). POST /v1/completions with the ids as prompt, max_tokens 1, temperature 0, ignore_eos; server --parallel 1 --context 131072 --prompt-cache-gib 0; one request at a time, with a short request whose first token differs between measurements so no prompt shares a prefix with the last; every record has cached_tokens 0 and prompt_tokens equal to N. One 2,048-token warm-up, then 3 rounds over the sizes, wall-clock medians, tok/s (s). mlx_lm takes the same arrays decoded as text, with the same checks.

Tokens M5 uniform, this branch M5 uniform, mlx_lm M5 oQ4e, this branch M5 oQ4e, mlx_lm M3 uniform, this branch M3 uniform, mlx_lm M3 oQ4e, this branch M3 oQ4e, mlx_lm
2,048 4,731.4 (0.43) 5,594.9 (0.37) 2,724.8 (0.75) 5,566.8 (0.37) 690.7 (2.97) 1,183.5 (1.73) 756.4 (2.71) 1,102.1 (1.86)
8,192 4,497.1 (1.82) 5,620.6 (1.46) 2,469.7 (3.32) 5,621.1 (1.46) 619.2 (13.23) 1,167.8 (7.01) 729.2 (11.23) 1,126.1 (7.27)
32,768 3,978.3 (8.24) 4,760.0 (6.88) 2,333.2 (14.04) 4,882.8 (6.71) 578.6 (56.63) 1,036.9 (31.60) 654.8 (50.05) 992.0 (33.03)
65,536 3,529.7 (18.57) 4,128.1 (15.88) 2,163.3 (30.29) 4,080.9 (16.06) 543.5 (120.59) 899.1 (72.89) 574.4 (114.10) 835.7 (78.42)

mlx_lm prefills faster than this branch at every size measured: on the M5 by 1.17× to 1.25× on the uniform checkpoint and 1.89× to 2.28× on oQ4e, whose 5-bit members run on the generic row kernels. The M5 oQ4e columns come from the same quiet window as its decode rows (build 8744853a9); a tip run beside the other job gave 1,685.6 to 1,376.8 tok/s. The M3 columns ran beside the other work. 128k and the native window (262,144) not run. This PR is about loading and serving the 9B exactly; its speed is for follow-ups.

Memory

Weights as stored: 5.6 GB (uniform) and 6.0 GB (oQ4e). Peak physical footprint (time -l of the server process): at --parallel 4 --context 65536 --prompt-cache-gib 0 after the bench and cancellation steps, 22.10 GB (uniform, M5), 22.14 GB (oQ4e, M5), 20.18 GB (uniform, M3), 20.21 GB (oQ4e, M3); prefill server at --context 131072 after the 64k prompts, 21.83 GB (uniform, M5), 20.01 GB (uniform, M3), 20.04 GB (oQ4e, M3); mlx_lm.server 13.25 GB (M5) and 16.23 GB (M3) on the uniform checkpoint. The qwen27 Metal backend reports no device allocations in /memory; the /memory snapshots after load and after the runs are in the receipts. RAM classes: M3 Ultra 256 GB, M5 Ultra 96 GB.

Host suite and lean check

  • zig build test test-golden -Dcpu=apple_m1 --summary all at the tip 6f0480d13 on the M3, from an empty cache: 594/594 steps, 9,200/9,255 tests passed, 55 skipped, 0 failed. Goldens: 322/335 equal with the 13 declared differences and 0 unexpected, and 55/55. Each new test went red and then green under zig build test as its commit was made: commits 1, 2 and 4 on main 5acf042, commit 3 on main 3ab49e3 with qwen27: chunk mixed-bit prompt passes at 128 rows #1822, and commit 5 there too. Commit 5's test didn't build without the pre-load check, and fails with the check removed.
  • zig build native -Dcpu=apple_m1 -Dversion=1.0.0: exit 0 at 6f0480d13 on the M3 (tensorfold-native sha256 49ca831fdc46193e…) and on the M5 (3172532aaad68d3e…). zig build dev-tools -Dcpu=apple_m1: exit 0.
  • python3 tools/zig/lean_check.py: 424 problems at 6f0480d13, main 9e8c368's 424; none in a file this branch touches. Touched files are under 600 lines (model.zig 248, gpu_weights.zig 370, simd_attention.zig 154, decode_round.zig 516, attention_plan.zig 470, attention.zig 464, dflash/config.zig 238, dflash/runtime_model.zig 309, qwen27_served.zig 515).
  • No kernel is new; qwen27: serve Qwen3.5-MoE (Ornith-1.5-122B) on Metal: routed experts, its DFlash2 drafter and the publisher's 4-bit weights #1125's commit changes the attention sources' constants, and qwen27's attention kernels are not in the kernel-contract registry at this base, so there is no contract SUMMARY line to post for them.

Platforms and what was not run

  • Run: Metal on an M3 Ultra (no tensor units) and an M5 Ultra (tensor units).
  • Not run: M1, M2 and M4 (no hardware; they take the M3's kernels). CUDA and HIP (no CUDA or HIP code changes).
  • Not run: drafting. The 9B's DFlash drafters (for example z-lab/Qwen3.5-9B-DFlash) have their own shape and are Qwen3.5-9B DFlash drafter in Zig engine #1851. Here a DFlash2 given with a 9B checkpoint is refused before load (commit 5), so drafted equals "draft": false isn't exercised beyond the serial check above.
  • Not run: the 27B served before and after. Its attention IR is identical (above) and no other 27B code path changes; on a 27B checkpoint the receipts would be:
./zig-out/native/bin/tensorfold-native serve <Qwen3.8-27B-MLX-4bit dir> --drafter <Qwen3.8-27B-DFlash2 dir> --port 8080 --parallel 4 --context 65536 --prompt-cache-gib 0
python3 tools/bench_openai.py http://127.0.0.1:8080 qwen27 --tokens 256 --reps 3 --temperatures 0 --output decode.json
python3 tools/bench_concurrent.py http://127.0.0.1:8080 qwen27 --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial --output concurrent.json
  • Not run: the fallback attention build of commit 2 on a chip whose prompt kernels don't build (neither Mac takes it). On such a chip: ./zig-out/bin/tf-qwen27-attention-precision --heads 16 --kv-heads 4 and the served checks above with the 9B.
  • Not run: tf-qwen27-widths at the 9B's projection shapes (the tool has the 27B's), 128k and native-window prefill, and quiet-machine decode on the M3: serve <9B dir> --parallel 1 --context 262144 --prompt-cache-gib 0 with the 131,072- and 262,143-token arrays.
  • Not run: atc overlap --json (ATC isn't installed here).

Checklist

  • Every token still goes through the lane rounds; no serial path, and nothing that needs drafts off.
  • A resumed prompt equals a fresh one, and concurrent equals solo (both checkpoints, both chips); drafted equals plain is not exercised (no 9B drafter, above).
  • No precision traded for speed: no reduction order changes; A_log is read at its stored precision.
  • Prompt processing no slower than the last release: no before-run exists for the 9B (main refuses it); the 27B's attention IR is identical, and its served prefill was not re-run.
  • The description says which platforms I ran (M3 Ultra, M5 Ultra) and which I could not.
  • New tests fail before the change and pass after it; the two tests whose premise changed now refuse GQA 5.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

@ashhart

ashhart commented Nov 1, 2026 •

Copy link
Copy Markdown
Owner

Thanks @grantoverton, this is a thorough qualification: seams, cancellation, resumed against fresh, float64
attention at a group of 4, and the 27B's attention IR shown identical. The fallback build without the KV-group
defines and the f32 A_log from mlx-vlm conversions are both real catches. Our own loader for #503 (#1817) touches the same
supported(), and we've compared the two: yours lands, with your geometry, fallback-build and A_log commits as they
are and your receipts as the 9B's qualification. We'll rebase #1817 onto yours and keep only what it adds, or close it. Two things we'd want either way:

  1. The oQ4e fail on the M5 (5 positions beyond the tie policy). Could you run the same judge on a mixed-width
    27B on the M5, on main with qwen27: chunk mixed-bit prompt passes at 128 rows #1822? If the 27B shows the same, it's the M5's 5-bit path or mlx_lm's M5 kernels,
    not the 9B. If you can, also compare those five positions against a float64 replay of the same rows.
  2. --drafter with a 9B target: please show, with a test or a receipt, that the 27B's DFlash2 (and a 9B
    DFlash2) given with a 9B checkpoint is refused by name before load.

This also waits on #1822 and #1125 landing. We'll take #1822 first.

@grantoverton

Copy link
Copy Markdown
Contributor Author

Thanks for the read, and for taking it with #1817. Both asks are below. Since #1822 has landed, I've also rebased the branch onto main 9e8c368b0 without its copy of #1822. It still carries #1125's commit, and it has one new commit for the second ask. #1851 and #1871 are rebased on top.

1. The mixed-width 27B on the M5, same judge

I built main 9e8c368b0, which has #1822 as b38f8f6cb, with zig build native tf-qwen27-run -Dcpu=apple_m1 -Dversion=1.0.0 on the M5 Ultra (tensorfold-native sha256 3383a4a9072e1fd1…, tf-qwen27-run da1c252ae8662c59…).

The checkpoint is an oQ4e requant of Blackfrost-AI/Qwen3.8-27B-ABLITERATED-BF16 @9d85770, which isn't published: affine 4-bit, group 64, with 96 modules at 5 bits and lm_head at 6. The 5-bit modules are all 48 GDN out_proj, 23 in_proj_a, 23 in_proj_b, one in_proj_z and embed_tokens. Shards sha256: 9874968f4766f140…, 11935fae589667f0…, 26255a37f5deb542…, effd1071b0a0806f….

The method is the same as for the 9B:

  • mlx_lm 0.32.3 (0d4d3d04f, mlx 0.32.5), on the same Mac, renders the four public teacher prompts. It prefills each in one pass and decodes 256 greedy tokens with EOS ignored, keeping its top two at every step.
  • Main's tf-qwen27-run --run --teacher --force-length --chunk 128 --compare-chunk 16 feeds those histories. The prompt state is byte-equal between chunk 128 and 16 on all four.
  • Main's qwen27_teacher_judge.py judges.

This checkpoint's chat template puts a default system message ahead of each prompt, so its prompts are 294 to 298 tokens. The 9B's were 24 to 28.

Chip Checkpoint Positions Top-1 equal Differences Beyond the tie policy Judge
M5 Ultra 27B, oQ4e (4/5/6-bit), main with #1822 1,024 1,018 6 2 fail
M5 Ultra 9B, oQ4e (4/5-bit), from #1826 1,024 1,011 13 5 fail
M5 Ultra 9B, uniform 4-bit, from #1826 1,024 1,021 3 0 pass

The 27B's two:

Prompt Position mlx_lm top 2 logits margin native top 2 logits margin
copy 15 248046, 198 26.0 / 25.5 0.5 198, 248046 26.125 / 26.0 0.125
arithmetic 153 248046, 198 22.75 / 22.5 0.25 198, 248046 22.875 / 22.875 0

Both are <|im_end|> against a newline. The other four differences are near-ties within the policy (story 49, 105, 174 and 233).

So on the M5 the mixed-width 27B fails the same judge on main with #1822, as the mixed-width 9B does, while #1826's uniform 4-bit 9B passes there. That puts it on the M5's mixed-width path, ours or mlx_lm's, and not on the 9B.

Float64 replay of the 9B's five positions

I replayed the same M5 mlx_lm histories through a small float64 reference. It is mlx_lm's own Qwen3.5 model code on the CPU, with every weight dequantized exactly and every op in float64. MLX's CPU paths for the transcendentals, conv1d, and the fused norm, RoPE and attention compute at float32 accuracy, so I replaced those. Every replaced op is within 1.3e-13 of numpy float64. I'm happy to share the script.

Prompt Position float64 top 2 float64 logits float64 margin mlx_lm pick native pick float64 top-1 is
copy 12 248068, 248044 1.9401 / 1.9011 0.0391 248044 248068 the native pick
arithmetic 166 760, 20003 24.4344 / 24.2877 0.1467 20003 760 the native pick
story 90 4723, 2919 16.4256 / 16.4143 0.0114 2919 4723 the native pick
story 122 27090, 2218 22.8421 / 22.8094 0.0327 11 (22.7412) 27090 the native pick
story 123 8720, 5983 16.3167 / 16.2519 0.0648 5983 264 (16.1971) neither (mlx_lm's pick is second)

At four of the five, float64 ranks this branch's token first. At story 123 it ranks mlx_lm's runner-up first, with mlx_lm's pick second and ours third.

Every float64 margin there is under one bf16 spacing at its magnitude except copy 12: 0.039 at a logit near 2, where the spacing is 0.0156. At that position mlx_lm reads 2.69 / 2.44 against float64's 1.94 / 1.90, and this branch reads 2.17 / 2.16.

  • The 13 differences on that checkpoint: float64 sides with this branch at 7, with mlx_lm at 5, and with neither at 1.
  • All 1,024 positions: float64's top-1 equals mlx_lm's at 1,008 and this branch's at 1,010.
  • Mean gap between each engine's top-1 logit and float64's logit for the same token: 0.0652 for mlx_lm and 0.0625 for this branch.

2. --drafter with a 9B target

Before, the refusal had a name but came after the target loaded: open() loaded the 9B's weights, and the drafter load then failed with TargetBinding (or NotStockDFlash2).

The new commit, 6f0480d13, reads both folders' config.json before either checkpoint loads:

  • A stock DFlash2 whose width, vocabulary or tapped depth belongs to another target is refused with DraftForAnotherTarget.
  • A DFlash2 config that isn't the stock 27B one keeps NotStockDFlash2.
  • The drafter load uses the same check.

The test uses two folders holding only config.json, so nothing could load weights:

  • The 27B target admits the 27B's DFlash2.
  • The 9B target refuses it with DraftForAnotherTarget.
  • The 9B target refuses a 9B-shaped DFlash2 with NotStockDFlash2. That config is the stock one at hidden 4096, intermediate 12288, 32 target layers and taps 1/8/15/22/29.

Before the fix the test didn't build, because there was no pre-load check. With the check removed, it fails on its first refusal.

Receipt on the M3: serve <mlx-community/Qwen3.5-9B-MLX-4bit @938d891> --drafter <dir> --parallel 1 --context 4096 --prompt-cache-gib 0, with the peak footprint from time -l:

Build --drafter Refused with Peak footprint
before (debe0dfc1) z-lab/Qwen3.8-27B-DFlash2 @50307d4 TargetBinding, after the 9B loaded 5.87 GB
before (debe0dfc1) 9B-shaped DFlash2 (config only) NotStockDFlash2, after the 9B loaded 5.87 GB
after (6f0480d13) z-lab/Qwen3.8-27B-DFlash2 @50307d4 DraftForAnotherTarget, before either loaded 0.25 GB
after (6f0480d13) 9B-shaped DFlash2 (config only) NotStockDFlash2, before either loaded 0.25 GB

For the first, the new message reads: --drafter <dir> drafts for another target: its width, vocabulary or tapped layers do not fit <checkpoint> (DraftForAnotherTarget). There's no published 9B DFlash2, so the 9B-shaped one is a config-only folder.

#1851 (the 9B's DFlash v1 drafter) is rebased on this. Its test adds that the 9B target admits the v1 drafter and the 27B target refuses it by name.

At the new tip 6f0480d13, on the M3:

  • zig build test test-golden -Dcpu=apple_m1 --summary all from an empty cache: 594/594 steps, 9,200/9,255 tests passed, 55 skipped.
  • Goldens: 322/335 equal with the 13 declared differences and 0 unexpected, and 55/55.
  • zig build native and zig build dev-tools exit 0.
  • lean_check.py: 424, main's 424.

The M5 also builds the tip (zig build native tf-qwen27-run, exit 0).

@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thanks @grantoverton. Our M5 Ultra run on Qwen3.8-27B (main against this branch, warm-up then base, new, new, base) is green: prompt tokens identical in all 20 rows, decode level to slightly up (2K 48.3 -> 48.9, 64K 39.0 -> 39.8 tok/s, runs overlapping), so the 27B is unchanged by the 9B's path. Main can't load the 9B, so that side rests on your receipts. Our gate passes too, so it lands in the next cycle.

BobClawblaw and others added 7 commits November 2, 2026 09:43
…both prompt kernels checked at a group of 16

Ornith-1.5-122B's attention is 32 query heads over 2 KV heads of 256, a KV
group of 16 where the 27B's is 6. A 16-row tile holds one group of up to 16
query heads, so Geometry.check admits 32 over 2 beside 24 over 4, and the
attention library takes the group and KV-head count as defines
(Q27_G, Q27_KVH, prepended by the runner) where the kernels wrote 6 and 4.
The 27B's kernels compile to the same code: the defines default to its 6
and 4.

tf-qwen27-attention-precision takes --heads/--kv-heads and --simd-prompt,
which runs attention_prompt_sg (the kernel chips before M5 take) on this
GPU's simdgroup matrices. M5 Ultra, 128-row prompt chunks from 0, 1000, 4000,
8100, 11000 and 12160 against float64 (queries in [-1, 1)):

  kernel                 24 over 4                 32 over 2
  attention_prompt_sg    max 3.906e-3, mean 6.14e-5  max 3.906e-3, mean 6.18e-5
  attention_prompt (M5)  max 4.093e-3, mean 6.24e-5  max 4.068e-3, mean 6.28e-5

The 24-over-4 simdgroup figures equal the M3 Ultra ones 993cebd quoted for it (3.906e-3,
6.14e-5).
No Mac before M5 was at hand: the simdgroup kernel ran on the M5's simdgroup
matrices, which those chips share.
…V group of 4

The Qwen3.5 9B's attention is 16 query heads over 4 KV heads of 256, a KV
group of 4. The attention library already takes the group and the KV-head
count as the Q27_G and Q27_KVH defines, so Geometry.check admits 16 over 4
beside 24 over 4 and 32 over 2. Its rank of a two-way split (8 over 2) is
not admitted.

Two tests used 16 over 4 as a shape no shader builds: the geometry test now
expects a group of 4 there and refuses 8 over 2 and 20 over 4 (GQA 5), and
the descriptor test refuses 20 query heads instead of 16.
Before M5, a chip that cannot build the prompt attention kernels falls back
to a library without them. That fallback compiled simd_attention.library's
text without the Q27_G and Q27_KVH defines, so its kernels took the 27B's
group of 6 whatever the checkpoint's geometry. simd_attention.grouped puts
the defines in front of either build, and the runner uses it for both.
mlx-vlm conversions of Qwen3.5 store linear_attn.A_log in f32 (for example
mlx-community/Qwen3.5-9B-MLX-4bit), and the loader asked for bf16 only, so
such a checkpoint refused with BadNativeTensor. The GDN kernels already
read either through the params' weight flags, which the round takes from
the tensor's own dtype, so the loader now takes A_log in bf16 or f32 and
the kernels read it as stored. A bf16 A_log loads as before.
On main a qwen3_5 checkpoint that is not the 2B opens the 27B's Metal engine,
whose supported() admits only the 27B's sizes, so the 9B was refused with
UnqualifiedQwenShape. The 9B is the same design: hidden 4096, intermediate
12288, 32 layers on the 27B's schedule (three Gated DeltaNet layers, then one
full-attention layer), 16 query heads over 4 KV heads of 256, and 16 key and
32 value GDN heads of 128, with the 27B's vocabulary and untied head.

supported() now takes either size from a short table; every other check is
unchanged. Test: the 27B and 9B configs open, and the 4B, the 9B's widths at
the 27B's depth and the 27B with the 9B's GDN heads refuse.
… either checkpoint loads

serve <9B> --drafter <the 27B DFlash2> loaded the 9B weights first and then
refused in the drafter load with TargetBinding: the drafter config was read
only once the target was resident, and the binding check had no name of its
own.

open() now reads both folders' config.json before the target loads. A stock
DFlash2 whose width, vocabulary or tapped depth is another target's is refused
with DraftForAnotherTarget, and a DFlash2 config that is not the stock 27B one
keeps its NotStockDFlash2. The drafter load uses the same check. Test: two
folders holding only config.json; the 27B target takes the 27B DFlash2, and
the 9B target refuses it with DraftForAnotherTarget and a 9B-shaped DFlash2
with NotStockDFlash2.
@ashhart
ashhart merged commit dad04f0 into ashhart:main Nov 2, 2026
5 of 9 checks passed
@ashhart

ashhart commented Nov 2, 2026

Copy link
Copy Markdown
Owner

Thank you @grantoverton. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

@ashhart

ashhart commented Nov 3, 2026

Copy link
Copy Markdown
Owner

Thanks @grantoverton! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants