[Nipale-ai]

What this changes

On an Ultra with tensor units (M5 Ultra), a Qwen3.8-27B prompt chunk of more than 128 rows runs each 4-bit projection through the reg lane kernel (tf_lane_reg) in one dispatch instead of 128-row dispatches, so the chunk sweeps each weight matrix once instead of once per 128 rows. Replies are token-identical; cold prefill throughput is about 10% to 12% higher from 2k to 64k tokens (time to first token 9% to 11% lower); decoded rows, verify rows and other chips are unchanged.

• Before: a prompt chunk (up to 2,048 rows on M5) runs each reg projection in 128-row dispatches (lane.reg_rows). Each dispatch sweeps the whole weight matrix, so a 2,048-row chunk reads gate/up's 89.1 MB of codes and 11.1 MB of scales and biases 16 times.
• lane_projection.Layout.rows (default reg_rows) sets the rows one reg encode may take, from reg_rows to max_reg_rows (2,048) in steps of 16. It sizes the sums scratch and the encode bound only; the kernel source is the same text (unit test).
• quant_gpu.longPrompt is the gate: the Metal device name ends in " Ultra" and tensorUnits() holds. model.zig sets Kernels.long_prompt once at load.
• quant_gpu.longKey: on such a chip, a prompt chunk of more than 128 rows takes its own lane key (Key.long), built with rows = max_reg_rows. Like every lane key, it compiles its pipelines on first use. Decoded rows, verify rows and prompt chunks of 128 rows or fewer keep their current keys and 128-row sums scratch (unit test).
• The row loop steps by the projection's max_rows: 2,048 for a long key, and 128 for other reg keys and 16 for non-reg keys, as before.

Every threadgroup computes the same output tile with the same K split, the same tensor-op calls, the same epilogue and the same reduction as before. Only the number of row blocks on a dispatch's y changes, so every output word is the same.

Other chips (M5 Max, M5 Pro, and chips without tensor units) build the same keys and layouts, take the same row steps and encode the same dispatches as main. I measured only on an M5 Ultra and have no M5 Max, so the gate keeps them on 128-row dispatches.

Receipt

• Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision:

  • Before is main 94b90c9; after is this branch, one commit on 94b90c9. Both built with zig build native -Dcpu=apple_m1 -Dversion=1.0.0 (Zig 0.17.0). No MLX or PyTorch in the served path; the clients ran on Python 3.13.14.
  • Mac Studio, Apple M5 Ultra (80 GPU cores), 256 GB, macOS 27.0.1 (26A434), Xcode 27.1 (Metal 32023.921).
  • Checkpoint: a local copy of TensorFold/Qwen3.8-27B-MLX-4bit (its model card names that repo). I did not record its revision; shard SHA-256 6cc1508e96fb, 83f2a20ca805, 31b8c91ef899; tokenizer.json 06b9509352d2.
  • Server for prompt and decode speed: tensorfold-native serve MODEL --name bench --context 66560 --parallel 1 --prompt-cache-gib 0 --no-update-check.

• Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false:

  • tools/bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial, server --context 16384 --parallel 4 --prompt-cache-gib 0, on this branch and on main: every concurrent reply's token_sha256 equals its request alone, 3/3 at one stream and 12/12 at four streams in each cell, and every alone run equals "draft": false. The cells are bench_openai's code and chat prompts, and a --prompts cell with two docs prompts of 1,885 and 8,004 tokens that take the long dispatches. The server accepted drafted tokens in the four-stream runs.
  • Served bytes: every cold prefill reply's token_sha (64 greedy tokens) is the same before and after at every size: 2k f4a149d4ab4a, 8k 5ca261bef56b, 32k e7beb06f7788, 64k d69cb44b4b30.
  • Kernel bytes: tf-core-lane-reg-check gains a long check. It packs random 4-bit codes, scales and biases at the 27B's seven prompt shapes with their prompt K splits (N x K, split: 34816 x 5120, 1; 5120 x 17408, 2; 10240 x 5120, 1; 12288 x 5120, 2; 6240 x 5120, 4; 5120 x 6144, 2; 2048 x 5120, 8) and at the two partial layouts of the opt-in Neural Engine share with TF_ANE_COLUMNS=3072 (gate/up without the first 96 tiles of each half; down over input groups 48 to 271, into fp32). It runs 129, 1,065 and 2,048 rows with own and with producer sums, and requires every output byte of one long dispatch to equal the 128-row dispatches: LONG CHECK PASS cases=54. The existing check still prints REG CHECK PASS cases=352.

• Decode speed before and after, from tools/bench_openai.py:

  • tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, a fresh server each, alternated A B B A (A is main, B is this branch), median decode in tok/s:
    • fibonacci-raw: 49.90, 49.62 before; 49.71, 49.84 after; token_sha a10c9b848ed9 in every run.
    • gpu-chat-no-think: 49.69, 49.79 before; 49.63, 49.80 after; token_sha a06fab348fbe in every run.
  • The 64-token replies of the cold prefill runs below decode at the same speed (median tok/s, before/after): 2k 49.57/49.59, 8k 47.62/47.60, 32k 44.60/44.68, 64k 39.39/39.67.

• Prompt speed before and after:

  • Cold ttft_s from tools/bench_cache.py --prompt-file ids-N.json --expected-prompt-tokens N --tokens 64 --repeats 2 --streams 1 --require-cold, four cold requests a bench_cache run (with the prompt cache off, the repeats and the concurrent request are cold too).
  • The server above, with a fresh server for each session, and each session served the prompt sizes in ascending order (2k, 8k, 32k, 64k). Every request reported cached_tokens = 0, and every bench_cache run qualified.
  • The prompts are fixed token arrays: the first N IDs of the repository's docs/*.md at 94b90c9, joined in case-insensitive name order (WINDOWS.md last), each file followed by two empty lines, tokenized by the checkpoint through /v1/tokenize with add_special_tokens: false. The arrays, the request body and the raw bench_cache JSON are attached: receipts-v3.zip.
  • Sessions alternated B A A B two times (B is this branch, A is main), with 60 s idle between servers, giving 16 cold samples per cell. Medians with [min-max]; the change and tok/s (N / median ttft_s) follow from the rounded medians:

  prompt	main ttft (s)	branch ttft (s)	change	main tok/s	branch tok/s
  2048	1.190 [1.183-1.205]	1.076 [1.075-1.101]	-9.6%	1721	1903
  8192	4.879 [4.854-4.890]	4.415 [4.411-4.429]	-9.5%	1679	1855
  32768	21.602 [21.428-22.431]	19.630 [19.615-20.149]	-9.1%	1517	1669
  65536	55.854 [51.674-57.874]	49.910 [44.836-52.144]	-10.6%	1173	1313

  • Each adjacent before/after pair of sessions agrees (main ttft_s over branch ttft_s, per-session medians): 1.101-1.108 at 2k, 1.100-1.106 at 8k, 1.094-1.103 at 32k and 1.111-1.181 at 64k. At 64k the first request of each run was 1.6-4.4 s faster than the three after it on both sides, so its samples and pairs spread more.
  • The recipe's M5 Ultra prefill figures (docs/recipes/qwen3.8-27b.md: 1325, 1269 and 1171 tok/s at 2k, 8k and 32k) come from its own method (100 greedy tokens, its own docs prompts); main measures higher here with the method above, and this change leaves the recipe alone.

• Memory:

  • Peak physical footprint of the server process (macOS footprint, phys_footprint_peak), read after each session's 64k run: 30.910-30.912 GiB on main and 30.899-30.904 GiB on this branch, four sessions each; after load, before any request, 30.535-30.536 GiB on both. 256 GB RAM, --context 66560.
  • Each long key adds a sums scratch of (K / 64) x 2,048 x 4 bytes: 2.1 MiB for down (K 17,408), 0.75 MiB at K 6,144 and 0.63 MiB at K 5,120. Decoded rows keep their 128-row keys. Nothing changes on other chips.
  • The branch's peaks above are about 0.01 GiB below main's, although the long keys add scratch; I did not trace why.

• Tests run, with pass, fail and skip counts:

  • zig build test test-golden -Dcpu=apple_m1 --summary all: this branch 8713/8766 passed, 53 skipped, 0 failed (from an empty local cache); main 8707/8760 passed, 53 skipped, 0 failed. Three new tests (the layout-rows test, the longPrompt test and the longKey test), each counted in two test binaries (+6).
  • zig build tf-core-lane-reg-check -Dcpu=apple_m1 and its run: both checks pass (above).
  • CI's no-GPU steps on this branch: zig build test-server-security test-pull-security 26/26 passed; zig build fuzz-cortheon 17/17 passed; zig build check-release-paths 59/59 steps succeeded; zig build test-kernel-contract 141/141 passed; Python flashnext_batch_serve_test.py 8/8 passed and build_options_test.py 5/5 passed (with ZIG set); zig build dev-tools -Dcpu=apple_m1 and zig build dev-tools -Doptimize=Debug -Dcpu=apple_m1: exit 0.
  • tools/zig/lean_check.py: 424 problems on main and on this branch; tools/zig/lean_new.py: 0 new.
  • zig fmt --check on the five changed files and git diff --check: clean.
  • Kernel contract: zig build kernel-contract -Dcpu=apple_m1 stops on this Mac before the gates with UnrecordedRotaryCompiler (its rotary probes need a tf-rotary-sweep record for this Metal compiler), on main and on this branch alike. tf-kernel-contract --only lane gives the same result on both: SUMMARY 6 kernels: 6 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems.

• Not run, and why:

  • The previous release (1.1.9) and a standard server (mlx_lm): not measured. There is no 1.1.9 build on this Mac, and the before is main 94b90c9, which already carries the changes since 1.1.9.
  • 128k and the native window: not measured; a 2k-to-64k session already takes about five to six minutes.
  • A resumed prompt: with --prompt-cache-gib 8, an 8k request and then its 32k extension reported cached_tokens = 0, and the server logged no cache hit, on main and on this branch alike. So I have no resume receipt; the replies' token_sha equal the cold runs'.
  • The checkpoint revision: not recorded at download; the SHA-256 values above identify the files.
  • Other platforms: M1 to M4, M5 Max and M5 Pro, not run; I ran this on the M5 Ultra Mac Studio only. The gate keeps them on the current path, and the host suite covers the gate (longPrompt test). CUDA: not run; this change touches no CUDA code.
  • CI steps I could not run as CI does: the NFKD table check needs the Unicode 16.0.0 of Python 3.14, and this Mac has Python 3.13 (Unicode 15.1.0), so the check refuses and its two version-bound tests skip (the other three pass); the Windows job needs a Windows machine; check-release-paths ran on macOS, while pull requests check it on Linux x86_64.

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. Drafted output and concurrent: yes, bench_concurrent --alone --serial above; a resumed prompt: no receipt (see "Not run, and why").
• No precision traded for speed. The bits do not change (byte check and token hashes above).
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. No kernel source changes; the Metal lane SUMMARY line is above.
• Prompt processing is no slower than the last release. Faster than main at every size measured; not measured against 1.1.9 (see "Not run, and why").
• 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. The unit tests need names this change adds, so they do not build on main, and they need no GPU. tf-core-lane-reg-check is a device tool and, as before, stops with TensorUnitsRequired without tensor units.
• 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.

Related

• perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648 bands the HIP prompt GEMM's blocks so row blocks share weights; this is the Metal reg kernel's counterpart for the 27B's prompt chunks, without a change to the launch order inside a dispatch.
• qwen27 DFlash: compile the drafter head's lane pipeline at load #1734 moves the lane key and the lane.Projection.init call in quant_gpu.zig into laneEntry; this change adds long to that key and rows to that call, so the two branches conflict there. If qwen27 DFlash: compile the drafter head's lane pipeline at load #1734 lands first, I will rebase this branch onto it.

Suggested changelog line

• Qwen3.8-27B on an M5 Ultra runs a prompt chunk's 4-bit projections in one dispatch each instead of 128-row dispatches, so a chunk sweeps each weight matrix once. Replies are token-identical. Cold prefill throughput is about 10% to 12% higher from 2k to 64k tokens; decode and other chips are unchanged.

[ashhart]

Thanks @Nipale-ai. Sweeping each weight matrix once per chunk instead of once per 128 rows is a clean idea, and gating it to Ultras with tensor units keeps other chips untouched. We've rebased it onto main (clean) with a changelog line, and it's running now on an M5 Ultra: Qwen3.8-27B MLX 4-bit with DFlash2, main against this, alternating fresh servers. That checks prompt tokens, cold prefill from 2k to 64k, decode after each prompt, and the stream sweep at 1, 4 and 16 streams. We'll post the numbers here. It would land after 1.2.0, which is being cut today.

[ashhart]

It holds up, @Nipale-ai. M5 Ultra, Qwen3.8-27B MLX 4-bit with DFlash2, main against this, alternating fresh servers (main, this, this, main):

cold prefill	main	this
2,048	1,699.9	1,877.0	+10.4%
8,192	1,659.8	1,829.3	+10.2%
16,384	1,605.2	1,767.5	+10.1%
32,768	1,501.3	1,650.4	+9.9%
65,536	1,323.9	1,450.2	+9.5%

• Prompt tokens equal to main (0 of 20 rows differ).
• Decode after each prompt is level.
• The stream sweep is +0.4% at one stream, +0.3% at sixteen, and -2.4% at four. The four-stream cell is inside run-to-run spread by our test, and it has moved like that on this box for unchanged builds all night.

It's ready, and lands right after 1.2.0. Thank you.

[ashhart]

A correction to my last comment, @Nipale-ai: I called this ready too soon. Reading every cell rather than the summary, one cell is slower in both runs of this PR and level in both runs of main. M5 Ultra, 27B + DFlash2, four greedy chat streams: 309.0 and 305.7 tok/s on main, 288.1 and 286.3 here (-6%). The other three cells at four streams are level or up, and the prefill gain stands (+9.5 to +10.4%). So we're holding it rather than landing it now.

We see the same -6% in that one cell from another change that alters how the 27B's Metal pipelines are keyed, and not from changes that don't. So our first guess is a pipeline built on first use inside that cell, perhaps the new Key.long entry for a chunk over 128 rows. That's a guess, not a finding. Could you check on your M5 Max whether a four-stream greedy chat round after a long prompt compiles a new lane pipeline (timing the first longKey lookup, or warming it at load) and whether the cell recovers if it does? We'll re-run it as soon as you push.

[ashhart]

Update on the four-stream cell I flagged: it isn't this PR. We ran main against main with only a comment changed, on the same M5 Ultra with the same harness. Greedy chat at four streams read 309.2/312.3 tok/s on the base binary and 286.2/287.7 on the new one, the same split as this PR's 309.0/305.7 against 288.1/286.3. So that cell measures our harness's second binary, not your change. Everything else on this PR is clean: tokens equal on all 20 prompt rows, prompt speed +9.5% to +10.4% from 2k to 64k, decode level, 1 and 16 streams level. It's queued to land now. Thank you @Nipale-ai, and sorry for the false alarm.

[ashhart]

Thank you @Nipale-ai. 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.

[Nipale-ai]

Thank you for checking every cell and landing it, @ashhart.

[ashhart]

Thanks @Nipale-ai! One dispatch per 4-bit projection for Qwen3.8-27B prompt chunks on M5 Ultra is in 1.2.1: https://github.com/ashhart/TensorFold/releases/tag/v1.2.1