[NSagan271]

What does this PR do?

Chunked prefill under TP, and streaming consumers that take a chunked prefill a chunk at a time. Stacked on #916, which expresses a chunk as one iteration of a graph Loop; this PR replaces its earlier version on #365, and takes over #371 (streaming chunks), whose completion machinery the Loop and producer-determined walks (#385) make unnecessary.

TP: followers run the leader's chunks

#916 runs a TP/SP node's prefills whole. Here only SP stays that way (_unchunkable_nodes, mstar/worker/engine_manager.py); under TP the leader alone cuts chunks and every follower runs the same ones.

• ScheduleTPNode carries the step's chunk ranges columnar, parallel to request_ids: chunk_starts, chunk_ends, chunk_totals and chunk_strides (-1 for a row that isn't chunked, all empty on a step with no chunked row). The wire codec encodes them typed.
• Followers stamp chunk_ranges from the head (MicroScheduler.tp_chunk_ranges) on both the serial and the speculative follow path, and derive the rows mid-prefill from the ranges, as the leader's _assign_chunks does. So a follower's postprocess holds a mid-prefill row's inputs and moves the chunk Loop by the leader's strides. Lists that don't match the head's rows raise before any row is popped.
• On a row's first chunk, a follower caps the Loop at the prompt's length (GraphRuntime.set_loop_caps), as the leader's scheduler does. The cap decides which step the Loop ends on, so every rank has to set the same one. A first chunk is the one whose stride is its end, since the Loop counts from 0 even past a cached prefix.
• On the speculative follow path, N's rows mid-prefill are handled as the leader's merge handles them: kept out of the loop-back prep and the iteration increment, and run from the input held since their first chunk.
• Each rank prepares a chunked row once and slices it per chunk (Engine._prepare_chunk, unchanged), and reserves ahead per row (reserve_whole_prefill) from the same ranges, so the ranks reserve alike. [Preemption 1/7] Slim Admission Control (A1) #869's allotments are off above one rank, so that reservation still matters here.

Streaming: a consumer can take a chunk at a time

Under #916 a chunked producer's streamed outputs leave on the final chunk, and a consumer that counts items joins the pieces into one (Connection.join_dim). A consumer that can use a prefix now takes each chunk's piece as it lands instead.

• Producer: ChunkingPolicy.per_chunk_outputs names the outputs a chunk sends as soon as it completes. _wrap_node_in_loop keeps those edges on the Loop body, where a streaming edge already routes every iteration, as a decode loop's does; no runtime change, Rust or Python.
• Consumer: ChunkPolicy.allow_partial_input() (default False), set with FixedChunkPolicy(allow_partial=True).
• Startup check (mstar/validation/placement.py): a per-chunk edge whose consumer does not allow partial input is refused, as is a per-chunk name that is not a streaming output of the node, and a streaming edge carrying the final chunk's output whose consumer neither joins (join_dim) nor allows partial input. join_dim only joins pieces that arrive in one pass, so it is not enough for a per-chunk edge.
• Qwen3-Omni: the Thinker sends thinker_states / thinker_mask per chunk, and the Talker takes them with allow_partial=True instead of join_dim=0, so it extends its KV as each Thinker chunk lands. Producer-determined walk transition #385's consumer_walk already assigns each pass its Talker walk, so every chunk runs talker_prefill and the conductor's trigger count is unchanged. A chunk with no kept rows (all system prompt or earlier turns) is not sent at all, so the Talker never runs a 0-row prefill.

What #371 needed and this does not: finished_graph_walk on edges, per-row completes_node / completes_walk through routing, the completion rids on RouteArg, the Rust frame fields and the StreamBuffer hold-back. The Loop decides which edges leave on which chunk, and the walk ends when the Loop does.

How was it tested?

On 4037bd2, which merges #916's 6f2d85f (main through #975, including #955), unless marked. This head (8d227ad) merges main through #916 (squashed as 4079d7f), #977 and #980; the diff against main is the same as it was against #916's last head (fb12141). On it this PR's own tests and #916's scheduler, postprocess and chunking tests pass (test/modular/test_{micro_scheduler,tp_follow_mid_prefill,kv_prefix_tp,chunked_prefill,chunk_stream_partial,qwen3_omni_talker_walk,qwen3omni_vision_chunks,stream_chunk_schedule,dropped_row_stream_stop,inline_tokens,spec_chain_key,loop_turns,chunk_loop_stop,admit_failure_handling,kv_offload}.py, test/rust/test_{runtime_parity,postprocess}.py: 630 passed, 37 skipped). A TP2 smoke on b16564a, the same code before #977 and #980 (Qwen3.5-9B, 2 x H200 per server, one session, against main at c6fcbe8) matches the runs below: parity 286 of 288 on both; c48 25.8 against 24.3 req/s, short-prompt TTFT p50 109 against 419 ms, 229 against 391 gaps over 50 ms; burst of 8 at 45.2 against 18.8 decode chunks/s, 0 against 16 stalls, TTFT p50 2.16 against 3.89 s; no failed requests, no hang, no server errors. The Qwen3-Omni runs below predate it. The port makes no Rust change of its own, and the TP-follow path replays the leader's chunks without going through #955's chain pick, so the merge needed no change here.

• ruff check .: clean.
• CPU Core ([Chunked Prefill] Chunked prefill as a graph Loop (replaces #365) #916's command, plus test/rust, OMP_NUM_THREADS=2, GPUs hidden): 221 failed, 9308 passed, 1084 skipped, 20 errors, all among current main's failures (9d161dc) on the same machine; test_eval_compare.py::test_whichever_arm_starts_first_on_cold_caches_still_starts fails on main and passed here.
• GPU tests, on e12da48 before the merge (the same paths plus test/integration, on 3 x H200, torch 2.12.1+cu132): 239 failed, 9570 passed, 530 skipped, 27 errors, every failure among [Chunked Prefill] Chunked prefill as a graph Loop (replaces #365) #916's head's on the same machine. Those fail the same way on main (see [Chunked Prefill] Chunked prefill as a graph Loop (replaces #365) #916).
• New tests:
  • test/modular/test_micro_scheduler.py: a follower stamps the leader's chunks and strides, a step with none, and lists that don't match the head.
  • test/modular/test_tp_follow_mid_prefill.py: a follower's speculative step with a row mid-prefill, on a real Worker.
  • test/modular/test_kv_prefix_tp.py: a follower's first chunk starts past rank 0's leased prefix.
  • test/modular/test_chunked_prefill.py: chunking stays on under TP and off under SP.
  • test/modular/test_chunk_stream_partial.py: per-chunk edges stay on the Loop body, the startup check's refusals, and join_dim still joining a final-chunk edge.
  • test/modular/test_qwen3_omni_talker_walk.py: Omni's real graph passes the check, an all-system-prompt chunk sends nothing, and a chunk sends just its kept rows.
  • test/rust/test_runtime_parity.py: on both runtimes, each chunk's piece routes on its own step and the token leaves only on the final chunk.

Qwen3.5-9B TP2

2 x H200, two sessions per arm, interleaved, prefix cache off, MSTAR_SHM_ARENA=1, Rust graph runtime, a fresh Inductor cache per boot. The workloads are #916's: the 200- and 3000-token prompts are mixed 4:1 with 256 tokens out; bursts of 16 decodes, then 8, 16, 32 and 32 prompts of 16k-30k tokens with 16 tokens out; and the 24-prompt parity set. Arms: this head with MSTAR_TP_ASYNC_SCHED 0 and 1, and #916's head, whose TP prefills run whole (with 0). Cells are session 1, session 2, in ms unless marked.

Workload	Metric	async 0	async 1	#916, unchunked
c48	req/s	18.3, 18.5	23.3, 23.3	18.4, 18.7
c48	TTFT p50, short	184, 182	88, 88	515, 379
c48	ITL p50	7.7, 6.7	4.6, 4.7	7.0, 7.6
c48	ITL p99.9	335, 334	58, 59	382, 328
5/s	gaps over 50 ms	1, 4	3, 1	627, 621
Burst of 8	stalls (gaps over 1 s)	16, 16	0, 0	16, 32
Burst of 8	decode rate, chunks/s	26	56	13
Burst of 16	stalls	16, 16	0, 0	48, 48
Burst of 16	worst gap	4432, 4301	510, 513	1753, 1690
Burst of 32	stalls	64, 72	0, 6	74, 80
Burst of 32	prefill, k tok/s	47	41	57

• Parity: async 1 and the unchunked head give identical greedy outputs on all 24 prompts in both sessions. Async 0 differs from every other run on one prompt in session 1, and matches them in session 2.
• No hang and no failed request in any run, at either setting.
• With async 1, chunking under TP behaves as it does on one GPU ([Chunked Prefill] Chunked prefill as a graph Loop (replaces #365) #916's Qwen3.5-4B numbers): about 55 chunks/s per decode during bursts, and no stalls up to bursts of 16.
• With async 0, chunking removes the 5/s gaps and halves short-prompt TTFT at c48, but not much else. c48 throughput and ITL match unchunked, and bursts still stall. Async 0 also logs about 1,800 "Not enough free pages" decode refusals per server, against 0 and 460 for async 1. The cause isn't traced. My hypothesis is that with async scheduling off a decode can't join a step until the chunk in flight lands, so decodes alternate with chunk steps; I haven't checked this against worker phases. Worker: MSTAR_TP_ASYNC_SCHED on by default #1075 has since made async scheduling the default.

Against main, after #1075

The same setup and workloads, with async scheduling on by default for both arms: this head (d7738b1, before #916's peek gating, which changes nothing under a token budget) against main (9d161dc). The parity set is #916's 288 prompts. Cells are session 1, session 2, in ms unless marked.

Workload	Metric	This head	main
parity, 288 prompts	right	286, 286	286, 286
c48	req/s	25.6, 25.6	24.5, 24.2
c48	TTFT p50, short / long	142 / 247, 142 / 217	345 / 345, 404 / 390
c48	ITL p99.9	105, 105	286, 289
c48	gaps over 50 ms	251, 230	501, 450
5/s	ITL p99	11.7, 11.7	17.6, 17.7
5/s	gaps over 50 ms	5, 5	434, 435
Bursts of 8-32	decode rate, chunks/s	43-45	18-26
Bursts of 8-32	stalls (gaps over 1 s)	0-16	16-112
Bursts of 8-32	prefill, k tok/s	50-51	49-58
Burst of 8	TTFT p50, s	2.17	3.9

Both arms gave identical answers in both sessions. This head had two failed 5/s requests in session 1, both client disconnects 0.2-0.4 ms after sending with nothing in the server log.

Qwen3-Omni, per-chunk Talker

qwen3omni_2gpu on 2 x H200 per server, the Thinker's budget at 16 tokens, so each prompt runs many Thinker chunks and the Talker extends its KV on each. Two sessions per arm: this head (per-chunk) against #916's head (the pieces joined into one item at the end). Speech out with Omni's system prompt.

• Speech matches its text as well as before. Each answer's audio was transcribed with whisper-large-v3-turbo and scored against that answer's own text (benchmark/answer_audio_wer.py), 384 answers per arm (benchmark.runner, the text set, c8, 192 per session). Mean WER is 0.169 for both, with 10 and 11 answers above 0.5 (Fisher p = 1.0), and every answer had audio.
• At 1/s (100 requests), this budget overloads the Thinker on both arms. TTFA p50 is 18-24 s for short and long prompts alike; this head runs 4-5% higher. RTF p50 is 0.42 on both sessions here, against 0.36 and 0.40 on [Chunked Prefill] Chunked prefill as a graph Loop (replaces #365) #916's head. My hypothesis is that the Talker now runs a step per Thinker chunk; I haven't measured it with worker phases. At the default budget (1024) a prompt seldom runs more than a few chunks, so this is an upper bound on the cost.

Open questions

• Prefix reuse under TP: KVManager.resolve_cached_prefix answers 0 until rank 0 has sent the step, so the leader measures a chunked row from 0 the first time it schedules it, and the row loses prefix reuse under TP. The ranks stay symmetric. Rank 0 leasing at measure time would fix it, which is a KV-lease change for its own PR.
• MSTAR_TP_ASYNC_SCHED=1 with speculation_mixes_walks on: a follower takes the speculation target's walk from the head's first row. With mid-prefill rows of several walks in one head, that could differ from the leader's choice. This predates this PR, and speculation_mixes_walks is off by default.

Developed with AI assistance (Claude).

🤖 Generated with Claude Code

[merceod]

Ran Qwen3.5-9B TP2 on klone H200s from the top of the stack with a log line per prepared and per settled chunk on both ranks: over the parity set (112 prompts of 20 to 3k tokens at c=1/8/32, greedy and seeded) both ranks produced 3411 chunk preparations and 3411 settlements in the same order with the same ranges, outputs were deterministic at c=1 and inside the usual batch floor at c=8/32, 70 aborts during or right after chunked prefills left the server clean, and late arrivals were served in under 140 ms with 40 and 64 requests decoding. Approving, with three notes.

• The SP refusal is lifted but there is no SP-capable chunkable model and no SP handling in the chunk path. Keeping the refusal for sp_size > 1 until there is one would be safer.
• tp_chunks zips with strict=False, so a length mismatch between the leader's chunk lists and the request ids would turn into a silent token-count mismatch (a hang on the collective) instead of an error. strict=True costs nothing here.
• Nothing checks that the follower's get_input_sequence_len agrees with the leader's ranges or that the prepared length matches the measured one; the ranks agreed in every chunk I logged, but a mismatch would hang the collective or drift the sampler rather than fail loudly. A cheap assert on the follower (range end <= prepared length) would turn that into an error.

The PR body says more thorough testing is in progress. The runs above cover TP2 parity, cancels and scheduling with async scheduling off, and the same with MSTAR_TP_ASYNC_SCHED=1 at a 256-token budget (11,944 chunk preparations per rank in identical order, 70 aborts clean, outputs unchanged afterwards).

[merceod]

strict=False here hides a length mismatch between chunk_starts/chunk_ends and the rids: it would become a token-count mismatch and a hang on the collective instead of an error. strict=True costs nothing.

[Gaurav-Shah05]

strict=True alone would raise on every step without a chunked row, including every decode-only step, because maybe_send_zmq_to_tp_followers sends both lists empty then. It works if the zip is skipped when the lists are empty.

[merceod]

This lifts the SP refusal, but there is no SP-capable chunkable model and the chunk path has no SP handling. Keeping the refusal for sp_size > 1 until there is one seems safer.

[Gaurav-Shah05]

LGTM. I ran qwen3_5_9b_tp2 and orpheus_tp2 w/ and wo/ MSTAR_TP_ASYNC_SCHED. Qwen3.5 gives the same greedy output on TP2 as on TP1 for prompts from 1.5k to 24k tokens, Orpheus samples the same tokens on TP2 with chunking on and off, and a client that disconnects mid-prefill leaves both ranks clean.

[NSagan271]

Force-pushed: this branch is now rebuilt on #916's new head (aaa4f38), where the chunk Loop counts tokens. The old head was 39715b2. Under TP, followers now also apply the leader's per-row strides (chunk_strides on ScheduleTPNode), and cap a row's Loop on its first chunk, as the leader does. The description has the TP2 and Qwen3-Omni results.

[Gaurav-Shah05]

Second pass after the core PR was changed. Still LGTM.

On 2xH100s against #916's head (6f2d85f), Qwen3.5-9B at TP2 with bursts of 16k-30k-token prompts beside 16 decodes, two sessions each:

Burst	Decode, tok/s per stream	Worst gap, ms	TTFT p50, s	Prefill, k tok/s
8	27 against 2	46 against 2065-3155	2.5 against 4.3-5.4	49 against 35-44
16	25-26 against 5	47-48 against 1627-1843	4.0 against 3.6-3.8	50 against 55-56
32	25 against 5	49-51 against 1867-1896	8.1 against 7.6-7.7	50 against 57

Neither session logs a Not enough free pages refusal, and the 24-prompt parity set gives identical outputs on both heads in every run. Qwen3-Omni T2S at 1/s with the per-chunk Talker brings long-prompt TTFA p50 to 240 and 241 ms against 273 and 281, and short-prompt TTFA p99 to 293 and 311 ms against 392 and 430, with RTF p50 unchanged.