Skip to content

Core support for oQ in Zig engine - #755

Merged
ashhart merged 6 commits into
ashhart:mainfrom
grantoverton:oqxe/core
Oct 14, 2026
Merged

ashhart merged 6 commits into
ashhart:mainfrom
grantoverton:oqxe/core

Conversation

@grantoverton

@grantoverton grantoverton commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

Rebased onto main c36b536 (1.0.9) on 2026-10-13: no conflicts in this change's code; every commit passes the gate again on the new base (host suite, goldens, lean).

Part of #716: the shared change. No family reads mixed widths on this branch, so no model's kernels, launches or bits change here; the family pull requests follow one at a time and stack on this one. The uniform paths (4/64 for the 27B and GLM-5.3-Flash, 6/32 for Flash Next) keep their kernels and launches: with default settings the 4-, 6- and 8-bit lane and affine_mm sources preprocess (metal -E) to main's text, and the row kernels' outputs at the widths main serves are byte-equal to main's.

What this changes

Six commits on main 0efbd447:

  1. MLX quantization: an override answers its module in any spelling (core/quant). oQ checkpoints write per-module widths as language_model.model.layers.N.…, while loaders ask for model.layers.N.… or the bare path, so an exact-key lookup missed every override and read the tensor at the global width. The table now stores each entry under its canonical module key (no .weight, no known prefix) and looks keys up the same way. Two spellings of one module that declare different widths refuse the table. Test that fails before: a 5-bit down_proj declared under the oQ prefix, read by the loader's spelling; on main it fails the packed shape check.
  2. MLX affine weights widen exactly (core/quant/widen.zig). Codes re-packed at more bits, scales and biases repeated for a smaller group; narrowing and a group that does not divide the source group are refused. No caller here; only Flash Next's hyper-connection down and inject pair uses it, where 0.6.6's QWeights.of widens.
  3. The row projection reads 3- and 5-bit MLX codes and can write into a wider output (row_projection). New branches only; the 2/4/6/8-bit paths are unchanged. callStrided writes one member of a fused stack into the fused rows. tf-row-widths-check (wired in zig/build/core_checks.zig, so build.zig stays under 600 lines) runs it against float64.
  4. Lane projections take every MLX affine width, an output row stride and 16-row members (lane_projection). 2- and 3-bit codes are staged as nibbles and 5-, 6- and 8-bit as bytes, the staging the Python lane matmul uses. 3- and 5-bit groups of 32 are an odd number of words a lane, so they are fetched a word at a time. Layout.y_stride sets the output row pitch in every kernel, including Let the core lane projection run on Macs without tensor units #679's simdgroup-matrix store, which keeps its column guard. Without tensor units the admitted layout is unchanged (6/32 on the plain kernel, Let the core lane projection run on Macs without tensor units #679); every other width still answers TensorUnitsRequired there, now checked at 2, 3 and 5 bits too.
  5. The prompt product affine_mm reads 2-bit MLX codes, exactly. A thread loads its 32 codes of a step as 8 bytes and writes each code into the bf16 tile as itself, with no offset (as at 8 bits), so the group sums are the 4- and 8-bit path's. A first try with the 4-bit "+128" offset stayed within the float64 bound, but a reference with scales 1% off no longer exceeded it at K = 4096, so I dropped it.
  6. tf-affine-ref, the device side of tests: quant-ref, the FP8 and NVFP4 products and experts against float64 references, no Python #633 for the shared prompt products: affine_mm at 2, 4 and 8 bits and simd_qmm at 4/64 against float64, and every row's bits in every window.

Exactness

Same checkpoint, backend, runtime and settings on both sides.

  • Uniform 27B, served, at this branch's tip (fresh build, below). Vontra/Qwen3.8-27B-MLX-4bit @70ae7fa with z-lab/Qwen3.8-27B-DFlash2 (model.safetensors sha256 67fc76d68dc5a941…), greedy:
    • --parallel 2 --context 65536: drafted equals "draft": false 9/9 and resumed equals fresh 2/2 (a 26,130-token prompt resumed at 26,123 cached tokens, same token hash); bench_concurrent.py --levels 1,2 --alone --serial: every alone and serial cell equal, 0 failed.
    • --parallel 4 --context 65536 --prompt-cache-gib 0, six servers alternating main and this branch (below): bench_concurrent.py --levels 1,4 --alone --serial every alone cell (3 at one stream, 12 at four) and serial cell equal, 0 failed, in every run; the code and chat token hashes (2b28aa2108de, cbdd31cd4736) are the same in every run of both binaries.
  • Uniform Flash Next, served (a full-size 6/32 requantization of the oQ3 checkpoint, same layout and names; the Flash Next PR describes it): main against this branch, greedy tokens identical on 6 prompts, plain and drafted, 2 repetitions each.
  • Rows. The row projection gives one row alone the same bits as that row inside 2-, 3-, 8- and 16-row windows, with relu2 and indexed slots too; a strided output equals the unstrided one and writes nothing outside its columns. affine_mm and simd_qmm: every row of a 129-row call equals the same row in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows, which crosses the 16-row fragments, the 32- and 64-row tiles and the gather's switch between them. So prompt rows and decode rows of these products are the same arithmetic.
  • Seams. This branch changes no state, cache or KV code. KV, recurrent state and final logits were not compared at chunk seams directly; the token-level resumed-equals-fresh runs above cover the uniform path.
  • Cancellation, reset and partial acceptance. No request lifecycle, scheduler or state code changes here, but the check ran on the uniform 27B in every served run of both binaries: a fresh "draft": false probe (each bench prompt, greedy, 128 tokens) gives the reference hash and the drafted probe equals it, with partial acceptance seen (accepted < drafted); then after a stream dropped mid-decode, a long prompt dropped mid-prompt, a stream dropped beside a running probe (whose hash is checked too) and a request past the window (refused, the reset path), the drafted and plain probes equal the fresh hashes (f5c9409cfe40, 73e9fd21b09e). 25 steps, pass, in all six runs.
  • Precision. Activations stay bf16 and group sums fp32 at every width. No existing sum changes order; the new widths sum in the existing order.

FP64 per class

M3 Ultra, no tensor units (nax.h's simdgroup-matrix layout), seeded synthetic MLX affine weights. The bound is the fp32 summation bound of the terms' magnitudes plus one bf16 ulp for bf16 outputs. Prompt and decode use identical input rows; since every row's bits are equal in every window, the decode (one row) and prompt (many rows) columns are the same numbers.

row_projection (decode before tensor units; three shapes, every column of every row), worst column per width and group:

Bits g32 |err| (of bound) g64 g128
2 1.20e-2 (0.850) 1.21e-2 (0.825) 7.77e-3 (0.879)
3 1.56e-2 (0.886) 1.56e-2 (0.894) 1.55e-2 (0.861)
4 2.82e-2 (0.865) 2.95e-2 (0.909) 3.04e-2 (0.909)
5 5.96e-2 (0.835) 6.18e-2 (0.856) 6.14e-2 (0.896)
6 1.23e-1 (0.892) 1.23e-1 (0.873) 1.23e-1 (0.940)
8 4.92e-1 (0.848) 4.98e-1 (0.870) 4.97e-1 (0.882)

64,512 values within the bound (worst 0.940, 6-bit group 128: bf16 output rounding); 408,192 bit-equal row/window pairs; 301,056 outputs byte-equal to main's kernels at the widths main serves. A fused stack of eight 48-row members at widths 4, 5, 4, 6, 3, 8, 2, 4 written through the stride: each member's columns equal its own call.

affine_mm and simd_qmm (prompt products), worst column (max |err|, and that over the column's largest |y|):

Product [192, 1024] [256, 4096]
affine_mm q2 g32, bf16 out 2.64e-2 (3.23e-3) 6.22e-2 (3.78e-3)
affine_mm q2 g64 2.87e-2 (3.51e-3) 3.12e-2 (2.47e-3)
affine_mm q4 g32 2.45e-2 (2.65e-3) 3.12e-2 (3.01e-3)
affine_mm q4 g64 1.56e-2 (2.89e-3) 3.12e-2 (2.53e-3)
affine_mm q8 g32 0.496 (3.41e-3) 0.968 (3.63e-3)
affine_mm q8 g64 0.453 (2.87e-3) 0.956 (3.38e-3)

Dense and gathered are equal. simd_qmm 4/64 at [40, 512], [256, 4096] and [6400, 256]: 1.19e-2 (2.57e-3), 3.13e-2 (2.87e-3), 1.54e-2 (3.74e-3), and its twin equals its tiles at 1 and 2 rows wherever twinExact says so. Every product is within 0.4935 of its bound (bf16 outputs 0.20 to 0.49; fp32 outputs at most 0.0081). Negative control, scales 1% off: 2.3 to 29 times the bound for bf16 outputs, 2.7 to 391 for fp32, 44 to 273 for simd_qmm.

Lane projection without tensor units (6/32 only): 14 layouts pass, 1,904 bit-equal row/window pairs, 208 windows equal to fn_lane's recorded bytes, 7,424 CPU float64 values (worst 0.0144 of the bound), 10 refusals; a 48-row member and two strided layouts at 6/32 equal the unstrided rows with the gaps untouched. Every width × group × {64, 48} rows packs MLX codes bit for bit with zero padding (host test). The tensor-unit arithmetic at the new widths needs an M5 (below).

Performance

Mac Studio, M3 Ultra (80-core GPU), 256 GB, macOS 26.4.1 (25E253), Zig 0.17.0. Fresh builds from clean detached worktrees, zig build native tf-qwen27-run -Dcpu=apple_m1; tensorfold-native sha256 6a42d48dc1b08a1d… (main 0efbd447), 1cc42957aea4b6f9… (this branch, 9325793f). Checkpoint Vontra/Qwen3.8-27B-MLX-4bit @70ae7fac with DFlash2. Clients are main's tools/ at 0efbd447 on Python 3.12.14, the server is the native binary; one model process on the machine, the file cache of other checkpoints evicted before each load, every run behind one GPU lock. Every receipt (manifests with binary hashes, decode.json, concurrent.json, prefill records with the exact token arrays and request bodies, /memory snapshots, time -l, load and cleanup timestamps) is in https://gist.github.com/grantoverton/70256846d366db6526a6db3e1b40a4e5 (index in 00-README.md; the host suite summary is test-summary.txt).

No oQ checkpoint has a before-run: main refuses them, and on this branch alone no family reads mixed widths yet. So every receipt here is on the uniform checkpoint, which both binaries load; the family PRs report the oQ checkpoints' speed as served decode only.

Decode and concurrency. tensorfold-native serve <checkpoint> --drafter <DFlash2> --parallel 4 --context 65536 --prompt-cache-gib 0; bench_openai.py … --tokens 256 --reps 3 --temperatures 0 (medians; prompts fibonacci-raw, code, and gpu-chat-no-think, prose, greedy) and bench_concurrent.py … --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial (aggregate tok/s). Six servers in this order: main, branch, branch, main, main, branch (two runs of the 27B branch came between the first pair and the rest).

Run bench_openai code chat concurrent code ×1 code ×4 chat ×1 chat ×4
main 1 95.49 70.00 95.4 114.8 69.8 95.0
branch 1 93.21 65.11 88.5 108.6 68.6 92.5
branch 2 94.07 69.42 94.5 112.1 68.7 94.8
main 2 94.43 68.39 93.1 111.8 68.3 95.0
main 3 94.21 68.38 93.0 112.2 69.8 95.2
branch 3 93.34 69.00 93.1 111.9 68.9 94.7

The first pair moved apart (main 1 fast, branch 1 slow, its chat reps spread 64.8 to 66.0); in the alternated rounds 2 and 3 (branch, main, main, branch) the branch is 0.6% below main on code (93.71 against 94.32) and 1.2% above on chat (69.21 against 68.39), and the concurrent cells agree within 1.6%. The kernels main serves are unchanged here (preprocess-identical sources, byte-equal row outputs), so I read these as run-to-run noise. Flash Next's uniform 6/32 path, main against this branch (an earlier receipt): plain decode 0.998 to 1.000 of main, drafted 0.998 to 1.004, prompt 1.001.

Cold prefill. Exact token arrays of 2,048, 8,192, 32,768 and 65,536 tokens: upstream's public prefill paragraph ("A lighthouse keeper writes in the log as the boats cross the bay.", tools/zig/qwen27_gate.py) repeated and tokenized once without special tokens by the checkpoint's own tokenizer.json, each array the first N ids (every array decodes and re-encodes to itself). POST /v1/completions with the ids as prompt, max_tokens 1, temperature 0, ignore_eos; server --parallel 1 --context 262144 --prompt-cache-gib 0, one request at a time, a fresh KV every measurement (every record has cached_tokens 0 and prompt_tokens equal to N). One 2,048-token warm-up, then 3 rounds over all sizes; wall-clock medians. The arrays, the exact request bodies (with sha256), every request's timings and cached-token counts, and the server's /memory after each are in the gist.

Tokens main tok/s (s) this branch tok/s (s) mlx_lm tok/s (s)
2,048 286.4 (7.15) 286.6 (7.15) 397.2 (5.16)
8,192 258.8 (31.66) 259.1 (31.62) 390.6 (20.97)
32,768 183.0 (179.04) 183.1 (178.92) 349.1 (93.86)
65,536 131.8 (497.15) 131.7 (497.50) 300.1 (218.36)

The three reps of each size are within 0.2% of each other on both binaries, and the branch is within 0.1% of main at every size, so I ran no second, reversed round. Peak physical footprint after the 64k prompts: 40.17 GB for both. mlx_lm takes text prompts (its server accepts no token arrays): the same arrays decoded, each record checking that prompt_tokens equals N and cached_tokens is 0; it prefills faster than the native engine at every size here, 1.39× at 2k and 2.28× at 64k (peak footprint 33.98 GB). 128k and the native window (262,144) were not run on this branch: its kernels for the uniform checkpoint are main's (above), and the 27B pull request records those two cells once each for main and for its own tip.

Standard server. mlx_lm 0.32.0 (mlx 0.32.3, transformers 5.19.0; python -m mlx_lm server --model <checkpoint> --prompt-cache-size 0), same checkpoint, clients and prompts: decode 37.32 tok/s code and 37.20 chat; bench_concurrent 37.2 / 36.6 (code ×1 / ×4) and 37.1 / 36.5 (chat), since mlx_lm serves this model's four streams one after another; peak memory footprint 16.16 GB. mlx_lm returns no token hashes, so its --alone/--serial cells read unverified, and it ignores ignore_eos. Cold prefill on mlx_lm is in the table above (text prompts that tokenize to the same arrays; the record checks prompt_tokens).

Memory

No buffer sizes or load plans change for any uniform checkpoint. Uniform 27B with DFlash2 at --parallel 4 --context 65536 --prompt-cache-gib 0: physical footprint 23.7 to 24.8 GB after load, lifetime peak physical footprint 38.16 GB after the bench and cancel steps, maximum resident set 17.9 GB (time -l), identical for main and this branch in every run. The qwen27 Metal backend reports no device allocations in /memory. RAM class: M3 Ultra, 256 GB. Load and cleanup boundaries are in each run's events.tsv; prefill memory at the 262,144-token context is in the prefill receipts.

Fidelity

Teacher-forced fidelity: n/a for this branch. No family reads the new widths yet, and the uniform paths give identical tokens (above). The family PRs carry fidelity against mlx_lm on the oQ checkpoints.

Host suite and lean check

  • zig build test test-golden -Dcpu=apple_m1 at every commit, last at the tip 9325793f: 158 of 160 steps; goldens 309/319 equal, the 10 declared differences, 0 unexpected. The one failing step is main's own: core.checkpoint_metal's test "a split checkpoint keeps later tensors in the file, packs the rest, and stages the same bytes" logs an error (checkpoint tensor text.weight stays in its file: read it through a Stage): get calls std.log.err before returning the error.DeferredTensor the test expects, and the test runner counts logged errors as failures. It fails on main 0efbd447 too, on the same machine: 158/160 steps, 1505/1513 tests passed, 8 skipped, this failure only. It came with Qwen3.8-27B on native 1.0.5: M6 (GPU generation 18) refused, and memory is too large for 32 GB #643 and shows on macOS.
  • New tests that fail before: the oQ-prefix override fixture; 3- and 5-bit row shapes validate and take whole words; every width has its own pipeline slot; a strided call writes rows at least a matrix wide; 2-, 3- and 5-bit lane formats validate, y_stride reaches the source and below N is refused, a 48-row member validates (cooperative and reg refuse it); affine_mm 2/64 and 2/32, dense and gathered, against exact host sums. widen.zig's tests are new with the file.
  • python3 tools/zig/lean_check.py: 536 problems against main's 539 at 0efbd447, no new findings. Three inherited ones are gone (over-long comments in row_projection.metal and row_projection.zig, shortened where I edited). Inherited findings left in files this branch touches: over-long comments in affine_mm.metal (4, plus one block comment), affine_mm.zig (7), row_projection.metal (3) and row_projection.zig (3). Every touched file is at or under 600 lines.

Platforms

  • Run: Metal on a Mac Studio M3 Ultra, 256 GB (80-core GPU, no tensor units), macOS 26.4.1, Xcode 26.6, Zig 0.17.0.
  • Also run, thanks to @sethforprivacy on an M5 Ultra (tensor units): test-core-lane, tf-affine-ref and tf-row-widths-check PASS.
  • And to @mrv777 on an M1 Max: the same three PASS, tf-row-widths-check's stock byte comparison included.
  • Not run: M1, M2 and M4 (no hardware; the M3 runs the same non-tensor-unit paths); CUDA (no CUDA code changes).
  • Not run: M5. On an M5 the lane projection at the new widths, its 48-row members and output stride, and affine_mm on the tensor units are untested. Commands:
zig build test-core-lane -Dcpu=apple_m1
zig build tf-affine-ref tf-row-widths-check -Dcpu=apple_m1
./zig-out/bin/tf-affine-ref
./zig-out/bin/tf-row-widths-check

I have no M5; @sethforprivacy has since run these (above).

Reproduction settings

Server: tensorfold-native serve <Qwen3.8-27B-MLX-4bit> --drafter <Qwen3.8-27B-DFlash2>, greedy, at --parallel 2 --context 65536 for the exactness checks, --parallel 4 --context 65536 --prompt-cache-gib 0 for the standard clients and cancellation, and --parallel 1 --context 262144 --prompt-cache-gib 0 for cold prefill. One model process on the machine at a time, every load and check program behind one GPU lock, and the file cache of other checkpoints evicted before each load. The machine's iogpu.wired_limit_mb is 240000, which I set as its owner; the 27B's load is far below either limit, and nothing else was raised.

Where the arithmetic differs from 0.6.6

Bit equality with 0.6.6's kernels is not the bar (#716); the float64 bound per product, row invariance and the exactness contract are.

  • 3- and 5-bit rows keep fp32 group sums, as the row kernel's other widths do. 0.6.6's per-family kernels differ (for example GLM's qmv_rows_b rounds group sums to bf16 at 5 and 6 bits); each family PR says which it replaces.
  • 2-bit prompt products use affine_mm with exact codes; 0.6.6 uses MLX's 2-bit quantized_matmul.
  • The lane staging (nibbles for 2 and 3 bits, bytes for 5, 6 and 8) is 0.6.6's.

The open question in #716

#716 asks whether the lane kernel should take a per-member output offset table, so a product grouped over non-adjacent members writes each member's rows in place, or run into scratch and copy. This branch adds neither to core's lane_projection: only the output row stride, which covers a contiguous run. The 27B and GLM-5.3-Flash never group non-adjacent members (each member of a mixed stack is its own product), so they need no map. Flash Next does group them, as 0.6.6's _stacked does, and its PR ships a column map in Flash Next's own lane kernel (FZ_MAP). If you'd rather have that map in core, or scratch and a copy, tell me which and I'll change it.

Checklist

  • 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 (uniform 27B, served, at this branch's tip; uniform Flash Next tokens).
  • No precision traded for speed. No existing sum changes; the new widths sum in the existing order.
  • Prompt processing is no slower than the last release. Cold prefill at 2k to 64k is within 0.1% of main 0efbd447 (above).
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • 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.

Compatibility with #626

Compatible with #626: quant.mlx.Config.overrides still holds the config's own keys, as on main, so #626's Table.width (which reads the map by the raw key, then with language_model. put back) finds every override it found before. The canonical module keys live in a separate modules index, which width(key) consults after the key as given. Two spellings of one module that declare different widths still refuse the table, now through error.UnsupportedQuantization, which quant.detect already returns on main, so the try quant.detect(...) in #626's Config.parse (error set shared.Error || error{UnsupportedQuantization}) compiles unchanged. I checked this by applying #626's diff on top of this branch: zig build test-qwen35 hip-host-test passes (the one skip needs TF_QWEN_DIR), and the full zig build test fails only on main's known core.checkpoint_metal test. A scratch test that reads an 8-bit override through #626's Table.width by its raw key, with and without the stripped prefix, passes here, and it reads 4 bits when the map is keyed canonically. Either PR can land first. Later, #626 could call q.width(key) to pick up the oQ spellings that an exact-key lookup misses.

The receipts above were taken at 9325793f. The only change since is this lookup and its error mapping, with no kernel code touched, re-checked with zig build test test-golden -Dcpu=apple_m1 and the lean check on the new tip.

@grantoverton grantoverton changed the title The shared affine products read every oQ width: override keys by module, exact widen, 2- to 8-bit rows, lanes and affine_mm Core support for oQ mixed-precision checkpoints (2- to 8-bit per layer) Oct 13, 2026
@grantoverton grantoverton changed the title Core support for oQ mixed-precision checkpoints (2- to 8-bit per layer) Core support for oQ mixed-precision checkpoints Oct 13, 2026
@grantoverton grantoverton changed the title Core support for oQ mixed-precision checkpoints Core support for oQ in Zig engine Oct 13, 2026
@sethforprivacy

Copy link
Copy Markdown
Contributor

Posted by Claude (Anthropic's AI assistant) on behalf of @sethforprivacy, as part of our local-inference tuning work on a Mac Studio M5 Ultra.

@grantoverton the M5 checks for this PR at 008abc0 all pass. Mac Studio M5 Ultra (80-core GPU, tensor units), 256 GB, macOS 27.0 (26A428), Xcode 27.0, Zig 0.17.0. Each branch was built clean from its PR head, zig build ... -Dcpu=apple_m1.

zig build test-core-lane            PASS 196 layouts, 26656 bit-equal row/window pairs, 208 frozen fn_lane byte-equal windows, 159744 CPU-reference values; max fp32 error 0.001635909, max roundoff-bound fraction 0.019091
                                    (and every b*g* n48/n64/n96 stride case: output rows N apart equal the unstrided rows, gaps untouched)
tf-affine-ref                       device Apple M5 Ultra, tensor units true
                                    PASS every product within 0.4935 of its float64 bound, every row's bits equal at every window
tf-row-widths-check                 PASS 64512 values within the fp32-sum bound (worst err/bound 0.9400), 408192 bit-equal pairs
                                    (no stock source given: the byte comparison was not run)

@mrv777

mrv777 commented Oct 13, 2026

Copy link
Copy Markdown
Contributor

M1 Max check at 008abc04: 32 GPU cores, 32 GB RAM, macOS 26.7, Zig 0.17.0 (-Dcpu=apple_m1), no tensor units.

test-core-lane      PASS 14 layouts; 1,904 bit-equal row/window pairs; 208 frozen windows; 7,424 CPU-reference values; 10 expected refusals
tf-row-widths-check PASS 64,512 values; worst err/bound 0.9400; 408,192 bit-equal pairs; 301,056 stock outputs byte-equal
tf-affine-ref       PASS every product within 0.4935 of its float64 bound; every row's bits equal at every window

The stock comparison used main at 0efbd447. Core-lane and row-width checks also pass Metal API + GPU Shader Validation. Affine passes API validation; Shader Validation inflates an unchanged simd_qmm variant from 19,456 to 38,912 bytes of threadgroup memory, exceeding M1's 32,768-byte limit.

test test-golden: 1,526/1,535 passed, 8 skipped, the existing GLM sparse-attention failure and checkpoint logged-error diagnostic remain. Goldens: 309/319 equal, 10 known differences, zero unexpected differences. Raw logs kept local.

oQ checkpoints write per-module widths as language_model.model.layers.N..., while loaders ask for model.layers.N...
or the bare path, so an exact-key lookup missed every override and read the tensor at the global width. The table
keeps each entry under the config's own key, as before, and also indexes it under its canonical module key (no
.weight suffix, no known prefix, the rule the 27B and Flash Next families use); `width` takes the key as given, then
its module. Two spellings of one module that declare different widths refuse the table as unsupported, the error
`detect` already returns, so its callers' error sets are unchanged; the same width twice is fine.

Test that fails before: the fixture reads a 5-bit down_proj declared under the oQ prefix by the loader spelling; on
main it fails the packed shape check (read at the global width). Also: every prefix form and a trailing .weight name
one module; aliases must agree; an override in any spelling answers every spelling and stays under its own key.

Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens
equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint
test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no
new findings.
… tables repeated for a smaller group

A projection of mixed-width members (the Flash Next hyper-connection down and inject pair, where 0.6.6's QWeights.of
widens) reads as one matrix at their common width, the most bits and the smallest group. Narrowing and a group that
does not divide the source group are refused. Used only where 0.6.6 widens; no caller on this commit.

Tests (widen.zig, new): every source width and every wider target dequantize to the same fp64 values, bit for bit;
a stack of experts widens row for row; the same width comes back as it is; narrowing bits, growing groups and
widths the kernels do not run are refused; the common width takes the most bits and the smallest group; fused
members keep each one's values in order.

Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens
equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint
test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no
new findings.
…wider output

Chips without tensor units now decode every MLX affine width (2, 3, 4, 5, 6, 8 bits) in the row kernel: 3- and 5-bit
rows unpack eight codes from three or five bytes, low code first, as MLX packs them. The 2-, 4-, 6- and 8-bit paths
are unchanged. A 3- or 5-bit row must end on a word, as the checkpoint stores it.

callStrided launches the row kernel with an output row stride, so one member of a fused stack (one run of equal
width) writes its columns straight into the stack output. The default call, its dims and its kernels are unchanged.

tf-row-widths-check (zig/build/core_checks.zig) runs it on the GPU: each width at groups of 32, 64 and 128 against
an fp64 product from the MLX reference dequantization, within the fp32-sum bound; one row alone against the same
row in 2-, 3-, 8- and 16-row windows, relu2, indexed slots and a strided output, bit for bit; a fused stack of
48-column members at mixed widths (4-bit among them, group 64) written through the stride at windows up to 48 rows,
each member as its own call; and, given the unchanged kernel source, the widths it already served byte for byte
against it.

Tests that fail before: "shapes the row kernels serve" now validates 3 and 5 bits (main refuses them),
"3- and 5-bit rows take whole words of codes", "every width has its own pipeline slot" and "a strided call writes
rows at least a matrix wide" (no callStrided on main).

tf-row-widths-check on an M3 Ultra (no tensor units), given main's row_projection.metal: PASS 64512 values within
the fp32-sum bound (worst err/bound 0.9400), 408192 bit-equal pairs, 301056 outputs byte-equal to main's kernels
at the widths main serves (3 and 5 bits have no main kernel to compare).

Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens
equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint
test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no
new findings.
…d 16-row members

Format accepts 2-, 3- and 5-bit codes beside 4, 6 and 8, in groups of 32, 64 or 128. The non-4-bit path widens 2- and
3-bit codes to nibbles for the bf16 x uint4 op and 5-, 6- and 8-bit codes to bytes for the uint8 op, the staging the
Python lane matmul uses, so each width keeps that arithmetic. 3- and 5-bit groups of 32 are an odd number of words a
lane; they are fetched a word at a time instead of in pairs, which would drop the last word.

Layout.y_stride (default N) sets the elements between output rows in every kernel (tf_lane, tf_lane_pair, tf_lane_reg),
so a member of a fused stack writes its columns of the fused rows; the output bound checks follow it. The store of
the simdgroup-matrix layout (Macs without tensor units, ashhart#679) takes the same stride and keeps its column guard.

N may be a multiple of 16 on the plain kernel: the last 32-column tile is packed with zero codes past N, its metadata
reads and stores are already guarded by N, and weightWords counts the padded tiles (sourceWords is MLX's W).
Cooperative and reg layouts still need N % 32 == 0. encodeSums checks its sums buffer by the group count, as ashhart#698
sizes the dispatch.

Without tensor units the admitted layout is unchanged: 6-bit group 32 on the plain kernel (ashhart#679); every other width
still answers TensorUnitsRequired there, now checked at 2, 3 and 5 bits too.

With the defaults the 4-, 6- and 8-bit kernels preprocess (metal -E) to the same text as main's for every group,
schedule, output, prepared-sums/reg layout and both fragment layouts tried (148 variants), so the uniform paths
keep their kernels. Host tests that fail before: 2, 3 and 5 bits validate; y_stride reaches the source and below N
is refused; a 48-row member validates (cooperative and reg refuse it); every width x group x {64, 48} rows packs
MLX codes bit for bit with zero padding. The device program covers every width, 48-row members and strided output
on tensor units; without them (M3 Ultra): PASS 14 layouts, 1904 bit-equal row/window pairs, 208 windows equal to
fn_lane's recorded bytes, 7424 CPU values (max roundoff-bound fraction 0.0144), 10 refusals, and a 48-row member
and two strided layouts at 6/32 equal the unstrided rows with the gaps untouched. fn_lane's recorded windows are
compared only on the layouts they were recorded at.

Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens
equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint
test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no
new findings.
affine_mm took 4- and 8-bit weights only, so a 2-bit matrix (oQ2's routed experts) had no tiled prompt product.
At 2 bits a thread loads its 32 codes of a step as 8 bytes and writes each code into the bf16 tile as itself, with
no offset (as at 8 bits; 0..3 is exact in bf16), so the group sums are the 4- and 8-bit path's. A first try with the
4-bit "+128" offset stayed within the float64 bound, but a reference with scales 1% off no longer exceeded it at
K = 4096, so the codes enter exactly. The 4- and 8-bit sources preprocess as before.

Test that fails before: "dense and gathered, each format, against exact sums on the host" adds 2/64 and 2/32, dense
and gathered; source() refused 2 bits.

The 4- and 8-bit sources preprocess (metal -E) to main's text at both groups, both outputs and both fragment
layouts (16 variants). tf-affine-ref (next commit) holds the 2-bit products to float64.

Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens
equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint
test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no
new findings.
…y row's bits in every window

The device side of ashhart#633's quant-ref for the shared prompt products, item 5 of ashhart#716. tf-affine-ref takes no
arguments and needs no model: seeded synthetic codes, scales, biases and inputs, packed as MLX affine. The engine
module now names core/affine_mm so a check program can reach it.

- affine_mm at 2, 4 and 8 bits, groups 32 and 64, bf16 and fp32 out, dense and as a one-expert gather, at
  [192, 1024] and [256, 4096].
- simd_qmm (4-bit, group 64) as tiles and as the scalar twin at 1 and 2 rows, at [40, 512], [256, 4096] and
  [6400, 256] (K chunks of 32, 16 and 8).

The reference dequantizes the same codes in float64 and multiplies in float64. Each product prints its worst
column (max abs error, and that error over the column's largest |y|) and its largest error as a fraction of the
bound: (group + K / group + 8) fp32 roundings of the terms' magnitudes as the kernels sum them (4-bit affine_mm
codes plus 128), plus one bf16 ulp for bf16 outputs. A reference with its scales 1% off must exceed the bound.

Then every row of a 129-row call must equal, byte for byte, the same row computed in windows of 1, 2, 3, 8, 16,
17, 33, 64 and 65 rows (every full window over the 129 rows, so each row is seen at every position), which
crosses the 16-row fragments, the 32- and 64-row tiles and the gather's switch between them. The twin must equal
the tiles wherever twinExact says it does, since only then do callers take it.

Receipts, M3 Ultra (no tensor units: nax.h's simdgroup-matrix layout), on main 0efbd44:

- every product within 0.4935 of its bound; every row's bits equal at every window, both kernels, every format;
  the twin equals the tiles at 1 and 2 rows on all three shapes, as twinExact says.
- worst columns, bf16 out (dense and gather equal), at [192, 1024] and [256, 4096]: affine_mm q2 g32 2.64e-2 (rel
  3.23e-3) and 6.22e-2 (3.78e-3), q2 g64 2.87e-2 (3.51e-3) and 3.12e-2 (2.47e-3), q4 g32 2.45e-2 (2.65e-3) and
  3.12e-2 (3.01e-3), q4 g64 1.56e-2 (2.89e-3) and 3.12e-2 (2.53e-3), q8 g32 0.496 (3.41e-3) and 0.968 (3.63e-3),
  q8 g64 0.453 (2.87e-3) and 0.956 (3.38e-3); simd_qmm 1.19e-2 (2.57e-3), 3.13e-2 (2.87e-3) and 1.54e-2 (3.74e-3).
  All under 2^-8 of the column's largest |y|; 0.20 to 0.49 of the bound.
- fp32 out: at most 0.0081 of the bound (q2 1.5e-6 to 7.7e-6, rel 1.9e-7 to 6.0e-7).
- scales 1% off: 2.3 to 29 times the bound for bf16 outputs, 2.7 to 391 for fp32, 44 to 273 for simd_qmm.

Not run: an M5, where affine_mm takes the tensor units; `zig build tf-affine-ref && ./zig-out/bin/tf-affine-ref`.

The fp32 bound is the worst case for any summation order, so a changed order still passes; the measured errors
sit two orders below it, and the bit checks are what pin the order.

Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens
equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint
test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no
new findings.
@ashhart
ashhart merged commit b0fa154 into ashhart:main Oct 14, 2026
4 checks passed
ThinkOffApp pushed a commit to ThinkOffApp/TensorFold that referenced this pull request Oct 14, 2026
ThinkOffApp pushed a commit to ThinkOffApp/TensorFold that referenced this pull request Oct 14, 2026
The contributor's six commits landed above with the same code; this keeps the PR's head in the history with no change to the tree.
@ashhart

ashhart commented Oct 14, 2026

Copy link
Copy Markdown
Owner

Thank you, @grantoverton. The shared Metal kernels and core/quant read every MLX width. At the widths 1.0.9 serves, the sources and the outputs stay 1.0.9's.

Your commits are in 1.0.10 under your name (6a6f75d1d, 8db1648fb, 4c00d57a0, 5fea98108, 63b5b5faf, dcc0efd3f). The CHANGELOG line is db657a491. Your head 6d2a62aa7 is merged at b0fa1540e with no change to the tree. The checks below ran at 008abc04f. widen lands beside the fp8 and nvfp4 imports 1.0.9 added. Both are kept.

The 4-, 6- and 8-bit lane and affine_mm sources preprocess to 1.0.9's text (168 layouts and formats). The row kernels give 1.0.9's bytes at 2, 4, 6 and 8 bits (301,056 outputs on an M3 Ultra and on an M5 Max). The lane kernels' outputs are 1.0.9's on every layout 1.0.9 runs (115 on an M5 Max, 13 on an M3 Ultra). Nemotron on an M3 Ultra, which compiles these row kernels, gives 1.0.9's tokens drafted, resumed and concurrent. @sethforprivacy ran the checks with test-core-lane on an M5 Ultra and @mrv777 on an M1 Max, both at 008abc04f, all passing. Thank you both.

@ashhart ashhart mentioned this pull request Oct 14, 2026
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.

4 participants