[kieranscotney]

What this changes

The DFlash drafter's head is the target's quantized lm_head trimmed to the draft vocabulary (dflash/runtime_head.zig).
It runs through the target's quant_gpu.Kernels.quant, which compiles one specialised lane-projection pipeline from
source the first time each (n, k, slices, ...) key is seen. No target layer has the trimmed head's n, so its library
was compiled inside the first draft round after a prompt.

• quant_gpu.zig: the lookup in quantOne (key, ranges, getOrPut, lane.Projection.init, joined) moves unchanged
  into laneEntry(p, sk, skip, cut). quantOne calls it with the same arguments. A new prepareDecode(p) calls
  laneEntry(p, p.slices, 0, .{ 0, 0 }), the key a decode-row quant of p builds. It encodes nothing and returns
  early for stacked or raw-layout weights and during prompt passes.
• dflash/runtime_backend.zig: Backend.init calls target.kernels.prepareDecode(head.linear) after Head.init.

The first head launch now finds the same pipeline object it would have built (same key, source and compile options).
Dispatches, buffers and arithmetic are unchanged.

No issue exists for this, so there is no ATC claim; happy to open one and claim it if you prefer. atc overlap lists
akol1's open stack #1433-#1437 as sharing quant_gpu.zig. Their hunks touch quantOne's split and raw paths, not the
lane lookup moved here, and this commit cherry-picks cleanly onto #1437.

Receipt

• Environment: TensorFold main at 94b90c9db; Apple M5 Max, 128 GB; Zig 0.17.0; Qwen3.8-27B MLX 4-bit target
  (revision 22d8d538) with the stock DFlash2 drafter at depth 1. Measurements came from the same change applied to the
  Yukon mlx.fast Qwen3.8-27B engine copy, whose quant_gpu lookup and drafter head match main's.
• Exactness: same pipeline objects, so no arithmetic changes. The exact-token correctness gate (64 steps of a
  1,065-token public long-copy prompt against the plain target's golden tokens) passed on every run.
• Decode speed: the first draft round after a 1,065-token prompt fell from 182-202 ms to about 29 ms under CPU load,
  and by about 25-40 ms on an idle machine. Encode of the head call was 179.7 ms of it, with a 1.8 ms GPU wait.
  Later rounds stay at about 4 ms. Over a 128-token decode window: +0.6% to +1.8% across seven alternating pairs on an
  idle machine, +4.7% under load.
• Prompt speed: n/a. The change only moves a pipeline compile to model load; prompt passes do not use this key.
• Tests run: zig build test -Dcpu=apple_m1 on this branch exits 0; summary 7 passed; 0 skipped; 0 failed. The
  failed command lines in the log come from tests that exercise failure paths, and main prints the same kind.
• Not run: CUDA and HIP (no hardware); Metal M1-M4 (only an M5 Max here); tools/bench_openai.py and
  tools/bench_concurrent.py (the numbers above come from the Yukon harness on the same engine code).

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
  equals the plain target's golden tokens on the exactness gate above; resumed and concurrent cases were not run. The
  change builds the same pipeline object earlier, so it cannot affect them.)
• No precision traded for speed. No sums change.
• No new or changed GPU kernel.
• Prompt processing is no slower than the last release.
• Platforms: Metal on M5 Max only. Not run on M1-M4, CUDA or HIP.
• New tests fail before the change, pass after it, and skip cleanly without their dependency. (No new test: the
  change is a compile-time move with no observable output difference.)
• 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]

Thanks @kieranscotney, another clean one. Moving a compile from the first draft round to load, with the same key, source and options, is exactly the kind of change we like. No ATC claim is needed for a fix this size. We've rebased it onto main (clean) with a changelog line, and are running the 27B with DFlash2 on an M5, main against this: tokens equal, decode, and the startup time, since the compile now lands there. An M3 Ultra run follows when that box is free again.

[ashhart]

@kieranscotney, our M5 Ultra run (Qwen3.8-27B MLX 4-bit with DFlash2, main against this, fresh servers, order main, this, this, main) found something we didn't expect from a compile move. Tokens are equal to main (0 of 20 prompt rows differ), but decode after some prompts is slower in both runs of this PR, and level in both runs of main:

decode tok/s after the prompt	main (run 1, run 4)	this (run 2, run 3)
2,048 tokens, first prompt	151.7, 152.2	143.0, 143.1 (-6%)
2,048 tokens, from offset 5,000	105.0, 106.1	99.7, 101.8 (-4%)
16,384 tokens, from offset 5,000	192.0, 191.6	173.7, 174.4 (-9%)
16,384 tokens, other offsets and repeats	level	level
8,192 tokens	101.6	103.2 (+1.6%)

The stream sweep is level at one stream (-0.4%) and -1.9% at four.

Since the pipeline should be the same object, our guess is that the one built at load differs from the one the first head launch would have built in some way the key doesn't capture. That could be a compile option, the function constants, or the library it joins. Or the early prepareDecode changes which pipeline later lookups find. Could you compare the two pipeline objects (their descriptors, or maxTotalThreadsPerThreadgroup and threadExecutionWidth) on your M5 Max, built at load against built lazily? A regression stops it landing as is, but the first-round saving is worth having, so we'd like to find it.

[ashhart]

Update from our M5 Ultra, @kieranscotney. We re-ran your change against main 94b90c9 in both orders (main, PR, PR, main and then PR, main, main, PR), fresh process per run, with the same checkpoint, drafter and seeds. All 768 replies are token-exact. Decode after a 16k prompt from offset 5000 is slower in both orders, -8.6% and -8.4%, with every PR run below every main run; after the first 2k prompt it's -5.2% in one order; 2k from offset 5000 doesn't reproduce. A host review found the same head key, generated source and pipeline inputs on both sides, so we can't point to a cause in your code yet. We've also just found that on this machine a comment-only change can move one cell by about 7%, because the binary's layout shifts. So we're checking whether that explains this before asking you for anything. It stays parked until then.

[ashhart]

Good news, @kieranscotney: we found the cause of the M5 slowdown, and it isn't your change. The engine calibrates its draft pricing once per executable and keeps it, and the two builds calibrated slightly different cost curves. With that, your build did one extra draft round at 16k (11 against main's 10) for the same exact reply. When we gave both executables main's retained prices, the two were level: -0.6% at 16k from offset 5000 and +0.2% on the first 2k, with the same rounds and proposals. That calibration spread is ours to fix, and it's in hand separately.

One thing before this can land: #1738 (just merged) rewrote the same projection path in quant_gpu.zig. The key gained a long field for long prompts, and the step now comes from the entry's max_rows. So your laneEntry needs to build the key and entry the way main now does. Could you rebase onto current main? Once that's pushed, we'll gate it and land it. Thank you for your patience on this one.

[kieranscotney]

Thanks @ashhart, and thanks for tracking down the calibration spread.

Rebased onto current main (after #1738); the branch is one commit again, f3b551c74.

• laneEntry(p, sk, skip, cut, long) now builds the key exactly as main's quantOne does, long included, and
  creates the entry with main's .rows = if (long) lane.max_reg_rows else lane.reg_rows.
• quantOne computes long with longKey(k.long_prompt, k.prompt, rows) as before, calls laneEntry, and takes
  its step from entry.value_ptr.max_rows.
• prepareDecode(p) passes long = false: it returns early on prompt kernels, and longKey is false for decoded
  rows, so it builds the same key a decode-row quant of the head builds.

zig build test -Dcpu=apple_m1 exits 0 on the rebased branch (Apple M5 Max). The failed command lines in the log
come from the failure-path tests, as on main.

[ashhart]

@kieranscotney, to save you the rebase: we put your commit on current main as lead/pr1734-rebased. Your commit is unchanged in authorship, and laneEntry now builds the key and entry the way #1738's quant path does, including the long field and the max_rows step. A second commit, by us with you as co-author, adds host tests: for short and long prompts and for decode, the entry prepared at load is the one the first head launch finds, with the same key. Those 5 tests pass, and nothing else changes.

If that looks right to you, either reset your branch to it, or just say so here and we'll land it from there with your authorship as it is. Either way, it's credited to you.

[kieranscotney]

Reset my branch to lead/pr1734-rebased. Thanks for the rebase and tests, good to land from here.

[ashhart]

Thanks @kieranscotney, and sorry for the wait. Our M3 Ultra run (Qwen3.8-27B MLX 4-bit with the DFlash2 drafter, main against this branch, warm-up then base, new, new, base) is green: prompt tokens identical in all 20 rows, decode level at every prompt size (2K 80.6 -> 80.4, 64K 19.5 -> 19.5 tok/s, runs overlapping), served chat at 1 to 8 streams within 1%, and sampled code at 8 streams up 5.2% (163.4 -> 172.0 tok/s, both new runs above both base runs). An independent review found the first draft round reuses the pipeline built at load, so draft and accepted tokens can't change, and a failed compile now refuses the drafter by name at start-up. It lands in the next cycle.

[ashhart]

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