GLM-5.3-Flash on two CUDA ranks: prompts admitted together (prefill_many)

[BobClawblaw]

GLM-5.3-Flash on two CUDA ranks: prompts admitted together (prefill_many)

Fresh text prompts that arrive together share one prompt pass, instead of filling one at a time between other streams' rounds. That was most of the family's first-token latency with many streams.

The branch is now up to date with main 76aafcf2, which carries #855, #979, #984 and #995. The diff against main is only this change.

What it does

• Lane host (lane_host_admit.zig, lanes/open.zig, backend.zig):
  • A backend that fills by chunks (prefill_step) can also batch prompts. Its many_rows gives the longest prompt it batches.
  • Shorter fresh prompts go together. Longer prompts, prompts with images, and prompts resuming a kept state still go to the fill queue.
  • A failed batched pass logs its error before each stream goes alone.
  • cuda_bound.zig forwards many_rows, and its forwarding test now compares integer fields too.
• The pass (cuda_prompts.zig, cuda_shared.zig):
  • One forward over every prompt's rows, at most 2,048 rows and 16 prompts.
  • With chunked fills on (the steps backend), a batch is one pass. A prompt that would overflow it fills by chunks, so other streams' decode rounds run between passes. This has a host test.
  • Without chunked fills, the lanes split a larger batch into passes with a host-tested helper, passes.
  • Row-wise work runs once. Each prompt's KDA chain and conv window, and its latent and sparse attention, run on its own sequence.
  • The head runs on each prompt's last row.
  • Then, per prompt: its first token, its first-draft source (stream i of the batch keeps hold row i, carried on the follow link), its MTP head's absorb, and its commit.
• DFlash2: the pass's taps feed each prompt's rows into its slot's ring in one ingest (Drafter.prompts), as a prompt chunk does.
• Ranks: a many step on the follow link carries each prompt's slot and length, the first hold row, then the ids.
• Exact alone or batched: a prompt pass's linears take glm_mm's one-slice blocks at any row count. Before, prompts under 128 rows summed K slices in another order.

Review fixes (all in this PR)

• First-draft rows across passes. A batch over 2,048 rows wrote every pass's hold rows from 0, so a later pass overwrote an earlier one's before their first drafts. Rows are now per stream across the batch.
• DFlash2 taps and ring context. The batched pass feeds them. The earlier version of this PR lacked that feed, and only a local combination had it.
• From a review pass after merging main:
  • The batched forward passes TF_GLM_PREFILL_REDUCE=ring like solo prompts.
  • A batch whose follow-link step would exceed the link is refused before any broadcast.
  • Every stream's capacity is checked before the first pass.
  • A comptime check ties the pass size to DFlash2's ingest and window.

The single-stream drop

The earlier "69.8 to 59.5 tok/s" came from the ruler's single-stream cell, which runs one prompt per phase.

• The one-slice prompt matmuls round that prompt differently, so greedy decoding writes a different reply (hash d958bc3258ec to 3113363b6ac2), and that reply drafts worse.
• With drafts off, both builds decode at the same speed: 32.2 tok/s.
• Over all 16 prompts the two builds are equal: code median 69.1 and 69.5.

The table below uses all 16 prompts at one stream.

Checks

Two DGX Sparks, EXL3 checkpoint, DFlash2, --parallel 16, recipe harness (plan glm-qual1014). Aggregate tok/s for prose / code / structured / list:

	main 76aafcf2	this branch
1 stream, all 16 prompts	42.8 / 67.2 / 55.3 / 53.1	42.5 / 67.6 / 56.1 / 55.8
8 streams	90.0 / 120.8 / 114.7 / 115.5	92.4 / 131.0 / 125.4 / 147.8
8 streams, first text	1.23-1.57 s	0.82-1.03 s
16 streams	107.2 / 119.6 / 119.4 / 129.6	117.4 / 134.1 / 147.9 / 158.6
16 streams, first text	2.69-3.38 s	1.34-1.66 s

Correctness on this branch (greedy; replies compared by the hash of their text):

Check	DFlash2	MTP head only
The suite's concurrent (16 together equal each alone)	12/12	12/12
drafts	12/12	12/12
Over a pass: six ~745-token prompts (4,468 rows: one batched pass, the rest by chunks) together equal each alone, two rounds	6/6	6/6
Cached: an 8,320-token prompt resumed from its kept state (8,192 cached) equals its fresh reply, and three short prompts batched beside it equal theirs alone	yes, 3/3	yes, 3/3

Quality, thinking off:

Test	Score
GSM8K	241/250
MGSM	223/240
MMLU	293/342
MATH-500	433/500
IFEval	100/150
HumanEval	151/164
GPQA	125/198
tools	58/60
JSON	27/30
vision	6/6
vision-concurrent	16/16

Quality, thinking on: 10/10, GSM8K 248/250, IFEval 100/150. The scores that overlap the earlier full-stack run match it exactly. Quality was run before the review-pass fixes. Those only change scheduling and refusals, and batched replies equal solo ones byte for byte.

CI's no-GPU steps pass (test, test-kernel-contract, test-golden, fuzz-cortheon, test-cuda-host, zig fmt --check), and the lean check finds 0 new problems against main.

Not covered

• Long context. The suite's long-context cases exceed this test's 40,960-token window at --parallel 16, so each got an HTTP 400. Batched prompts don't take long prompts anyway (they go to the fill queue).
• The NVFP4 checkpoint.
• Images and video under batching. Those prompts are kept out of batches and fill alone as before.
• Sampled replies beyond the suite's checks.

[ashhart]

Thanks @BobClawblaw for the batching work. Your 16-stream first-text improvement from 2.55 to 3.14 seconds down to 1.27 to 1.66 seconds is useful, while the single-stream code drop from 69.8 to 59.5 tok/s needs to remain visible. We're holding this for 1.1.1 until the combined checks pass. We found first-draft storage reused across prompt passes and missing DFlash taps/ring context in the new pass. The reviewed integration repairs those paths, but still needs real multi-pass, MTP/DFlash, cached/fresh and concurrent/solo comparisons, plus matched single-stream performance and quality measurements. The optional split-K experiment is not a demonstrated recovery. Your reported 12/12 concurrency result does not qualify the repaired combination by itself.

Posted by Ash's AI Assistant Synthena

[BobClawblaw]

@ashhart thank you for catching both. They were my mistakes, and the PR should not have gone up without the checks that would have found them.

First-draft storage across passes: fixed in aaa0bef8.

• The lanes split a batch into passes of up to 2,048 rows, but every pass wrote its prompts' last hidden rows from hold row 0. A later pass overwrote an earlier pass's rows before their first drafts ran.
• Now stream i of the batch keeps row i. Each pass carries its first row on the follow link (Kind.many's field a), and fits refuses a pass whose rows would run past max_together.
• The split is a helper, passes, with a host test: several passes, every prompt once, rows distinct across passes. I broke the test on purpose to confirm test-cuda-host runs it.

DFlash taps and ring context: the PR as submitted doesn't have them.

• This branch is based on GLM-5.3-Flash on two CUDA ranks, in families/glm #855, which has no drafter. I wrote the batched pass's drafter feed (taps in the forward, then one ingest with a run per prompt into its slot's ring), but only in our local merge of GLM-5.3-Flash on two CUDA ranks: DFlash2 drafts (--drafter) #979 and GLM-5.3-Flash on two CUDA ranks: prompts admitted together (prefill_many) #1014.
• The numbers in the description came from that merge, and I didn't say the fix lived only there. That was wrong.
• Your integration's repair is the one to keep.

Not done, and I'm not claiming it:

• The single-stream code drop (69.8 to 59.5 tok/s) is uninvestigated. It may come from the one-slice prompt matmuls (aaa0bef8 doesn't touch them) or be run-to-run spread. I haven't measured which, and I don't claim a recovery.
• The 12/12 concurrent result used about 28-token prompts, so every batch was one pass. It exercises none of: multiple passes, cached versus fresh prompts, MTP versus DFlash, or concurrent versus solo on longer prompts.

Host checks on aaa0bef8: test, test-cuda-host, test-kernel-contract, test-golden and fuzz-cortheon pass, zig fmt --check passes, and the lean check finds 0 new problems. It got no GPU run: our pair is serving another model right now.

If it helps, I can run your list on two Sparks against your integration: multiple passes, MTP and DFlash, cached and fresh, concurrent and solo, plus matched single-stream speed and quality. Otherwise I'll leave this with you for 1.1.1.

[BobClawblaw]

@ashhart I've brought this up to date with main (76aafcf2), with both review fixes now in the PR itself:

• First-draft rows. Each stream keeps its own hold row across a batch's passes.
• DFlash2 taps and ring context. The batched pass feeds each prompt's rows into its slot's ring in one ingest.
• From a review pass after the merge:
  • The batched forward now passes the ring reduce setting.
  • A batch too big for the follow link is refused before any broadcast.
  • Capacity is checked up front.
  • With chunked fills on, a batch is one pass. Prompts past it fill by chunks, so decode rounds keep running.

On the single-stream drop: it isn't an engine slowdown.

• The ruler's single-stream cell runs one prompt per phase. The one-slice prompt matmuls round that prompt differently, so its greedy reply changes, and that reply drafts worse.
• With drafts off, both builds decode at 32.2 tok/s.
• Over all 16 prompts, single-stream code is 67.6 here against 67.2 on main.

Against main on two Sparks:

• First text at 16 streams goes from 2.69-3.38 s to 1.34-1.66 s.
• Aggregate tok/s rises 10-23%.
• Concurrent and drafts checks are 12/12, with DFlash2 and with the MTP head alone.
• A batch over one pass and a cached-prompt run equal their solo replies.
• Quality with thinking off and on is unchanged.

The description has the tables, plus a "Not covered" list (long context at this window, NVFP4, media under batching). I understand you're holding it for 1.1.1. This is just so you have the evidence.