[gilby]

What this changes

Fixes #1570. CUDA GLM takes prompt chunks of 4,096 rows instead of 2,048 (cuda_state.prefill_rows). DFlash2's MAX_CTX, which sizes the context rows one ingest takes, now follows prefill_rows instead of a separate 2048: changing the chunk alone faults at startup with a drafter loaded at TP=3 (CUDA_ERROR_ILLEGAL_ADDRESS in cuStreamSynchronize), because the ingest of a 4,096-row chunk overruns the 2,048-row tap and context buffers. The output is unchanged.

Receipt

• Environment: three DGX Spark (GB10, 128 GB), driver 580, CUDA 13.0, ConnectX-7 RoCE through a switch, NCCL over IB; nvidia/GLM-5.3-Flash-FP4-Local @ 76348ddb, --context 500000 --parallel 4 --prompt-cache-gib 0 --kv-cache fp8, TF_GLM_INDEX_CUDA=1; zig build native -Dsm=121. Each table's rows come from one session.

• Prompt speed, cold first token, tools/bench_cache.py --require-cold --repeats 1 --streams 1 --tokens 16 (a fresh random prefix each run), TP=3, no drafter, main a881b6f:

  prompt tokens	main	4,096 rows
  32,768	22.01 s	19.11 s (−13.2%)
  41,943	26.74 s	24.89 s (−6.9%)
  166,700	104.68 s	97.08 s (−7.3%)

• TP=2 (two ranks), no drafter, main a881b6f: 32,768 tokens 24.92 → 22.84 s, 131,072 tokens 98.98 → 92.54 s. Startup weights and workspace a rank 93.24 → 95.18 GiB; four 500K windows still fit (11.1 GiB left after them); lowest MemAvailable on a rank during the 131K prompt 12.9 GiB.

• With the DFlash2 drafter at TP=3, main 7c06889 against this branch: startup 67.19 → 68.97 GiB a rank, four 500K windows fit (36.4 GiB left). 41,943 tokens 26.83 → 25.92 s and 166,700 tokens 103.71 → 96.69 s to the first token, the same 64-token completions (same hashes), and draft acceptance level (34/59 and 33/58, 27/76 and 28/74).

• Exactness, tools/bench_concurrent.py --levels 1,4 --tokens 256 --reps 3 --temperatures 1.0,0 --alone --serial: 0 unequal and 0 failed in every cell, without the drafter and with it, both at TP=3; with the drafter, all 540 completion hashes equal main's with the drafter. Without the drafter, 64-token completions after 41,943- and 166,700-token prompts hash the same at 2,048 and 4,096 rows.

• Decode speed: level (with the drafter at TP=3, code T=1 one stream 63.5 vs 63.6 tok/s).

• Tests: zig build test-cuda-host on Linux aarch64 (GB10) at 7c06889 with this change: 270/270 steps, 3,843 of 3,852 tests passed, 9 skipped. zig fmt --check clean on both files. python3 tools/zig/lean_check.py: 425 problems, the same as main.

• Trade-off: at TP=2 with the drafter, the larger chunk and drafter buffers leave 9.90 GiB for streams where --parallel 4 at a 500K window needs 10.44 GiB, so that setup serves --parallel 3 (or a shorter window) where main fits four. TP=3, and TP=2 without the drafter, still fit four.

• Not run: the exactness bench at TP=2 (the row split is the same code at any rank count). Metal is not affected (CUDA family path).

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false and concurrent equals solo (the bench above). A prompt split at other rows gives the same bits, so a resumed prompt reads the same caches as a fresh one.
• No precision traded for speed. The bits do not change.
• No GPU kernel added or changed (two constants).
• Prompt processing is no slower than the last release.
• Platforms: CUDA on GB10 at TP=3 and TP=2. Metal not affected.
• New tests fail before the change, pass after it: none added; the drafter fault needs a real drafter checkpoint on CUDA ranks.
• 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]

Thank you for this, and for the thorough receipts. The cold-prompt gain is clear, and it's good to see the completions match main with and without the drafter.

Before we can land it, there are a few things to sort out.

1. The TP=2 drafter setup loses a stream. You noted that at TP=2 with the DFlash2 drafter, --parallel 4 at a 500K window no longer fits, where it does on main. We don't take a capacity regression on any setup, even in exchange for a speedup elsewhere. Most of the extra memory is the prompt workspace and the index ring (index_ring = 2 * prefill_rows), not the drafter's buffers, so shrinking the drafter side alone won't recover it. Two ways we'd be happy with:

   • pick the chunk at startup: 4,096 rows when the requested streams and window fit, and 2,048 otherwise. The native startup already lets a family report its prompt rows at load, so the snapshot grid and buffers could follow that value.
   • or keep 2,048 as the default and make 4,096 an opt-in setting.

2. Warm turns need a receipt. The prompt cache keeps and resumes states only at multiples of the chunk (snap_grid follows prompt_rows, and resume needs at % prefill_rows == 0). With 4,096-row chunks, a warm turn can re-fill up to 4,095 rows instead of up to 2,047. All the receipts so far ran with --prompt-cache-gib 0 and --require-cold. Could you post first-token times with the prompt cache on, against main? For example, a multi-turn chat that adds about 500, 1,500 and 3,000 tokens a turn after 10K and 40K histories.

3. Decoding streams beside a long prompt need a receipt too. Each long-prompt step now runs about twice as long, so streams that are already decoding wait about twice as long for their next round while a long prompt fills. Could you post p50 and p99 inter-token gaps for three decoding streams while a 32K prompt arrives, against main, at the default prefill share?

4. Small fixes:

   • cuda_engine.zig still documents prefill as taking "2,048-row chunks".
   • The prompt profiler's last bucket is still labelled "prompt chunk to 2048 rows".
   • The commit is authored as gilby@users.noreply.github.com. Could you re-commit with the ID form 880013+gilby@users.noreply.github.com, which your earlier commits here use?

Once those are in, we'll run the device cells and land it.

[Bizuayeu]

2. Warm turns need a receipt. […] Could you post first-token times with the prompt cache on, against main? For example, a multi-turn chat that adds about 500, 1,500 and 3,000 tokens a turn after 10K and 40K histories.

3. Decoding streams beside a long prompt need a receipt too. […] Could you post p50 and p99 inter-token gaps for three decoding streams while a 32K prompt arrives, against main, at the default prefill share?

(#1576 (comment))

The receipts we'd want:

• the longest wait between a decoding stream's tokens while a long prompt fills, at 2,048 and 4,096;
• the TP=2 fit and its first-token times, as you note;
• concurrent equal to solo and resumed equal to fresh at both sizes.

(#1570 (comment))

@gilby I had the same chunk size under test here, so these are yours to use if they save you a run. At 4,096 rows a warm turn takes 10 to 52% longer, and a decoding stream's longest wait goes from 1.73 to 3.07 s at TP=2 and from 1.45 to 2.52 s at TP=3. The bits match at both sizes. So these receipts support the opt-in setting, not the default.

Setup. main a0d24896f against a local build of it with prefill_rows = 4096. That build is not this branch: MAX_CTX stays 2,048 and only the drafter's tap rows follow prefill_rows. No DFlash2 drafter is loaded, so item 1's drafter setup is not measured here. GLM-5.3-Flash XSFP4 at 88736f916f89, --parallel 4, KV pages on, the MTP drafts on, the default prefill share, a 524K window. TP=2 on two GB10, at the default prompt cache, which the engine sizes at 4.8 GiB at 2,048 rows and 2.8 GiB at 4,096. TP=3 on three, at --prompt-cache-gib 2. Launched 2,048, 4,096, 4,096, 2,048 at each TP. No launch waited on the heat guard.

Warm turns. A chat with a fresh nonce, its history sent cold, then four turns that each add the tokens shown, 32 tokens a reply. Each cell is the mean first token and, in brackets, the mean rows re-filled, over 16 turns at TP=2 and 8 at TP=3:

History + a turn's addition	TP=2, 2,048	TP=2, 4,096		TP=3, 2,048	TP=3, 4,096
10K + 500	1.38 s (1,650)	1.92 s (2,672)	+39%	1.14 s (1,647)	1.59 s (2,671)	+40%
10K + 1,500	2.02 s (2,571)	2.23 s (3,082)	+11%	1.68 s (2,569)	1.85 s (3,085)	+10%
10K + 3,000	2.83 s (3,697)	3.36 s (4,722)	+18%	2.36 s (3,701)	2.77 s (4,726)	+17%
40K + 500	1.23 s (1,374)	1.88 s (2,399)	+52%	1.02 s (1,374)	1.54 s (2,397)	+51%
40K + 1,500	2.01 s (2,478)	2.93 s (4,014)	+46%	1.66 s (2,478)	2.41 s (4,014)	+45%
40K + 3,000	3.02 s (3,888)	3.55 s (4,909)	+18%	2.51 s (3,886)	2.92 s (4,910)	+17%

Every turn resumed at the deepest point of its own grid (cached_tokens checked turn by turn), so TP=2's two cache sizes didn't change what was kept. The first send, cold with the cache on, gains 7.00 → 6.85 s at 10K and 27.88 → 26.93 s at 40K at TP=2, and 5.82 → 5.63 s and 23.25 → 22.15 s at TP=3 (−2 to −5%).

Correction: I first set that against the −6 to −8% from 8K to 128K at --prompt-cache-gib 0 and called them the same filler. They aren't: those cold prompts are Python source from reuse_served.py ids, and these turns are a repeated filler line. So the two gains don't compare, and the cache may have nothing to do with the gap. I'm measuring both prompts at --prompt-cache-gib 0 on one build and will post that here.

Three decodes beside a 32K prompt. Three drafts-off decodes of up to 1,200 tokens, and 4 s later one cold 32K prompt, over 300 and 162 gaps at TP=2 and 192 and 96 at TP=3:

	TP=2, 2,048	TP=2, 4,096	TP=3, 2,048	TP=3, 4,096
Gap before it arrives, p50 / p99 (ms)	45 / 90	45 / 90	37 / 73	37 / 73
Gap while it fills, p50 (ms)	1,460	2,764	1,217	2,275
p99 (ms)	1,726	3,066	1,450	2,517
Longest (ms)	1,731	3,066	1,451	2,517
Tokens a stream during the fill	16.7	9.0	16.0	8.0
The prompt's fill (s)	23.47	22.27	19.54	18.30

TP=2 fit and first tokens, at --prompt-cache-gib 0, no drafter: the model takes 90.59 → 92.53 GiB, the KV pool 1,796 → 1,468 pages, of which one 524K window takes 1,024, and the lowest MemAvailable stays at 11 to 12 GiB. Cold first token: 8K 6.02 → 5.66 s, 32K 23.7 → 21.9 s, 128K 97.1–98.5 → 89.9 s.

TP=3 at a 1M window holds one whole window fewer, with no drafter loaded. The pool goes 6,215 → 5,948 pages and a 1M window takes 2,048, so the streams a 1M window each would serve (kv_pool.windows) go from three to two. The planner holds those windows before it gives the cache anything, so at that window --prompt-cache-gib 2 gets 0.4 GiB on main and all 2 GiB at 4,096 rows, and main's 40K turns went cold. That's why the TP=3 cells above use 524K. --parallel 4 still fits here, so a startup choice that looks only at whether the streams fit would keep 4,096 rows and this loss; it would need to look at those windows too. The opt-in covers it.

Exactness at both sizes. Prompts of about 8K, 32K and 64K tokens, greedy and sampled (seed 1234, 64 tokens, ignore_eos), each sent fresh, again (resumed at 6,144, 30,720 and 63,488 tokens at 2,048 rows, and 4,096, 28,672 and 61,440 at 4,096), then twice at once: at TP=2 and at TP=3, all eight replies to a prompt, four at each size, have one token_sha256. At TP=2, which your receipt lists as not run, tools/bench_concurrent.py --levels 1,4 --tokens 256 --reps 3 --temperatures 1.0,0 --alone --serial gives 60 of 60 equal alone and 20 of 20 equal "draft": false at each size, 0 failed, and the hashes match between the sizes. On top of #1732 at TP=3, 70 replies, 18 of them resumed, hash the same as on #1732 alone, and a 64K repeat goes 1.74 → 2.82 s as it resumes at 61,440 instead of 63,488.

Why it gains. Per row, tf_nvfp4_prompt_swiglu and _down take 28 to 40% less time at 4,096 rows at both rank widths, while the expert items per row fall 8 to 24%, and ncu has their L2 hit rate going from 45–56% to 68–71%. So the gain comes from weights a longer launch reads again, which is why a pair, running each chunk's kernels on their own, doesn't reach it and the two stack as in #1570.

For the device cells. Three things that fail or go stale at 4,096 rows, which you may want before the run:

• zig build test-glm-prompts-host, which CI doesn't run, goes 7/7 → 6/7 on this diff: glm_prompt_host.zig expects DrafterContextTooWide from a literal 2044, and MAX_CTX is now 4,096.
• tf-cuda-test glm-lanes-cover fails four lines of its own assumptions: "a cancel at ask 3" and "the cancel was heard at ask 3" take the prompt as three chunks, and "the pass kept 1 states at its marks" and "the pass resumed at 2048" take 2,048 as a grid point. With those read from prompt_rows, it fails only "copy drafts, sampled", as main's runner does at TP=3.
• glm_paged_tests' lengths have nothing on the new chunk's edges (4,095, 4,096 and 4,097 pass here), and cuda_lanes.zig's "2048 grid" comment and two in the tests join the two in item 4.

The test changes are about twenty lines; I can post them as a patch if that helps.

On the shape. #1554's C would fit GLM too: 4,096 rows a pass only while no other stream decodes, as #1423's pair chooses, and the grid left at 2,048, as #720 does for Nemotron with pass_rows. A warm turn would then resume where main does and a decoding stream would see main's passes. It would not help item 1 or the 1M window, as the buffers are still sized for the wide pass, so the startup choice would still apply. I haven't built it.

[ashhart]

Thanks @Bizuayeu, these are exactly the receipts we asked for, and they answer it clearly.

At 4,096 rows, against 2,048, on GLM-5.3-Flash:

• Warm turns get slower first tokens at both widths: +10 to +18% when a turn adds 1,500 or 3,000 tokens, and +39 to +52% when it adds 500. That's the common chat case.
• Decoding streams beside a long prompt wait almost twice as long. At TP=2 the gap while a 32K prompt fills goes from 1.46 s to 2.76 s at p50, and from 1.73 s to 3.07 s at p99. Each stream gets 9 tokens out during the fill instead of 16.7.
• The cold 32K fill gains 5 to 6% (23.47 s to 22.27 s at TP=2; 19.54 s to 18.30 s at TP=3).

So 2,048 stays the default: a default can't get slower for warm turns or for streams already decoding. We'd still take 4,096 as an opt-in setting for servers that mostly see cold long prompts, as on #1554 and #1570, with these tables in the recipe so people can choose. If you'd like to reshape the PR that way, it's welcome. The DFlash2 context half can go separately if it stands on its own.

[Bizuayeu]

So 2,048 stays the default: a default can't get slower for warm turns or for streams already decoding. We'd still take 4,096 as an opt-in setting for servers that mostly see cold long prompts, as on #1554 and #1570, with these tables in the recipe so people can choose. If you'd like to reshape the PR that way, it's welcome.

(#1576 (comment))

Thanks, that settles the default. The PR and the claim on #1570 are @gilby's, so the reshape is theirs unless they'd rather hand it over.

One question on the shape, since it changes what the recipe's tables say: should the opt-in be a fixed 4,096, or #1554's C? I have C as a local prototype on main, TF_GLM_PROMPT_WIDE=1: 4,096 rows a pass only while no other stream decodes, with the grid, the index rings and prompt_rows left at 2,048, and not with a drafter, so it needs no DFlash2 change. Off, it holds nothing extra. On, it holds the same prompt workspace as 4,096 rows, and warm turns and decoding streams should stay at main's numbers. That last part is my reading of the code: the host tests pass, but it hasn't run on a GPU yet. If C is the shape you'd want, I'll measure it and post the same tables here.

[Bizuayeu]

Measured it anyway, since the hosts were free: #1554's C on GLM at TP=3, in the shape of #1733's table. One build on main 94b90c9db, launched unset, =1, =1, unset, in the setup above (three GB10, 524K window, --prompt-cache-gib 2, KV pages on, the MTP drafts on).

	unset	TF_GLM_PROMPT_WIDE=1	Fixed 4,096 (above)
First token, one stream, 7,822 / 32,453 / 136,169 tokens	4.78 / 18.92 / 80.00 s	4.65 / 18.03 / 76.22 s (−2.6 / −4.7 / −4.7%)
Warm turns, the six rows above		−1.8 to +0.2%	+10 to +51%
Gap between a decoding stream's tokens while a 32K prompt fills, p50 / longest	1,218 / 1,456 ms	1,217 / 1,454 ms	2,275 / 2,517 ms
First token of a short request sent 6 to 8 s into a 32K prompt, mean (range)	1.25 s (0.70 to 1.69)	1.68 s (1.06 to 2.60)
That prompt's own first token	19.34 s	18.86 s
Model and scratch at startup, a rank	63.34 GiB	64.90 GiB (+1.56)	64.93 GiB
Lowest MemAvailable on rank 0	22 to 23 GiB	21 GiB

• Arrival cost: this is what C adds, as on flashnext cuda: TF_FLASHNEXT_CHUNK_SOLO, a pass's rows chosen as it starts (#1554 option C) #1733. A pass lasts about 2.3 s at 4,096 rows against 1.2 s at 2,048, going by the gaps above, and that's the longest a request arriving during one waits before its own prompt runs. So it stays opt-in. The rest of the table is equal or faster, and unset is main's path.
• Exactness: glm-wide-check, new in the branch, steps eight prompts wide and a chunk at a time on one engine at TP=3: the first token, the next window's logits and tokens, the rings' bases, every other state byte and the KDA records, on the copy each reads next as a wide pass flips them once, are equal. glm-lanes-concurrent passes, and glm-lanes-cover fails only "copy drafts, sampled", as main's runner does at TP=3. For prompts of about 8K, 32K and 64K tokens, greedy and sampled, each sent fresh, again and twice at once, all eight replies to a prompt, four at each setting, have one token_sha256. All 96 warm turns resumed on the 2,048 grid.
• Tests: test-cuda-host 270/270 steps, zig build test 445/445, compile-cuda-tests 10/10, tools/zig/lean_check.py with nothing new.
• Not run: TP=2, and a DFlash2 drafter, as a step stays at 2,048 rows with one loaded. It holds the wide scratch whenever it's on, so the 1M window's loss above would stand with it on, though not unset.

It's two commits, local so far: the lanes' host checks moved to their own file, as #1423's first commit does, and the step, which I'd rename after #1733's variable. It overlaps #1423 in Leader.chunk, the lanes and the fingerprint. A wide step and a paired step are two ways of taking two chunks, and this one doesn't yet pair wide chunks, the stacked form in #1570. The PR and the claim stay @gilby's: I can push the branch for you to take, or open it if you'd rather hand this over.

[Bizuayeu]

TP=2 for the table above, and the re-measure from the receipts' correction: I had called the −6 to −8% at --prompt-cache-gib 0 the same filler, and those prompts are Python source. Same build as above on two GB10, launched unset, =1, =1, unset, at a 524K window, KV pages on, the MTP drafts on. The cache is the default here, where TP=3 had 2 GiB on both sides.

	unset	TF_GLM_PROMPT_WIDE=1	Fixed 4,096 (above)
Warm turns, the six rows above	1.24 to 3.04 s	−1.7 to +0.2%	+11 to +52%
Gap between a decoding stream's tokens while a 32K prompt fills, p50 / longest	1,464 / 1,738 ms	1,463 / 1,725 ms	2,764 / 3,066 ms
First token of a short request sent 6 to 8 s into a 32K prompt, mean (range)	1.25 s (0.83 to 1.71)	2.30 s (1.26 to 3.29)
That prompt's own first token	23.53 s	23.00 s
Decode, bench_openai, each launch's mean of two prompts	40.46, 40.13 tok/s	40.31, 40.36 tok/s
Model and scratch at startup, a rank	90.59 GiB	92.51 GiB (+1.92)	92.53 GiB
KV pool at --prompt-cache-gib 0, of which a 524K window takes 1,024	1,851 pages	1,542 pages	1,796 → 1,468 on a0d24896f
Default prompt cache	5.2 GiB	3.2 GiB	4.8 → 2.8 GiB on a0d24896f
Lowest MemAvailable, rank 0 / rank 1	12 / 12 GiB	12 / 11 to 12 GiB

Cold first token, one stream. Cache 0 is one launch a side:

Prompt, cache	About 8K	About 32K	About 130K
Python source (8,192 / 32,768 / 131,072 tokens), 0, unset → =1	6.01 → 5.66 s (−5.9%)	23.72 → 22.05 s (−7.1%)	97.10 → 90.67 s (−6.6%)
Python source, 0, main → fixed 4,096 (above)	6.02 → 5.66 s (−6.0%)	23.71 → 21.91 s (−7.6%)	97.34 → 89.88 s (−7.7%)
Filler (about 7,820 / 32,450 / 136,170 tokens), 0, unset → =1	5.71 → 5.52 s (−3.3%)	22.70 → 21.91 s (−3.5%)	96.27 → 93.65 s (−2.7%)
Filler, default, unset → =1	5.70 → 5.65 s (−0.9%)	22.76 → 22.04 s (−3.2%)	96.65 → 93.52 s (−3.2%)

• With =1 at TP=2, the gap is the prompt, not the cache. At cache 0, source gains 7.1% at 32K and the filler 3.5%. On the filler, cache 0 gains 3.5% and the default cache 3.2%. With =1 a token costs the same on both prompts at 32K, 673 and 675 µs of first token. Unset it costs more on source, 724 against 700 µs. So the filler understates this, the TP=3 table above included: I haven't run =1 on source there. 8K is the one place the cache shows. Its one mark, at 6,144, ends a wide step early, so I read the rows from 4,096 on as staying narrow.
• Wide against fixed: close on the same prompts. =1's repeats are 5.66, 21.95 and 90.32 s, against 5.66, 21.91 and 89.88 s on an older main. Its first send at 32K and 128K is 0.3 and 1.0 s slower than those repeats, and that's in the means above. Unset doesn't show that and neither did fixed. I haven't found out why.
• Arrival cost: larger than at TP=3, as a pass lasts about 2.8 s at 4,096 rows against 1.5 s at 2,048, going by the gaps above. These are eight requests at four fixed moments, so the range says more than the mean.
• Exactness: at TP=2 glm-wide-check and glm-lanes-concurrent pass, and glm-lanes-cover passes whole, "copy drafts, sampled" included. For prompts of about 8K, 32K and 64K tokens, greedy and sampled, each sent fresh, again and twice at once, all eight replies to a prompt, four at each setting, have one token_sha256. All 96 warm turns resumed on the 2,048 grid. The three source prompts have one hash across unset, =1 and both builds above.
• Not run: =1 on source at TP=3, and a DFlash2 drafter.

[Bizuayeu]

TP=3 on source, which the last comment left open. Same build, launched unset, =1, =1, unset, at --prompt-cache-gib 0, a 524K window, KV pages on and the MTP drafts on.

Prompt	About 8K	About 32K	About 130K
Python source (8,192 / 32,768 / 131,072 tokens), unset → =1	5.00 → 4.65 s (−6.9%)	19.65 → 18.03 s (−8.3%)	80.23 → 73.71 s (−8.1%)
Python source, main → fixed 4,096 (a0d24896f, a 1M window)	5.02 → 4.67 s (−7.0%)	19.69 → 18.06 s (−8.2%)	80.45 → 73.74 s (−8.3%)
Filler (about 7,820 / 32,450 / 136,170 tokens), unset → =1	4.73 → 4.50 s (−4.9%)	18.85 → 17.92 s (−4.9%)	79.94 → 76.38 s (−4.5%)

• On source =1 matches fixed, so the −4.7% in the first table is the filler's, not the setting's limit.
• The slower first send from TP=2 doesn't show here: first against repeats is 18.02 against 18.03 s at 32K and 73.70 against 73.71 s at 128K. TP=2 had one launch a side at cache 0, and launches here differ by 0.8% at 128K (73.42 and 73.98 s), so that may be the launch rather than the setting. I haven't rerun it.
• The three source prompts have one hash across all four sides, and no launch waited on heat.

With #1733 in, this could take the same shape: TF_GLM_CHUNK_SOLO, opt-in, these tables in the recipe. Should that go on #1576 or in a PR of its own?