lanes: prompt steps for the 27B and behind --drafter, and a share of wall time for prompt work by jschmied · Pull Request #1188 · ashhart/TensorFold

jschmied · 2026-10-22T05:47:44Z

What this changes

A long cold prompt on the native CUDA engine stops every other stream. Two causes, three commits.

3aa8718d (lanes.drafted). The drafter's wrapper offered the lane host no prefill_step. Behind --drafter, every family's prompt therefore ran whole at admission, even where the target steps its prompts. Over a target with prefill_step, reported fill and features, the wrapper now steps too: it runs the target's step, then gives the drafter that step's rows (all but the prompt's last, as the streamed prefill does).
bd64eb6d (qwen CUDA). The Qwen3.5 family ran a whole prompt pass inside prefill. prefill_step now runs the same pass one cut a call: the same nextCut boundaries, forwards and marks, in the same order. The pass reports its fill, and a cancel lands between cuts.
34560aa8 (lane host). The host alternated one prompt step with one decode round. TENSORFOLD_PREFILL_SHARE=p in (0, 1] gives prompt steps a share p of wall time while streams decode:
- a step that took t owes the decoders t (1 - p) / p of rounds;
- a prompt with at most 512 rows left never waits;
- only the order of whole steps changes, never a cut;
- p = 1, the default, is today's alternation.

Stacking. On main @ f8d254c5e (it contains chaog992's lone-driver halt fix, 010850d74). Branch prefill-interleave.

Receipt

Environment. One GB10 (sm_121), CUDA 13.0, Zig 0.17.0.
- 27B: Qwen3.8-27B MLX 4-bit.
- Flash Next: nvidia/Qwen3.8-Flash-Next-FP4-Local @ b9ce18a6, served from a symlink folder with a rewritten index (see 1.1.2 refuses nvidia/Qwen3.8-Flash-Next-FP4-Local: its index names shards under tensors/ #1187).
- --parallel 4, no prompt cache, --no-drafts unless noted.
The test. Stream A is decoding a long reply. 5 s in, B sends a cold prompt; 0.5 s later, C sends a 34-token chat. Measured: B's and C's first token, A's longest pause, and A's rate while B fills.

Baseline (main).

	B prompt	B first token	C first token	A longest pause	A during B (before)
27B	35,139	157.2 s	156.9 s	157.4 s	0 tok/s (13.9)
27B	97,709	1,352 s	1,351.7 s	1,352 s	0 (14.0)
Flash Next (1,536-row steps)	35,140	27.3 s (25.6 alone)	1.06 s	1.30 s	1.0 (22.8)
Flash Next	97,715	77.2 s	1.56 s	1.82 s	0.9 (23.6)

The 27B stops everyone for the whole prompt. Flash Next already steps, but its decoders get one token per 1.1 s step: 4 % of their rate for the whole prefill.

This PR.

	p	B first token	C first token	A longest pause	A during B
27B, 35k	1	179.0 s (+14 %)	0.34 s	1.39 s	1.5 tok/s
27B, 35k	0.5	323.9 s (2.06x)	0.19 s	1.39 s	6.9 (49 % of 14.0)
Flash Next, 35k	0.5	51.6 s (2.0x alone)	0.96 s	1.26 s	10.9 (50 % of 21.8)
Flash Next, 35k	0.75	35.0 s (1.37x)	0.95 s	1.25 s	5.7 (26 %)

Predictions written before each run:

Every result after the change landed in its predicted range.
The 27B baseline missed high. I predicted 15-45 s at 35k and 450-900 s at 98k; main's 128-row prompt chunks run at about 224 tokens a second at 35k, and slower further in.
Exactness. Token hashes against main's plain run (1 stream, no drafts, no cache):
- 27B with stepped prompts, p = 1 and 0.5, 4 streams with the prompt cache, run twice (resumed at 4,736 and 2,816 tokens): 48 of 48.
- 27B with the DFlash2 drafter through the stepped wrapper, p = 0.5: 24 of 24.
- Flash Next at p = 0.5 with MTP drafts, 4 streams and the prompt cache, twice (resumed at 4,854 and 2,864): 24 of 24.
Tests.
- A new host test (drafted_step_test.zig) runs a stepping target's prompts through the wrapper three steps at a time; the replies equal the whole prefill's.
- The fake target now holds a cold stepped pass's rows as each step makes them.
- Full zig build test -Dnvcc=…: 276/276 steps, 3454/3469 passed, 15 skipped (missing libjpeg/libwebp and hardware), exit 0.
- lean_check: 0 new. fmt clean.
Contract. No new kernel.

Not run, and why

Metal: no Mac. Metal's Flash Next steps since 1.1.2 (Flash Next on Metal (1.1.0): a decoding stream stops for the whole of another stream's prefill #1032), and the share applies to any stepping family behind the lane host.
A decoder's pause below one step: it is still one whole step (1.25 s for Flash Next's 1,536 rows, up to 1.4 s for one 27B cut). Going lower needs smaller steps on each family's own cut grid.

Gotchas and overlap

qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080 (BobClawblaw, 4096-row chunks for the 27B) and qwen CUDA: TF_DRAFT_NODES above 15 serves; a stop or cancel lands between prompt chunks #1078 (cancels between chunks) change the same prompt loop in families/qwen3_5/cuda_native.zig. bd64eb6d is a thin split of that loop: whichever lands first, the other rebases. With 4096-row chunks, a 27B step becomes the whole wait a decoder sees (about 4 s).
feat(drama3): batch CUDA prompt prefill in bounded rounds #1120 (drama3: bounded prompt rounds with fixed cuts) follows the same pattern for another family.
Drafted.take now cuts follow to the rows it absorbs, as the whole prefill's absorb does.

Questions for @ashhart

The default p. It is 1 today: the old alternation, so no default behaviour changes. p = 0.5 halves both sides; 0.75 keeps decoders at about a quarter while the long prompt is 1.37x slower. Do you want a default other than 1, chosen by family or by measured cost?
A flag. Should this be a CLI flag (--prefill-share) rather than an environment variable?

jschmied · 2026-10-22T09:39:40Z

Metal check, filling the "Metal: no Mac" line in the description. zig build test --summary all on an M3 Pro (36 GB), macOS 26.7.1, zig 0.17.0, each commit in its own clean worktree:

commit	steps	tests passed / skipped / failed
`main` @ `bd264368a`	395/395	5712 / 21 / 0
this PR @ `34560aa8f`	395/395	5701 / 21 / 0

Both exit 0. The 11 fewer tests come from this branch's base, which is 10 commits behind that main. The same 21 tests are skipped on both. This is the build and host-test suite only: no model was served on the Mac, so the share itself is still measured on CUDA alone.

ashhart · 2026-10-24T10:16:13Z

Thanks @jschmied. We rebased #1188 onto current main. The rebase conflicted with #1078's cancel check, and we kept your side, since chunk() keeps the check. We then ran it against main on a GB10 with Qwen3.8-27B MLX 4-bit, in alternating base/new/new/base runs.

Tokens are identical to main.
The served sweep: 1 stream 13.9 → 13.9 tok/s (+0.0%), 2 streams 26.1 → 26.1 (−0.2%), 4 streams 48.7 → 48.1 (−1.2%), 8 streams 85.4 → 82.0 (−4.0%).

An earlier single run showed the same 8-stream loss (−3.6%), so it reproduces. The cost grows with stream count, which points at the lane host's share of wall time for prompt work (34560aa). Each decode round pays for it even once the prompts are done.

Could the prompt share apply only while a prompt is actually waiting, so steady-state decode rounds run as they do on main? The first two commits (prompt steps behind --drafter, and Qwen's stepped prefill) look independent of this, and we'd be glad to land them on their own if they measure flat. We'll land the whole PR once 8-stream decode is within noise of main and a short request behind a cold prompt still gets its faster first token.

- Synthema, Ash's AI Assistant.

lanes.drafted.Drafted offered the lane host no prefill_step. Behind --drafter, every family's prompt therefore ran whole at admission, and every other stream waited for all of it, even where the target steps its prompts. Over a target with prefill_step, reported fill and features, the wrapper now steps too: each step runs the target's, then gives the drafter that step's rows (all but the prompt's last, as the streamed prefill does). The drafter opens with the stream's first step. Drafted.take now cuts `follow` to the rows it absorbs, as the whole prefill's absorb does. The fake target holds a cold stepped pass's rows as each step makes them. A new host test runs a stepping target's prompts through the wrapper three steps at a time; the replies equal the whole prefill's.

The Qwen3.5 CUDA family ran a whole prompt pass inside prefill. The lane host admitted it whole, so every decoding stream and every queued short chat waited for all of it: 157 s behind a 35k cold prompt on a GB10, and 22.5 min behind a 97.7k one. prefill_step now runs the same pass one cut a call: the same nextCut boundaries, forwards and marks, in the same order. prefill runs the same steps to the end. The pass reports its fill, and a cancel lands between cuts. Bits are unchanged: 48 of 48 token hashes equal main's plain replies at four streams with the prompt cache (resumed at 4,736 and 2,816 tokens), and 24 of 24 with the DFlash2 drafter. Behind the 35k prompt, a short chat's first token comes in 0.34 s and a decoding stream's longest pause is 1.39 s; the long prompt's own first token takes 14 % longer (179 s).

The lane host alternated one prompt step with one decode round. Decoders waited at most one step, but got one token per step: on Flash Next (1,536-row steps), a stream decoding at 22 tokens a second fell to 1 for a 35k prompt's whole prefill. TENSORFOLD_PREFILL_SHARE=p in (0, 1] gives prompt steps a share p of wall time while streams decode: a step that took t owes the decoders t (1 - p) / p of rounds before the next step. A prompt with at most 512 rows left never waits, so a short chat still goes before a long prompt's next step. Only the order of whole steps changes, never a cut, so resumed and concurrent replies keep their bits. p = 1 (the default) is today's alternation. GB10, a 35k cold prompt behind one decoding stream: - Flash Next at p = 0.5: the stream keeps 10.9 of its 21.8 tokens a second (1.0 at p = 1); the long prompt's first token takes 51.6 s (25.6 s alone); a short chat's takes 0.96 s. - Qwen3.8-27B at p = 0.5: 6.9 of 14.0 tokens a second; 324 s for the long prompt; 0.19 s for the chat. Flash Next's replies equal main's plain ones at p = 0.5 with MTP drafts, four streams and the prompt cache: 24 of 24.

…en no prompt waits On 8 streams the stepped prompt pass cost decode speed against main: main's qwen CUDA backend prefilled every admitted prompt whole before the next round, while the lane host filled one step per round, so a burst of short prompts started its streams a round apart. Every round also read the clock twice for the share. A prompt whose whole pass is at most 512 rows now goes in back to back with the other short ones before the next round, as a whole pass at admission did, in the host's fill order (foreground, aging, least remaining). Long prompts still take one step a round, under the share. A round reads the clock only while it pays a step's debt down; with no prompt waiting the loop is main's. GB10, Qwen3.8-27B MLX 4-bit, --no-drafts, bench_concurrent (256 tokens, 4 cells, 3 reps), fresh servers in the order main/new/new/main/old/main: - aggregate tok/s, main (3 runs) -> this commit (2 runs): 1 stream 15.07 -> 15.00, 2: 28.27 -> 28.26, 4: 52.97 -> 52.77, 8: 93.72 -> 93.51 (-0.2 %; main's runs span 93.25 to 94.35). The previous head: 90.22 at 8 (-3.7 %). - token hashes equal main's in every cell. - a 69-token chat sent 3 s into a cold 12k-token prompt: first token in 0.26 to 0.30 s, against 19.6 to 19.9 s on main; the long prompt's own first token 23.5 to 23.7 s, against 22.5 to 22.7 s.

jschmied · 2026-10-24T13:36:52Z

Thanks for the rebase and the runs. I rebased onto current main the same way (our side; chunk() checks isCancelled before each cut, and a cancelled pass's cache goes back through release). I also added a commit for the 8-stream loss. Most of the loss wasn't the share's bookkeeping. Main's qwen CUDA backend prefills every admitted prompt whole before the next round, but stepped prompts got one step per round, so 8 short prompts arriving together started decoding a round apart. Now any prompt of at most 512 rows goes in whole, back to back, before the next round. Long prompts still take one step per round under the share. A round only reads the clock while it pays down a step's debt, so with no prompt waiting the loop is the same as main's. That fix is in the lane host, so I wouldn't land the first two commits without it.

GB10, Qwen3.8-27B MLX 4-bit, --no-drafts, bench_concurrent.py --levels 1,2,4,8 --tokens 256 (4 cells, 3 reps), a fresh server per run, order main/new/new/main/old/main, aggregate tok/s:

streams	main (3 runs)	new (2 runs)		previous head (1 run)
1	15.07	15.00	-0.5%	14.70
2	28.27	28.26	-0.0%	27.77
4	52.97	52.77	-0.4%	51.60
8	93.72	93.51	-0.2%	90.22 (-3.7%)

At 8 streams, main's own runs span 93.25 to 94.35.
Token hashes match main's in every cell.
A 69-token chat sent 3 s into a cold 12k-token prompt gets its first token in 0.26 to 0.30 s (main: 19.6 to 19.9 s).
The long prompt's own first token: 23.5 to 23.7 s (main: 22.5 to 22.7 s).

… a restored fake pass holds its rows The prompt share counted paused streams as decoders, so with every stream paused a long prompt's step created debt that nothing paid down; it now counts runnable streams. A share below 0.01 falls back to 1 and a step's debt saturates instead of overflowing. The fake backend holds a restored stream's kept rows from its first step, so the drafter wrapper's per-step feature reads find them.

jschmied · 2026-10-24T14:14:37Z

One more commit (8de3712) from a second review pass, no change to the scheduling policy or the numbers above:

The share counts runnable streams only. With every decoder paused, a long prompt's step no longer builds debt that nothing pays down.
TENSORFOLD_PREFILL_SHARE below 0.01 (or NaN/inf) falls back to 1, and the debt add saturates, so a tiny share can't overflow the i64 conversion.
The fake backend holds a restored prompt's kept rows from its first step, so the drafter wrapper's per-step feature reads find them. A new test (partial and full restore, cancel after step 2) failed with RowsNotHeld before.

ashhart · 2026-10-25T04:54:46Z

Landed on main in 2ff0191..464824f, with your five records under your authorship. Thank you, @jschmied, and thanks for chasing down the 8-stream loss.

We replayed it onto current main, where #1228's interrupted-prefill rules had landed in the meantime. One maintainer child reconciles the two: a backend failure during a stepped prompt still ends failed, a cancel keeps its last complete cut, and the drafter forwards its cancel hook.

On a GB10 with Qwen3.8-27B MLX 4-bit, alternating base/new/new/base:

Tokens: identical (0 of 20 rows differ).
Prefill: 2k to 64k within 0.3%.
Served sweep: flat at 1, 2, 4 and 8 streams (8 streams 87.4 → 87.4). The earlier −4.0% is gone.

- Synthema, Ash's AI Assistant.

jschmied mentioned this pull request Oct 22, 2026

Flash Next CUDA: NVFP4 experts on prompt tiles with the decode form's sums, bit for bit #1143

Closed

9 tasks done

Sapphire-Rapids mentioned this pull request Oct 24, 2026

Preserve completed CUDA prefill across cancellation and retry #1228

Closed

7 of 9 tasks

jschmied added 4 commits October 24, 2026 12:33

jschmied force-pushed the prefill-interleave branch from 34560aa to 8154927 Compare October 24, 2026 13:36

ashhart closed this Oct 25, 2026

This was referenced Oct 25, 2026

gemma4 CUDA: a stepped drafted pass keeps its mark after the drafter absorbed the chunk #1286

Closed

lane host: three fairness fixes to #1188's prompt steps #1287

Closed

Chris9121 mentioned this pull request Oct 26, 2026

GLM-5.3 Metal path: DSA indexer scan 100-300 ms per 1K context (85x the CUDA rate of #259) #1054

Closed

Navigation Menu

Repository navigation

lanes: prompt steps for the 27B and behind --drafter, and a share of wall time for prompt work - #1188

lanes: prompt steps for the 27B and behind --drafter, and a share of wall time for prompt work#1188
jschmied wants to merge 5 commits into
ashhart:mainfrom
jschmied:prefill-interleave

jschmied commented Oct 22, 2026

Uh oh!

jschmied commented Oct 22, 2026

Uh oh!

ashhart commented Oct 24, 2026

Uh oh!

jschmied commented Oct 24, 2026

Uh oh!

jschmied commented Oct 24, 2026

Uh oh!

ashhart commented Oct 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jschmied commented Oct 22, 2026

What this changes

Receipt

Not run, and why

Gotchas and overlap

Questions for @ashhart

Uh oh!

jschmied commented Oct 22, 2026

Uh oh!

ashhart commented Oct 24, 2026

Uh oh!

jschmied commented Oct 24, 2026

Uh oh!

jschmied commented Oct 24, 2026

Uh oh!

ashhart commented Oct 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants