Admission control for KV resources by Gaurav-Shah05 · Pull Request #400 · mstar-project/mstar

Gaurav-Shah05 · 2026-10-10T23:02:32Z

This is a new feature in M*, and it addresses #307. For this implementation, I have to guarantee that a request, once admitted, never runs out of KV pages. The idea is the banker's algorithm with one resource. A request is admitted only if every admitted request, itself included, can still get all the pages it may need. Because admitted requests decode together, they must all be able to finish at the same time, not one after another. That makes this the conservative form.

Problem

On main, a request is admitted as soon as its prompt fits. Nothing checks that the pool can also hold the tokens it will generate. When the pool runs out, the worker pauses the batch or moves a request's KV to host memory. With a shared prefix, this can deadlock. In one test, 32 requests share a 59-page prompt in a 96-page pool, and none of them finishes within 600 s.

How it works

Each model declares, in Model.kv_slots, how many tokens a request writes to each KV label.
count_pages turns that into pages, adding room for max_tokens of output and one speculative step.
Requests wait in a queue in arrival order, and only the first one can be admitted.
It is admitted once its pages fit in the free pages plus the cached pages that can be evicted.
The check is a single comparison against running totals.
A reservation is a promise, not pages. Cached prefixes stay in the pool until a request needs those pages.
When a request is admitted, its cached prefix is held for it, so eviction can't take it before prefill.
Pages that a waiting request will reuse are evicted only after idle pages.
Under TP, rank 0 decides which requests to admit, and the other ranks admit the same ones.
With CFG parallelism, the leader's cache decides for the guidance caches.
A request the pool can never hold fails at ingest, with a 400 that says what to change.

Trade-off

Each request reserves room for its worst case, all of max_tokens, so no admitted request is ever paused. When the pool is tight, fewer requests run at once than on main, and the rest wait.

In the Bagel stress run, chat requests wait behind image requests, and chat TTFT p50 rises from 0.45 s to 17.7 s. In return, no request hangs (57 do on main), and throughput is 1.6x.
Pi0.5 on a quarter pool completes 6% fewer requests per second.

Reserving less would need preemption: pausing a request mid-flight, freeing its pages and resuming it later.

Offload as a decision: rank 0 picks the request to move to host, and host pages count as room.
Backfilling and peak-fit admission, with predicted output lengths (kv: peak-fit admission with EASY backfill, plan cache, parked waiters, and KV release at completion #366).
Prefix caching under TP (GLM-5.2/GLM-5.3 PR 10: engine: prefix cache at TP > 1, recurrent snapshots, speculative sampling #405).

There's one open question that I would like the reviewer's feedback on. Since each model's count is written by hand, I want views on deriving it instead, either from the layout #DIGESTS declares or from each walk's input lengths.

Testing

On CPU, test/modular passes at every commit, apart from the 8 tests that fail on main.
With debug asserts on, six models on a quarter pool completed every request without exceeding their reservations.

H100, tight pool	`main`	This PR
Shared prefix, 96 pages, c32	0/32 in 600 s	32/32 in 8.8 s
Bagel stress, 512 pages, c24	57 hung, 0.263 ok/s	0 hung, 0.428 ok/s	1.63x
↳ chat TTFT p50	0.45 s	17.7 s	0.03x
Bagel chat, no `max_tokens`, c24	5 hung, 0.912 ok/s	0 hung, 4.337 ok/s	4.76x
↳ TTFT p50	0.59 s	2.37 s	0.25x
Chatterbox, ¼ pool, c32	90/80 of 128	128/128 of 128
Higgs-Audio, ¼ pool, c32	0/96	96/96
Higgs-Audio, shipped pool, c32	15.6 req/s	30.3 req/s	1.94x
Pi0.5, ¼ pool, c80	40.2 req/s	37.7 req/s	0.94x
Whisper large, TP2, c16	2.45 req/s	4.45 req/s	1.82x

Where the pool isn't tight, throughput matches main on Orpheus, Cosmos3, V-JEPA2-AC, Chatterbox, Qwen3.5 and Qwen3-TTS.
Engine time per decode step is 1.7% lower on Orpheus and 2.8 to 3.6% lower on Bagel.
On kvshare, the prefix-cache hit rate stays within 0.6 points of main. At 2,048 pages, main deadlocks and this PR completes all 713 requests.

Checklist

ruff check . passes
Added or updated tests / docs where relevant

The pool has to know how far a request can grow before it lets it in. A model that returns no config for a KV cache, as Whisper doesn't, now gets a default one to carry it.

Eviction takes leaves only, so a page only the index holds is out of reach while a page below it is held. Admission counts what eviction can reach, and asks what would match without re-stamping the LRU.

…amps nothing

A live child would keep eviction from a page the pool counts as free, so the room admission counted could vanish after it admitted into it.

…er it fills a page

An image request writes image tokens, latents and guidance copies that no prompt key describes, so the pool can only size it from the model's own count. The conductor hands each cache its share, next to the prefix keys.

A running request keeps needing pages as it decodes, and nothing checked that the pool could supply them: the 96-page repro holds all 32 of its requests. The pool now counts what each admitted request and each padding row may still take against what is free or evictable.

… of growth

Under CFG parallel each worker builds its own cache from a spec naming all three LLM nodes, so a request's reservation counted the guidance nodes' labels on every worker.

Ready takes the reservation, so two requests in one scan can't share the same room, and a request behind one that doesn't fit waits too. Only a request the model or its keys count is held to it; its hit is leased as it is admitted.

The decoder's stream is the forced prompt plus what it decodes, and the audio context is always the whole window, so both caches are counted exactly and admitted by reservation.

A victim that came back into room the running set had reserved was the thrash on #280's shared-host-copy branch. Its reservation now outlives the offload, grown by the pages the index had lent it, which the reload takes back as its own copies.

…meanwhile

A model that stamps no prompt_slots let a client's own reach the pool, which admits by them: a small count takes room others were promised, and a malformed one raises in the gate.

A batched replay gives its padding rows their capture span, so each took a page per label on first replay and kept it: counted as a reservation, that left the 96-page repro 15 pages for requests. The rows' tokens are discarded, which is what the sink page is for. The padding caps go with it.

Bagel counts an attachment's image tokens, which are not positions, so a request with enough images passed max_seq_len and was reserved short.

Every worker builds its cache from one spec naming all three LLM nodes, so no pool decided and Bagel ran there without admission. The spec names its leader; the guidance caches reserve what it admitted and check no room of their own.

… caches follow

No preplanned step carries a lease: a preplan runs only on a batched capture, and every keyed walk is a packed one.

So the stamp every cache admits by is the one its tests exercise, not a copy of the loop.

…ounts

…to the sink Bagel runs its guidance branches for a whole batch once one row needs them, so a text request prefilled beside an image request wrote its prompt to cfg_img as well and forked it onto cfg_text: twice or more the pages its reservation counted, which the S1b gate caught. The row never reads those labels back. A label outside the request's prompt_slots now plans onto the sink page, takes no page, keeps no length, and is not forked onto.

stephen-dwq

The admission gate as it is loses the benefits that paged-KV provides with respect to internal fragmentation of per-request allocations: because admission gates on worst-case request size, we are logically allocating (reserving) much more KV space than a request is likely to use. This results in severe internal fragmentation and poor KV utilization, as well as low concurrency.

I have two solutions in mind:

We admit a request if there is enough free KV such that it's entire prefill fits, and each active request still has a free page; in the event that we still run out of room, we have a preemption system that offloads/drops a request (choosing metric TBD) so that other requests can continue; this is contingent on a working preemption system
We can also adopt a method of 'pausing' a request, so that a sequential bankers algorithm works: we would have a buffer of KV enough that the the active request with smallest remaining worst-case space can finish alone, and if we reach the point where we would tap the buffer, we pause everyone else and allow this request to finish executing and release its pages, and then resume concurrent operation until this happens again

Gaurav-Shah05 · 2026-10-20T04:29:39Z

@stephen-dwq We should measure how the fragmenation and KV utilisation w/ concurrency change when admission control is turned on. trt-llm's admission control policy is also simillar: the default policy admits requests when they can reserve their worst case, and it doesn't preempt. Nonetheless, I agree that we should look for a better option, and I am currently leaning towards supporting preemption recompute machinery over allowing requests to pause mid-decode.

Here is how I'm thinking it would work. Requests will fall into two classes. Strict requests, like the ones streaming audio out, plus, for now, requests with images or audio in their KV, which can't be rebuilt yet. They still reserve their worst case and never pause. Every other request is preemptible: it can be preempted and rebuilt later, so it reserves only its prompt, which the one part whose size is known at admission. When a request needs a page, it takes a free one, or the prefix cache evicts its oldest page that no running request holds. Only when neither works is the most recently arrived best-effort request preempted:

Its pages become cache, and only what the shortage needs is evicted, starting from its tail.
It's requeued under its arrival time. Admission is FIFO, so it lands at the head. Running requests keep running, and only new arrivals wait behind it.
When pages free up, it's readmitted: a cache probe leases its remaining pages, and its lost tail is prefilled in the next step.

For this to work on every model, the counts have to be right, and the hand-written per-model counts in this PR are flaky and brittle, which is a known limitation. Instead, each model will declare a per-request layout of what each walk writes into each KV label. Admission counts, prefix keys and whether a request can be rebuilt all come from it. A new model then gets admission and preemption without touching the engine.

Preempting is cheap but a rebuild is only as cheap as what the victim still has cached. Whatever that was evicted gets recomputed, and that prefill takes part of the next step's token budget from the running requests. So I'd add a host tier for the KV cache which is a needed feature for M* anyways. Pages are copied to the host ahead of eviction, keyed by prefix, and reloaded asynchronously when the victim is readmitted. Pages that are never evicted never cross PCIe, and the same keys let disk or remote tiers plug in later.

And for the same reason, the remaining prefix caching follow-ups should land before preemption. Qwen3-Omni with its decode keyed, Qwen3.5 after the recurrent snapshots from #405, and encoder outputs cached by digest.

Prefix caching at TP > 1 doesn't depend on this PR, so I'm moving #408 onto main without it, with @vasilevklart's fixes on top. If his GLM TP8 run still diverges, the behavior can exists for now until till this work finishes (and I think it will by the time GLM stack is reviewed and tested)

This PR stays parked for now.

The gate refused a row whose KV did not fit, intending to let it sit a step out rather than refuse the whole step. Two reports show it does not hold up: - Decode deadlocks. Every decode row needs a page for its next token, so a pool held by decode rows refuses all of them; no step is built, nothing runs, and nothing is ever evicted or offloaded to break the cycle. The gate runs before the admit-error path, which is the only thing that triggers an offload. - Every row of a pass sees the same free count, with no running sum, which is what let eight chunked prompts in at once on a 48-page pool. has_room does not exist on main; it arrives with this PR. So dropping it returns to main's behaviour -- the admit-error path refuses the step, OOM backoff holds rows, and the retry is smaller -- and cannot regress against it. Real per-row admission control belongs in #400, which reserves a request's whole uncached prompt up front; this gate was a stop-gap for it. Keeps the declaration itself: get_input_sequence_len, InputSeqLenInfo and engine.input_sequence_len stay, since chunked prefill sizes its chunks off them. Only the readiness gate goes. Co-Authored-By: Claude <noreply@anthropic.com>

A chunk admitted only its own tokens, so on a tight pool several long prompts got in at once, each held the pages of the chunks it had written, and all waited on pages the others held (the 48-page wedge in review). StepContext.chunk_tokens_after carries the tokens a chunked row runs after this step, and KV reserves them with the step: the first chunk admits the whole walk, as an unchunked prompt does on main, and later chunks take no more. A prompt that does not fit is refused before it holds anything; admission control proper is #400's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

arun-elr

took a pass w/ a TP lens now that #835 is in. today the gate's a no-op above 1 rank (_keyed_label returns None w/o an index), but after a rebase on main it won't be:

rank 0 leases the hit at readiness in _try_reserve, and only rank 0 holds it. w/ #835 followers adopt at send, so a follower that evicted those pages in between exits in adopt_cached_prefix. maybe reserve the full count above 1 rank and let decide_cached_prefix take the lease?
the debug reservation asserts run on followers too. if one trips there, only that rank raises and rank 0 hangs in the collective
nit: test_qwen3_omni_kv_slots doesn't collect on a cpu box w/o flashinfer (talker.py imports it at module level)
small overlap w/ #754 in _unready_backlog_rids, one line either way

ran the 24 touched test files on its head, both graph runtimes: 147 pass, 4 fail here on missing torchcodec + the collection error above

arun-elr · 2026-10-24T00:26:04Z

+        if self._outstanding + need > self._supply(leasing=hit):
+            return _WAIT
+        if hit:
+            self._lease(rid, stream, rooted)


self._lease(rid, stream, rooted)

above 1 rank this lease lives on rank 0 only. w/ #835 followers take hits at send, so a follower can evict these pages first and then exit on adopt. maybe skip the lease when world size > 1 and reserve w/o the hit?

this pr will be closed soon, and its accounting is moving into the slim version of it. I'll handle this there.

sounds good, I'll take a look at the slim one when it's up and rerun the TP cases there.

arun-elr · 2026-10-24T00:26:04Z

+                    rid: (held[rid], res.pages)
+                    for rid, res in self._reserved.items() if held[rid] > res.pages
+                }
+                assert not over, f"requests took more pages than they reserved: {over}"


assert not over, f"requests took more pages than they reserved: {over}"

this (and the outstanding check below) runs on TP followers too (_leads). if it fires on a follower, only that rank raises and rank 0 waits in the collective. maybe use the SystemExit #835 uses for divergence?

In the slim version of the admission control pr which I'll open soon, this will exit the same way as 835

The gate refused a row whose KV did not fit, intending to let it sit a step out rather than refuse the whole step. Two reports show it does not hold up: - Decode deadlocks. Every decode row needs a page for its next token, so a pool held by decode rows refuses all of them; no step is built, nothing runs, and nothing is ever evicted or offloaded to break the cycle. The gate runs before the admit-error path, which is the only thing that triggers an offload. - Every row of a pass sees the same free count, with no running sum, which is what let eight chunked prompts in at once on a 48-page pool. has_room does not exist on main; it arrives with this PR. So dropping it returns to main's behaviour -- the admit-error path refuses the step, OOM backoff holds rows, and the retry is smaller -- and cannot regress against it. Real per-row admission control belongs in #400, which reserves a request's whole uncached prompt up front; this gate was a stop-gap for it. Keeps the declaration itself: get_input_sequence_len, InputSeqLenInfo and engine.input_sequence_len stay, since chunked prefill sizes its chunks off them. Only the readiness gate goes. Co-Authored-By: Claude <noreply@anthropic.com>

A chunk admitted only its own tokens, so on a tight pool several long prompts got in at once, each held the pages of the chunks it had written, and all waited on pages the others held (the 48-page wedge in review). StepContext.chunk_tokens_after carries the tokens a chunked row runs after this step, and KV reserves them with the step: the first chunk admits the whole walk, as an unchunked prompt does on main, and later chunks take no more. A prompt that does not fit is refused before it holds anything; admission control proper is #400's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Gaurav-Shah05 · 2026-10-25T23:23:03Z

opened updated version in #869

* Admission by row: leave out a row whose KV does not fit - NodeSubmodule.get_input_sequence_len(walk, fwd_info, inputs) reports a row's token count from input shapes, before prepare_inputs; None (the default) keeps today's behaviour. InputSeqLenInfo can give a resource a different segment length than the row's. - Resource.has_room(rid, node, walk, segment_len) answers whether a step adding that many tokens could be admitted now; KVManager counts the pages each needed label grows by against free pages plus the pages the prefix index alone holds. - StepRunner.admit_retrieve / Engine.check_ready take the optional InputSeqLenInfo; no room is reported as not ready, not as a failure. - MicroScheduler measures popped rows of nodes that report lengths: a fresh row without room goes back to its ready queue, a backlogged one stays parked, and the rest of the step runs. The check is per row: a step of rows that each fit can still fail admission, which the existing hold path handles. - Orpheus (prefill, decode) and BAGEL (prefill_text, decode) report lengths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: report Thinker row lengths for admission by row Every Thinker walk's KV span is known from its inputs before prepare_inputs: text from its token ids, audio and vision from the encoder's embeds plus the two sentinels around them, decode one token. The embeds arrive on the Thinker's input edges, so admission can leave out a row whose KV does not fit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: report LLM row lengths for admission by row Text prefill rows are their token count and decode rows one token, known before prepare_inputs, so admission can leave out a row whose KV does not fit. A vision row's span depends on its image grids; it is admitted whole. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Backlog: measure a parked row once and check it once per pass A backlogged row was checked for readiness, then measured and checked again with its length, on every pass it waited -- every speculation step under a split decode set. Its inputs do not change while it waits, so the batch keeps each row's length, and the one readiness check takes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Update docs * Drop the has_room admission gate from scheduling The gate refused a row whose KV did not fit, intending to let it sit a step out rather than refuse the whole step. Two reports show it does not hold up: - Decode deadlocks. Every decode row needs a page for its next token, so a pool held by decode rows refuses all of them; no step is built, nothing runs, and nothing is ever evicted or offloaded to break the cycle. The gate runs before the admit-error path, which is the only thing that triggers an offload. - Every row of a pass sees the same free count, with no running sum, which is what let eight chunked prompts in at once on a 48-page pool. has_room does not exist on main; it arrives with this PR. So dropping it returns to main's behaviour -- the admit-error path refuses the step, OOM backoff holds rows, and the retry is smaller -- and cannot regress against it. Real per-row admission control belongs in #400, which reserves a request's whole uncached prompt up front; this gate was a stop-gap for it. Keeps the declaration itself: get_input_sequence_len, InputSeqLenInfo and engine.input_sequence_len stay, since chunked prefill sizes its chunks off them. Only the readiness gate goes. Co-Authored-By: Claude <noreply@anthropic.com> * Scheduling: try the next node when a step comes out empty Only the first round-robin pick was tried, so a node whose step came out empty left the worker idle for the pass -- and because such a node is never marked scheduled, round robin picked it again on every later pass too. Under a colocated model that starves the node which would unblock it: on orpheus_colocated, snac_decoder waits behind the LLM, and the LLM's pages only come back once snac_decoder finishes the request. _select_node_rr becomes _node_rr_order, returning the whole order rather than its minimum. A key already tried off its backlog above is skipped; a key that declines takes no rows, so moving on costs nothing. Co-Authored-By: Claude <noreply@anthropic.com> * docs: tighten get_input_sequence_len, add per_request_info to declare_step The length a model reports is the whole input; the engine is what subtracts the cached prefix, so say so rather than leaving the reader to guess which side cuts it. None covers a walk that declares no length, not only absent inputs. Drops the claim that the scheduler uses it to leave out a row whose KV does not fit -- that gate is gone. declare_step grew per_request_info when the non-final-chunk sampling fix landed (#353); the signature here never caught up. Co-Authored-By: Claude <noreply@anthropic.com> * BAGEL: an absent text input is undeclared, not a KeyError get_input_sequence_len's contract is to return None when the walk declares no length, and every other model that implements it -- orpheus, qwen3_5, qwen3_omni -- guards the key before reading it. BAGEL was the one that indexed blind, so a prefill_text row that reached here without its text input took down the measuring pass instead of simply going unmeasured. Co-Authored-By: Claude <noreply@anthropic.com> * Pre-plan a speculative packed step from its rows' spans A packed step (a mixed or prompt step) had no token count before prepare_inputs, so it could neither lease a capture nor pre-plan, and planned inline on the gpu thread behind the step in flight. Its rows' spans are now measured when the speculation is armed: a row in hand from its inputs (get_input_sequence_len), a row continuing from the step in flight from the submodule's prediction (speculative_input_seq_len, None by default). The step leases its packed bucket on those spans and declares over host-only rows of them (declare_inputs_from_spans, None by default), on the plan thread like a batched step. The resources promote a stage by its rows and slot, not its spans, so after prepare_inputs the engine compares the prepared spans with the planned ones; a mismatch (a cached prefix cut a prompt, a row left) drops the stage and the packed lease, and the step leases and plans inline. MSTAR_PRE_PLAN_PACKED=0 turns this off. Qwen3.5 opts in for its mixed and text prefill walks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus and BAGEL: pre-plan a speculative packed step from spans Both declare a text step from each row's span and walk; BAGEL also from its guidance flag, so declare_inputs_from_spans now takes each row's forward info. A step declared from spans whose capture key is not its lease's is not pre-planned, so a hook that misreports guidance cannot stage a plan for the wrong capture. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Pre-plan: estimate a continuing row's span from its last one speculative_input_seq_len takes the row's span in the step in flight (-1 if it ran on another node), and by default repeats it, the rule chunked prefill's merge budget uses; a walk that fixes its span says so (decode: one token, which a prompt's last chunk continuing as decode would otherwise get wrong). ExecutingBatch.seq_len_of reads a prepared step's row spans, for the worker to hand over. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: run a long input across several steps - NodeSubmodule.supports_chunked_prefill(walk) / max_batch_tokens(walk); NodeInputs.chunk_start / chunk_total / is_final_chunk. A chunked walk must report get_input_sequence_len; chunking under TP/SP is refused at load for now. - MicroScheduler keeps chunk_progress per (rid, node). FIFOBatchBuilder divides max_batch_tokens max-min fairly (fill_token_budget): rows under an equal share go in whole, chunkable rows split what is left, others sit out; a step is never empty. Chunk ranges ride on ScheduledBatch, through the speculation merge too; a chunkable row's room check is for one step. - Engine prepares a chunked row once, cuts each chunk with split_inputs, and keeps the whole input until the final chunk lands, so a refused admit retries without re-reading advanced counters. A non-final chunk's row is left out of the sampler's kept rows. Chunked rows skip the prefix cache for now. - Worker leaves a non-final chunk's row out of input cleanup, stop checks and routing: its outputs wait in ChunkOutputAccumulator (its own file) and its node goes back to the ready queue once the chunk lands, so chunk k+1 never runs before chunk k; the final chunk routes the concatenation. A mid-prefill row never continues a speculation chain. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus: chunk long prefills prefill can be split; max_batch_tokens is 512 for prefill and mixed steps. postprocess drops a non-final chunk's sampled token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: chunk long Thinker text and audio prefills prefill_text and prefill_audio can be split, with a 1024-token budget for them and thinker_mixed steps. split_inputs also cuts masks_for_talker on its sequence dim; row lengths come from the Thinker's get_input_sequence_len (cp/2); only the final chunk of the last prefill keeps its token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Refine estimated sequence length of continuing requests; add output accumulation policy * Qwen3-Omni: chunk long Thinker vision prefills prefill_vision can be split like text and audio, with its own 2048-token budget (vision steps run alone, at bs=1). split_inputs also cuts the per-layer deepstack features and the talker masks with the chunk's tokens, and refuses tensor inputs it does not know how to cut. A vision span advances the MRoPE counter by its 3D-grid extent, not its token count, and its chunks' positions are cut from the whole span's, so only the final chunk declares the advance. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a final-chunk-only output mode ChunkedPrefillOutputMode.FINAL keeps an edge's final-chunk tensors alone; earlier chunks' are dropped at hold time rather than held and discarded. It does not replace a model dropping a non-final chunk's sampled token in postprocess, since check_stop reads outputs before they are held. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: chunk long text prefills prefill_text can be split, with a 1024-token budget for it and mixed_text (placeholder). A chunk keeps its row's requires_cfg flag. A guided prompt's pre-forward main -> cfg_text fork (cfg_text keeps the pre-text context) happens on its first chunk only, via KVStep.fork_rids, and only the final chunk's sampled token is kept. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: start a chunked row past its cached prefix The scheduler's first measurement of a keyed, chunkable row resolves its cached prefix (the KV lease, under the manager's lock, as admit_retrieve does) and reports only the rest; the row's chunks then start past it, and its first chunk applies the prefix with the whole prepared input. A model can refuse a row before it is prepared (reuses_cached_prefix; BAGEL refuses a guided prompt, which writes two labels). Also gives test_kv_prefix_stage's stub batch the wrapped request info. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Worker: read the speculated step's row lengths off its executing batch The speculation merge read them off the scheduled batch, which has no node_batch, so every speculative step raised and failed its requests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: a prefill step planned as a decode keeps every token An eager step with one token per row gets FlashInfer's decode plan, which sets no qo_indptr; select_last_hidden then failed, and BAGEL's copy of it asserted. Chunking makes such steps common (a one-token final chunk). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a serving config can set token budgets max_batch_tokens: per node, one budget for every walk the model budgets, or per walk; null turns chunking off. Refuses unknown nodes and walks, and walks the model sets no budget for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: correct the FINAL mode's note on sampled tokens Held rows leave the batch before the stop check, and the engine takes non-final chunks out of the sampler's kept rows, so neither reads a non-final chunk's token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * ruff * Qwen3.5: chunk long text prefills prefill_text can be split, with a 1024-token budget for it and mixed steps, under the largest prefill capture with a mixed step's decode rows. Chunks resume the GDN state from the request's slot and take positions from the counter, as successive prefill steps already do; only the final chunk of the last prefill keeps its token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: one chunking policy per walk NodeSubmodule.get_chunking_policy(walk) returns a ChunkingPolicy (chunkable, max_batch_tokens, output_policies, reuses_cached_prefix) in place of four separate hooks, so a submodule declares a walk's chunking in one place. The engine asks once per (node, walk) and keeps its accessors. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a row's progress is a ChunkProgress(done, total) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: final-chunk sampling is the submodule's call keep_final_chunk_samples(sampler_step, request_ids, inputs) drops the rows mid-prefill from a SamplerStep's kept rows. An autoregressive submodule that chunks calls it in declare_step; the engine no longer rewrites every step, since not every chunked node samples. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Recurrent state: scope a step's forks to some of its rows RecurrentStep.fork_rids, as KVStep's: the rows the step's forks apply to, None for every row. Nothing forks a recurrent pool yet; without it, a chunked prompt's later chunks would copy the source over the fork target again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: the runtime never speculates into a chunkable walk A row entering a chunkable (node, walk) is measured and cut when it is scheduled. A continuing row's inputs are the in-flight step's outputs, so scheduled ahead its first prefill would run whole, past the token budget. The worker loads its submodules first (EngineManager.load_submodules), derives the chunkable (node, walk) pairs from their chunking policies, and hands them to the graph runtime's constructor as disable_spec_node_walks; speculate_node, Python and Rust, never returns one as a target. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: the chunking engine fake follows _SizedEngine's rename Dropping the has_room gate one branch down renamed _SizedEngine's `too_big` parameter to `unready`, and the stack rebase replayed this branch's commits on top without updating the call site -- the two lines never overlap textually, so git merged it clean and every _ChunkingEngine test died in construction with a TypeError rather than on anything it meant to assert. Co-Authored-By: Claude <noreply@anthropic.com> * Chunked prefill over combined walks: tests follow the own-walk merge and the uniform path A speculation merge takes only its own walk, so the budget test's continuing row is a prompt. The uniform-inputs fakes carry chunk ranges and row lengths, and a uniform step records no row lengths: one token a row is what an unrecorded row is taken for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: main's new engine fakes carry the chunking fields test_dropped_row_stream_stop, test_kokoro_chunk_loop and the Cosmos3 step guard build an Engine by hand; chunked prefill reads its per-walk policy cache and each batch's chunk ranges, so they get both. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the packed pre-plan and the row lengths admission never read Pre-planning a packed speculative step from row spans measured ahead measured slower than planning it inline: at c32 on Qwen3.5-0.8B it saved ~196 us/iter of inline planning and added ~278 us/iter of plan-thread wait, with no gain at c64 or c128. This removes it (MSTAR_PRE_PLAN_PACKED, declare_inputs_from_spans, measure_ahead_spans and the packed-lease bookkeeping), and the per-row InputSeqLenInfo passed through check_ready to admit_retrieve, which ignored it. Kept for chunking: speculative_input_seq_len (it sizes the token budget left beside continuing rows), ExecutingBatch.seq_len_of and the scheduler's measured row lengths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Speculation merge: estimate the continuing rows' spans only when sizing a budget The merge estimated every continuing row's span (a submodule call per row, plus a dict over the in-flight batch) before scheduling, though the token budget only needs it once fresh rows are merged in -- which a steady decode merge rarely has. get_next_batch now takes the estimate as a callable and asks for it only then. About 8 us/step at c16 and 50 us/step at c128 off the speculate phase on Qwen3.5-0.8B. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: skip the per-row chunk bookkeeping on steps without chunks A split copies row_lens only when it holds some; the speculation candidates and input cleanup check incomplete_node_rids only when a row is mid-prefill; the batch builder collects its backlogged rids only when the token budget dropped a row; and a speculation copies chunk_ranges and incomplete_node_rids only when they are set. All per-row work on every decode step that never chunks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: admit a chunked prompt's whole walk on its first chunk A chunk admitted only its own tokens, so on a tight pool several long prompts got in at once, each held the pages of the chunks it had written, and all waited on pages the others held (the 48-page wedge in review). StepContext.chunk_tokens_after carries the tokens a chunked row runs after this step, and KV reserves them with the step: the first chunk admits the whole walk, as an unchunked prompt does on main, and later chunks take no more. A prompt that does not fit is refused before it holds anything; admission control proper is #400's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: fail a row whose prepared input is not the length it was cut by The scheduler cuts a row by the length get_input_sequence_len reported, and nothing checked that prepare_inputs agreed (PR #365 review): a mismatch would drop tokens or run a chunk past the end. chunk_ranges now carries the measured total with each [start, end), and the first chunk's prepare compares it with the prepared input; a mismatch raises, which fails that row and lets the rest of the step run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Positions: fork only the rows KV forked KVStep.fork_rids scopes a step's forks to some of its rows (a chunked guided prompt forks on its first chunk and its last), and KV honoured it, but positions applied every fork to every row: a guided row on a later chunk, sharing a step with another prompt's first chunk, had its cfg_text counter re-forked to main's while its cfg_text KV stayed put (PR #365 review). KV's plan output now carries fork_rids, and positions fork the same rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a step of held chunks still routes the rows the cache served _settle_chunks ended postprocess when no row that ran was left, but #768 routes a row the prefix cache served whole, which is not among them: a step holding only non-final chunks beside a served row would never complete that row's walk. It now counts cached_rids, as main's own check after the stop check does. Adds the _settle_chunks tests the review asked for: holding a non-final chunk, joining and releasing on the final one, and the early end. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: join a final chunk's outputs on the side stream, after step N only The join ran on the current stream, behind step N+1 that the gpu thread had already launched, while every copy of its result -- the check-stop D2H and the store's -- waited only on N's completion event and could read it before it was written (17 of 64 Qwen3-Omni answers with broken audio over SHM at c16, PR #365 review). The join now runs on the worker's side stream after N's event; the current stream waits on it, the joined tensors are recorded on it, and the step's completion event moves past the join, so those copies wait for it without waiting for N+1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: per-walk budgets inside a combined walk, behind can_override_mixed_max_tokens A walk inside a combined walk runs in the combined walk's steps, and the scheduler only ever read the combined walk's budget, so a per-walk override did nothing -- and the overrun warning then compared each step against it (PR #365 review). Now: - By default a walk's own budget applies only to a walk outside any combined walk, and a per-walk override of one inside is refused at load. The combined walk can be named directly (LLM: {mixed: 512}). - A node listed in the serving config's can_override_mixed_max_tokens caps each member walk's rows within a step by that walk's budget, its own or the model's, under the step's budget; null keeps that walk's prompts whole. - The overrun warning compares a step against the budget it was scheduled under, the combined walk's for a step relabelled to one walk, and names both causes of an overrun. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: chunked prefill in adding_models Step 4 The hooks a walk opts in with (get_chunking_policy, get_input_sequence_len, split_inputs), how a chunked row's outputs are held and joined, and the two things a chunking AR node does itself that nothing checks at load: wrap its sampler step in keep_final_chunk_samples, and drop a non-final chunk's token or mark its edge FINAL (PR #365 review). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: a chunked row refused, cancelled or failed between chunks The paths the review found untested: an OOM refusal holds the row and its retry runs the same range, since nothing landed; and a request removed mid-prompt leaves no progress, cached prefix, held outputs or prepared input behind in the scheduler or the engine. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: a 2048-token default chunking budget The 1024-token budget was sized to fit a 2048-token prefill capture, but the capture ladder now runs to 8192. Swept on the stack (long mix of 4 ~200-token prompts per ~3k-token one, two restarts each), 1024 cost 3-10% closed-loop throughput at c48 (0.8B the most) in extra steps and fewer speculated ones; 4096 recovered it but gave back most of the short-prompt TTFT gain and doubled decode ITL p99 at c48 on 4B and 9B TP2. 2048 keeps throughput within about 1% of chunking off with most of the TTFT gain. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: chunked prefill's state in ChunkingManagement The engine held seven attributes for chunking and combined-walk budgets: the per-walk policy cache, the serving config's budget overrides, the combined walk of each walk, the nodes whose walks cap their rows in a mixed step, the walk-cap cache, the overrun counter, and the whole prepared input of each row mid-prefill. They and the methods that only touch them (policy lookup, budgets, walk caps, the overrun warning, releasing a row's input) move to ChunkingManagement in mstar/engine/chunking.py, as SubmoduleManagement holds a node's. The engine keeps one-line delegates for what the scheduler and worker call. An engine fake now sets one attribute instead of a list that grew with every change to the engine's chunking. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a cut row runs at least a floor of tokens When decode rows used up a step's budget, max-min fair share cut a prompt to a token or two a step (or, after a speculation merge's continuing rows spent it all, to the one-token head rule). A cut row now runs at least min(128, the node's configured budget) tokens, or all it has left if that is fewer. When the share is under the floor, only as many rows as what is left holds are cut, sharing it; if it holds none, one row runs over budget by less than the floor. Rows that fit whole -- decode rows, short prompts, a last chunk -- are untouched, and check_budget allows the floor's slack. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: vision prefills take the Thinker's one chunking budget A Thinker step of any walk runs in thinker_mixed and so under its budget; the separate vision budget was read only as a can_override_mixed_max_tokens cap, where it sat above the step budget and did nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a row mid-prefill takes its next chunk in the speculative step The speculation merge builds step N+1 while N runs, and a row whose chunk is in N has its node off the ready queue until that chunk lands, so beside decode rows a prompt ran a chunk every other step. The merge is now handed N's rows that need another chunk (MidPrefillRow, with where the next chunk starts: the end of N's) and the builder takes them first, without a pop. When N's chunk lands, a row whose next chunk is in the step launched after N stays off the ready queue; any other goes back on it, so a speculation that is cleared, or a step refused at admit (which re-queues all its rows), needs no undo. A row the merge leaves out is neither parked nor returned. Only a merge that may take the row's walk does this: with speculation_mixes_walks off, a decode chain keeps to decode rows. From Gaurav Shah's version on #430, with the next chunk's start passed in rather than kept as a second progress cursor. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Respond to bot review comments: TP/SP refusal, stream block reuse, documentation * [Graph][Test] Parity: a Loop output read off an EMPTY_DESTINATION body edge Chunked prefill as a Loop gives the chunked node no loop-back edge and no destination of its own, so the Loop's declared output is the only thing that carries a chunk's result on. Nothing covered that shape: the existing loop cases all cache from a loop-back edge. The two runtimes cache loop outputs from different sources -- Python from the node's declared `outputs` list, Rust from the routed-edge slice -- and an edge that routes nowhere is where those sets could differ. They don't, but an empty `Loop.outputs` would drop the final chunk silently rather than raise, so it is worth a test rather than a reading of both. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Declare chunking on the model, and run a chunkable node in a Loop Chunked prefill's state lives in six places and each model implements four to six hooks for it. Expressing a chunk as one iteration of a graph Loop removes most of that: the Loop already owns per-request iteration state, re-queues its node, accumulates or drops per-iteration outputs, and is speculated like decode. This is the first step -- the declaration and the wrapping -- with no behaviour change. The policy moves from the submodule to the model because the graph build needs it: `get_graph_walk_graphs` is read long before any submodule loads, so a submodule-level answer is unreachable at the point the node has to be wrapped. A policy now names the walks it covers, and the model passes its walk dict through `apply_chunking`. `apply_chunking` is called by the model rather than applied to its result from outside, since `get_graph_walk_graphs` is what declares the graph and some twenty call sites read it directly. Forgetting the call would report a prompt done after its first chunk, so `_check_chunkable_nodes_are_looped` refuses it at startup instead. Policies and budgets are set separately: the policy is the model's and fixed, the budget is the deployment's and layered on top. The Loop holds `max_iters=1` for now. The worker still drives chunking and holds a non-final chunk's row back without completing its node, so the loop is only ever completed once, by the final chunk. It becomes a real bound when the engine takes over ending the loop. Nested chunk loops are refused at graph build: a chunkable node inside an existing Loop orders its completion cascade differently in the two runtimes. No model declares one. Tested on 2 x H100 (CPU-only paths; the box's GPUs were in use): - `ruff check .` clean. - `test/modular` against this branch's base (f28ce39a), both CPU-only: 6106 passed / 11 failed / 1 error, against 6104 / 11 / 1. Identical failure sets, all pre-existing. The two extra tests are an untracked shell script in the working tree that `test_launch_script_socket_prefix` parametrises over. - `test_startup_graph_validation`: 95 shipped configs pass, which is what checks that `apply_chunking` wrapped the right node in each of the four chunkable models. - Graph-runtime cost of the added Loop entity, interleaved A/B in one process on the Rust runtime, 4000 requests x 9 rounds per arm: -0.07 us/request against a 4-5 us per-arm spread, on ~21 us of graph work per request -- below the measurement floor. End-to-end still to run on a quiet machine. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Move a chunked node's outputs onto its Loop, and seed the loop-back A chunk is one iteration of the node's Loop, so the Loop is what should decide what a chunk's output means. The node's real edges move onto it: an edge the policy gathers collects every chunk's tensors, and any other carries the final chunk's alone, earlier ones dropped as the loop advances. That is what the worker's accumulator and side-stream join do by hand today. The body keeps a bare same-named edge to EMPTY_DESTINATION, because the Loop caches by edge name and so needs the name emitted. The copy drops every flag: persisting a per-chunk tick would stage a copy of each chunk for a remote reader, and counting one would bill the conductor a token per chunk. It also routes `chunk_tick_name` back to itself. Speculation offers a next-iteration target only for a node that routes an edge back to itself, so without it every chunk-only step is a pipeline bubble. A loop-back has to be declared as an input too, which leaves the node unready for its FIRST chunk -- nothing has produced it yet -- so the conductor seeds one on every entry into the walk. Re-entry included: a walk entered twice (BAGEL think-mode) takes its inputs through `_send_partition_inputs`, not `initial_inputs`, and a seed on the first entry alone would hang the second. The tick is free: the submodule produces no tensor for it, and a declared signal with no tensor still satisfies readiness. `max_iters` is now the policy's `max_chunks` rather than 1. It is a runaway guard, not the real count -- the engine ends the loop on the final chunk in the next commit -- and a prompt cannot have more chunks than tokens, so the model's longest prompt is a safe value. Tested (CPU-only; the box's GPUs are in use by other tenants): - `ruff check .` clean. - New parity cases over the real `_wrap_node_in_loop` output, not a copy of it, on both runtimes: first-chunk readiness, that the seed is what provides it, next-chunk readiness and speculability (`is_new_loop_iter=True`), that only the final chunk's token leaves the loop, and walk re-entry. - `test/rust/`: 39 failed / 426 passed, against 39 / 416 on the branch base -- same failures, all pre-existing, plus the 10 new tests. - `test_startup_graph_validation`: 95 shipped configs pass. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Let the Loop drive chunking, and drop the worker's parallel machinery A chunk is one iteration of the node's Loop, so the Loop decides when the prompt is done and what a chunk's output means. The worker's own version of that goes: the output accumulator, the side-stream join, the push-back onto the ready queue, and the surgery that took a non-final chunk's row out of the batch before postprocess could see it. `Worker._end_finished_chunk_loops` ends the loop for a row that just ran its last chunk. Keyed on the walk having a Loop rather than on the row having a chunk range: a row with NO range ran whole, which is also a final chunk -- a walk with no token budget is never measured or cut, and without this its Loop ran to max_iters while the request held its KV, which looks like pool exhaustion rather than a scheduling bug. What the Loop does not absorb is ~10 lines: until its last chunk a row's sampled token is not a real sample, so it decides no stop (a spurious one would cut the prompt short mid-prefill) and keys no prefix. That was enforced before by removing the row from the batch; now it is stated where the reason is. A streaming consumer needs the pieces joined, which is `Connection.join_dim`. Its ChunkPolicy counts ITEMS and a chunked producer sends one per chunk, so to Qwen3-Omni's FixedChunkPolicy(1) Talker the pieces looked like separate items and it read only the first -- which was empty, the opening chunk being all system prompt, every position of which the Thinker drops. Joined on arrival at the consumer, which owns those tensors outright: the producer-side join this replaces cat'd on the default stream behind step N+1 while a reader waited only on N. A consumer that can take the pieces as they come should do that instead and set no join_dim. The join registers in the producer's units, one per piece, because that is what `signal_done` reports. Counting a group as one left the buffer permanently short of the producer's total, so it never agreed the producer had finished and the Talker's speech came out truncated. Tested on 2 x H100 (slurm), CPU suite on a quiet box: - `ruff check .` clean. - `test/modular test/rust` against the branch base: 6536 passed, 46 failed, 1 error -- NO new failures; all 39 test/rust and the rest are pre-existing. Four base failures pass here: timeout-based tests that fail only under load (one reads 18s quiet, 180s under a concurrent GPU job). - Qwen3-Omni t2s, 192 requests per arm, Thinker prefill unchunked vs cut at 16 tokens so every prompt takes several chunks. Per-answer WER: each answer's audio transcribed with whisper-large-v3-turbo and scored against that same answer's text, over the spoken prefix -- the thinker writes about twice what the talker speaks, so scoring the whole text reads 0.54 on correct audio. | arm | n | WER | p50 | p90 | >0.5 | spoken/text | |---|---|---|---|---|---|---| | unchunked | 192 | 0.110 | 0.078 | 0.298 | 1 | 0.48 | | 16-token | 192 | 0.102 | 0.061 | 0.266 | 2 | 0.47 | Fisher exact p = 1.000 on the garbled counts; the unchunked arm's 1/192 is the model's own floor (#365 measured ~1/64 with chunking off). 192/192 wavs both arms, no failed requests. A budget that leaves prompts in one chunk is a false pass: the join takes a single-piece fast path and never runs. Confirm chunking happened first. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Key the chunk Loop set on the graph, and restore the prefix chain Two defects from review. The worker took ``_chunk_loop_walks`` from the model's chunking POLICY, but ``apply_chunking`` builds the Loop from that policy at graph-build time while the policy-derived set excludes nodes this deployment chunks nowhere. Under TP/SP the Loop is in the walk's graph and the worker's set omits it, so nothing ends it and it spins to ``max_iters`` holding the request's KV -- a hang on every TP deployment of a chunkable model, which the shipped-config check cannot see because it only verifies the Loop exists. Read off the graphs the worker was given instead (``_chunk_loops_in``). The policy says where a Loop may be built; the graph says where one is, and cannot disagree with itself whatever rule excluded a node. ``extend_prefix_chains`` had been dropped from postprocess entirely, so nothing keyed what a request generated and cross-request reuse silently stopped learning. Restored, with the rows mid-prefill skipped: a non-final chunk's sampled token is not a real sample, so keying the chain on it keys a token the request never generated. Tested on 2 x H100, CPU only: `ruff check .` clean; 386 passed over the chunk Loop, scheduler, chunking, admit-failure, prefix-stage and runtime-parity suites. New cases: the Loop ended for a walk whose deployment chunks nothing, the set read off wrapped and unwrapped graphs, a decode Loop not mistaken for a chunk Loop, and a guard that postprocess hands the mid-prefill rows to the prefix chain as skips -- that call was dropped once and nothing noticed. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Speculate into a chunk Loop, cutting the next chunk from a lookahead A chunk-only step had nothing queued behind it: the runtime never offered a chunkable (node, walk) as a speculation target, so the decode rows waiting on it paid the whole step. The filter was there because a speculated chunk k+1 needs where chunk k ENDS, and chunk k has not landed -- ``_assign_chunks`` would fall back to the end of k-1 and re-run k's tokens into the KV. ``ChunkProgress`` now carries that end. ``done`` is unchanged: the ground truth, advanced in postprocess when a chunk lands. Beside it a ``ChunkCursor`` records where the chunk just scheduled begins and ends, kept apart for a speculative schedule so a cleared speculation leaves no trace. A chunk is cut from the lookahead when there is one and from ``done`` otherwise, and landing promotes the speculative cursor -- the step it was cut for is the one now in flight, so it is no longer a guess. The lookahead is only ever ahead of ``done``, so every abandon path just discards it. Three reach it: the worker's push-back (admit refusal, OOM), the scheduler's own hold backoff, and a speculation cleared before it runs. The hold is the one a test caught -- it never touches the worker, so a held row retried past tokens it had not run. ``MidPrefillRow`` needed no change: its construction was never gated on walk mixing, only on the spec target being the chunked node, which the filter made impossible. TP is unaffected -- chunking is off for TP/SP nodes until #367 -- so the lockstep question does not arise here. Tested on 2 x H100, CPU only: `ruff check .` clean; 454 passed over the runtime-parity, scheduler, chunking, chunk-Loop, speculation, admit-failure, resident-delta and worker-drain suites. New cases: a speculative chunk cut after the one in flight, a cleared speculation leaving the next chunk where it was, and a landed chunk promoting the speculative cursor. NOT YET MEASURED. The point of this is closing the bubble, and that claim needs a `benchmark/worker_phases` breakdown showing it; the machine this was written on could not give a reliable allocation. Until then, treat the win as unverified -- correctness is tested, the performance argument is not. Note for bisecting: this removes a Rust constructor argument, so it and the commit before it need different `maturin develop --release` builds. A stale extension fails at request time with a TypeError, not at import. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Mark a chunk Loop's last iteration so speculation skips it A prompt short enough to run in one chunk still goes through a one-iteration chunk Loop, whose loop-back makes the prefill node a speculation target after every prefill step. Every row then dropped out of that speculated step (its loop-back never arrives), so decode lost its overlapped step: on Qwen3.5-4B short prompts the branch was -1.5% tok/s at c16 and -3.8% at c48 against the main it merged, ITL p95 11.1 -> 14.2 ms, with ~500 "Speculation: dropped rids" warnings per run. The runtime now takes last-iteration marks: `mark_last_loop_iters(rids, walks, loops)` says a row's scheduled iteration is its last, and `_can_continue_loop` honours it. Only speculation reads a mark; the loop still ends by its stop, so a stale mark can cost at most one unspeculated step. A mark goes when the loop's stop is applied, on `clear_last_loop_iters`, and when the request is removed. `speculate_node` takes every row of the step, grouped by walk, rather than one sample rid per walk: a loop-back target is kept only if some row can continue, so a walk whose first row is on its last iteration still speculates for the rest. The scheduler marks rows running their final chunk (or no range) when it builds a step, and `drop_chunk_lookahead` clears them, since a refused step's retry is marked again. Python and Rust runtimes both carry it. Measured on 2 x H200, Qwen3.5-4B, short prompts: c16 3576 tok/s vs 3576 on main, c48 8027 vs 8043, and no dropped-rids warnings. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Don't publish a stopped loop for a cache-served row A chunkable walk the prefix cache serves whole runs no iteration, but its chunk Loop is still stopped. Publishing after that stop looked up the row's forward-pass info, which a row that did not run does not have: a KeyError that failed every request in the step. Seen on Qwen3-Omni with prefix caching. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Default max_chunks to effectively no cap `ChunkingPolicy.max_chunks` becomes the chunk Loop's `max_iters`, and it defaulted to 1. BAGEL and Orpheus never set it, so any prompt cut into chunks ended its prefill walk after the first chunk; the conductor then started decode with the rest of the prompt still to run, and decode steps carried it ("plan label 'main' carries 929 tokens but its captured bucket holds 16"). The count is only a runaway guard: the engine ends the Loop on the final chunk, and the number of chunks is not known until the budget cuts each step. So the default is now 2**31 - 1, which fits the runtimes' u32 `max_iters`. Tested on 2 x H200, prefix cache off: Orpheus with a 64-token budget went from 93/287/188 failed requests (c16/c48/Poisson) to 0; BAGEL short, mixed, long and image workloads, closed loop and Poisson, ran with 0 failed requests and no server tracebacks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Measure a walk the cache holds whole at 0 tokens A chunkable keyed row is measured without its cached prefix, but only a partial match was taken off: a walk the cache holds whole was measured at its full length. The budget then cut it into chunks from 0, and a chunk cut from 0 writes over the lease the walks before it were served from, which fails the request whenever those walks end inside a page ("the walks the cache served end inside a page (43 slots), so their KV cannot be kept"). Qwen3-Omni image requests hit it whenever an image repeated: the text walk is served (43 slots), then the vision walk, held whole, was chunked. Measured at 0 it takes no budget and no range, and prepare serves it from the cache like any unchunked row. Tested on 1 x H200, Qwen3-Omni colocated, prefix cache on, images repeating, c8: 8 of 48 image requests failed before, 0 after, with no server tracebacks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Refuse a Thinker input split_inputs cannot cut The merge took main's `split_inputs`, which cuts the inputs it knows and silently drops any other, where this branch's refused one. A dropped key would leave a chunk missing an input it was prepared with, so refuse it again. `test_qwen3omni_vision_chunks.py::test_an_input_it_cannot_cut_is_refused` failed since the merge and passes again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Drop the explicit max_chunks from Qwen3.5 and Qwen3-Omni The default no longer caps a prompt's chunks, so the per-model value (each model's max_position_embeddings) is redundant. The docs described max_chunks as something to size to the longest prompt; they now say to leave it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Split a step's prompt budget in start order The token budget was a max-min fair share across every waiting prompt, so with several long prompts each advanced a little every step: all of them finished late, and their KV pages were held the whole time. The default builder now fills the budget with the prompts that fit whole (shortest first), then gives what is left to prompts in the order their first chunk was cut, cutting at most one. A prompt mid-prefill keeps its start position, so it finishes before a later one begins. The equal share moves to FairShareBatchBuilder, selected with `batch_builder: fair_share`. On its own this trades ITL for TTFT (H200, Qwen3.5-4B, against the equal share): long prompts c32 TTFT p50 1.3-1.4 s -> 0.37 s but ITL p95 43 -> 80 ms; very long c16 TTFT p50 3.3-3.8 -> 3.0 s, ITL p95 86-100 -> 130 ms. A prompt left out of a step now waits as a backlog, and the step after a prompt step serves it alongside the decode rows. The next commit gives that step to the decode rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Let ready work go first after a chain that cannot continue The fairness peek (`has_ready_excluding`) yields a speculation chain to other ready work only once the chain has run a step past its first. A chain that could not continue at all skipped it, and the fallback `get_next_batch` was unrestricted, so a (node, walk) with a backlog won. A prompt-only chunk step is such a chain (rows mid-prefill are not speculated), so under the start-order builder the next step was another prompt's chunk with the decode rows riding along through the prefill kernels. The peek now applies there too: with other ready work, the next step is that work, giving two decode steps per prompt step. With `MSTAR_SPEC_PEEK_FOR_FAIRNESS=0`, or nothing else ready, nothing changes. H200, two rounds, start order with this against the previous equal share: Qwen3.5-4B very long c16 ITL p95 86-100 -> 49 ms; long c32 TTFT p50 and ITL p95 flat (1.36 s, 42 ms). Qwen3.5-4B mixed c48 ITL p99 flat, +1.2% tok/s; BAGEL mixed c48 ITL p95 54-69 -> 44-45 ms, +2.6% tok/s; Qwen3-Omni mixed c24 ITL p95 80 -> 64 ms, +1.5% tok/s. TTFT on long prompts is in the next commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Park the prompts a step's budget leaves out, in line A fresh row the token budget left out went back to its ready queue, and a backlogged one was parked behind the rows the row cap pushed out. The Rust runtime scans ready rows in handle order and the backlog is served first, so either way a prompt rejoined behind ones that arrived after it. Under start order, which leaves out most prompts every step, the first prompts of a burst started last. Every row the budget leaves out now parks ahead of the cap's leftovers, so prompts start in arrival order. H200, two rounds, against the previous commit: BAGEL long c32 TTFT p95 44.2-44.4 -> 7.2 s, Qwen3.5-4B very long c16 28.9 -> 6.3-6.5 s, with mean TTFT and E2E unchanged (BAGEL 5.8 -> 5.9 s, 7.5 -> 7.6 s; Qwen 5.2 -> 5.1 s, 7.4 s). TTFT p50 rises to match (0.7 -> 6.2 s, 1.2-1.5 -> 5.4 s): the low median came from the prompts left waiting. Against the equal share before start order, E2E mean drops 4-8% (BAGEL 7.9 -> 7.6 s, Qwen 7.9-8.4 -> 7.4 s) and TTFT p50 rises (2.8-4.2 -> 6.2 s, 3.3-4.7 -> 5.4 s). Qwen3.5-4B long c32 and mixed c48 are flat; BAGEL mixed c48 is -1.3% tok/s, TTFT p95 2.3 -> 2.9 s, p99 4.6-6.6 -> 4.2 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [GDN] Size a captured mixed step's decode half to its bucket `GDNMixedWrapper` runs its decode half on every row it is sized for, and a captured wrapper was sized to the largest capture (`_cg_max_bs`), the floor the other wrappers take for their buffers. So every mixed bucket ran the GDN decode kernel on 128 rows: on Qwen3.5-4B at c48, ~2.6 ms of a mixed step was decode work on rows addressing the null slot. The split wrapper now takes the bucket's rows; the plain decode and prefill wrappers keep the floor, since their kernels run on the step's rows whatever their buffers hold. The dropped rows address only the null slot: outputs and every real slot's state are bit-identical to the 128-row wrapper (`test_the_decode_halfs_width_changes_no_real_rows_result`). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Engine] Order a split walk's steps with its prompt rows last A resource that runs a combined walk on split kernels (CombinedSplit) reads the step's one-token rows on a decode kernel and its prompt rows on a prefill kernel. The engine records each split's prefill walks at load, refusing two resources that disagree, and the worker's one batch builder puts those walks' rows last, keeping each side's order. The kernels can then read the decode rows as the step's first tokens in place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [GDN, Attention] Read a mixed step's decode rows in place GDNMixedWrapper and FlashInferMixedWrapper gathered every token into decode-then-prompt order and scattered the outputs back, per layer (index_select / index_copy_ over the whole step). With the engine putting the decode rows first, the decode half reads the leading run of one-token rows as the step's first tokens and a bs-row torch.where merges its outputs in place. Under capture the prefill half takes every token, the run as one sink row ahead of the prompt rows; an eager step's prefill half takes the tokens after the run, so it writes prefix-cache checkpoints only for the prompt rows its checkpoint buffer is sized to. Row order only bounds capacity: a decode row behind a prompt row runs on the prefill half, which is still correct while the step holds no more prompt rows than the bucket was captured for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Attention] Keep a mixed step's padding rows out of the decode split budget FlashInfer's tensor-core decode splits long KV rows across CTAs within a budget of 2 CTAs per SM over the KV heads (66 rows on an H200 with 4 KV heads), and every row it plans costs at least one, padding rows included. A 64-row mixed bucket holding ~46 decode rows left no room to split, so its long-KV rows ran unsplit: 328 us a layer against 135 us for the same rows in a 48-row decode bucket. The split's decode half now runs the same FA2 paged kernel planned through the prefill wrapper, where a row with no query costs nothing: the run's rows take one token each and one row carries the rest. Measured on 46 rows with 15% at 4k-16k KV: 159 us at 64 rows. The CUDA-core option (decode_tensor_cores: false) keeps the decode wrapper. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] docs: cut the model author's chunking section down Keep what a model author does: declare policies, wrap the walks with apply_chunking, the policy fields, join_dim, the two length/cut hooks and the final-chunk sampling rules. Drop engine internals (how resources read reserve_whole_prefill, why join_dim joins on arrival) and the stale note that the engine holds and joins non-final outputs, which the Loop does now. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] reserve_whole_prefill as a per-row setting StepContext.reserve_whole_prefill was one bool for the step, the AND over its chunked rows, so a step mixing a walk that opts out with one that opts in turned the reservation off for both. It is now a PerReqSetting (SET, UNSET, or PARTIAL with the rows it holds for), and the KV manager reserves ahead per row. No model opts out today, so no shipped step changes. Also cuts StepContext's comment on it to three lines. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A serving config can set min_chunk_tokens The chunk floor was a module constant. It is now the serving config's min_chunk_tokens (default 128), set on the engine's ChunkingManagement with the token budgets and validated there. Both readers take it from the engine: the micro-scheduler's cut and check_budget's overrun tolerance. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Comment fixes from review - ChunkingPolicy: the multi-line field comments move into the docstring. - PythonGraphRuntime._combined_walk_of gets its own comment back, which a note about the removed disable_spec_node_walks had replaced; the same leftover goes from speculate_node's docstring and from the Rust runtime's struct and constructor, where rustdoc attached it to the next field. - MicroScheduler: drop a comment whose field went with the accumulator, and finish ChunkProgress.cursor_spec's. - Worker._advance_chunks keeps only what it does now. - Engine: a blank line before _prepare_chunk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Name a chunk Loop and its tick from (node, walk) alone _wrap_node_in_loop built the tick name from the loop name it was handed, a second spelling of chunk_tick_name that matched only while every caller passed chunk_loop_name(node, walk). If they drifted, the conductor would seed an input the node does not declare and the first chunk would never become ready. It now takes neither name and calls chunk_loop_name and chunk_tick_name itself. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Graph] Rename last-iteration marks to do-not-speculate marks mark_last_loop_iters / clear_last_loop_iters read as if they changed how a loop is scheduled, but only speculation reads them: the loop still ends by its stop. They become mark_do_not_speculate_loops / clear_do_not_speculate_loops, with the runtime state, the scheduler's _mark_final_chunks (now _do_not_speculate_final_chunks), tests and docstrings following. Python and Rust runtimes both, which needs a rebuilt extension. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Load the submodules in EngineManager.build again The worker loaded its submodules before building the graph runtime so the runtime could be told which (node, walk) pairs never to speculate into. That argument went when speculation into a chunk Loop landed, and nothing else needs the submodules before the runtime exists, so the split and its stale comment go: build loads them, as on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Move the chunk Loop bookkeeping into ChunkLoops Worker had grown four chunking handlers. They move to ChunkLoops in mstar/worker/chunk_loops.py: the (node, walk) pairs read off the graphs at load (Worker.chunk_loops.walks, and walks_by_node for the scheduler), ending a Loop on its final chunk (end_finished) and recording what each chunk ran (advance). The worker calls them at the same points of postprocess as before, on the same thread; no behaviour change. Tests drive a real ChunkLoops over a real ScheduledBatch instead of a bare Worker. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] A decode step beside a chunk Loop pays one lookup, not one per row ChunkLoops.end_finished walked every row of every step on a node with a chunk Loop, decode steps included, only to find decode is not a chunk walk. It now returns early for a single-walk step whose walk has no chunk Loop, as the scheduler's _do_not_speculate_final_chunks already does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Speculate a prompt's next chunk on its own A step whose rows were all mid-prefill never speculated: the worker took the held rows out of the candidates and built the merge that carries their next chunk only behind a continuing row, so with walk mixing off (the default) a prompt's chunk k+1 waited for chunk k to land. A speculation can now be made of those rows alone, its capture group taken from one of them, and the merge keeps to their walk. So the step after a prompt step still goes to other ready work, the fairness peek runs for a prompt chain from its first step: two prompt steps in a row only when nothing else is ready. `speculate_mid_prefill` (serving config, per node or per walk, default on) turns it off; off, the next chunk waits for the current one to land, in a mixed step too. H200, Qwen3.5-4B, two rounds, PR head / on / off: a lone very long prompt's TTFT p50 455-456 / 437-441 / 451-455 ms (p95 582-585 / 564-566 / 580-582); a lone long prompt 119-121 / 116 / 117-118 ms. Long c32, very long c16 and mixed c48 are flat in tok/s, TTFT and ITL. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Skip the chunk bookkeeping when nothing is chunked _assign_chunks returns at once when no row is budgeted, in progress or holding a cached prefix (4.9 -> 0.1 us for a 64-row decode batch), and drop_chunk_lookahead walks chunk_progress once against a set of the rids rather than once per rid (16 -> 0.9 us for 64 rows, 8 prompts in progress). A StreamBuffer entry that is a lone tensor, every item of an unjoined stream, skips the joined path (0.88 -> 0.53 us per item; main is 0.37). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: main's request-added and speculation-readiness fakes carry the chunking fields Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Key a step's prefix chains before its withheld outputs are dropped Main indexes the stop check's host copy before dropping withheld and empty-span outputs; the merge had moved the indexing after the drop, which on a step without a completion event (where the host copy is the outputs themselves) left a withheld token unkeyed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A cut row waits a step rather than run the step over its budget The row a step's budget cuts took at least the chunk floor even when less than that was left, so a step that finished one prompt could run up to the floor over its budget: 1990 + 128 tokens under 2048, say. Qwen3.5's prefill captures go 2048 -> 4096, so that step padded to the next bucket up, and the budget check allowed exactly the floor's slack, so it was never logged. A row now waits for the next step when less than the floor is left, under both start order and the equal share; the floor still sets the smallest chunk. The budget check warns on any step over its budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A chunk Loop counts the tokens its chunks ran Chunk progress lived in the scheduler (ChunkProgress and its lookahead cursors), because the Loop only counted iterations up to a runaway guard and could not tell when a prompt was done. So the worker ended the Loop by hand on the final chunk (ChunkLoops.end_finished), both runtimes carried do-not-speculate marks for the final chunk, and every abandon path had to drop the lookahead or the retry would skip tokens. The Loop now counts tokens. Each row's completion moves its enclosing loop by a stride (RouteInput.loop_strides): a chunk's ChunkRange.stride, the END stride for a chunk-walk row that ran whole, one otherwise. A prompt's first chunk caps its Loop at the prompt's length (GraphRuntime.set_loop_caps, which assigns rather than lowers, so the walk's next prompt sets its own). The Loop then ends itself on the last chunk, and speculation eligibility takes the in-flight step's strides (speculate_node, prep_spec_rids, prep_spec_rows), so nothing speculates past it. The scheduler resumes a popped or backlogged row mid-prefill from the Loop's count (GraphRuntime.loop_progress), and a speculative one from the end of the chunk in flight. Nothing is cut ahead of what landed, so nothing has to be undone when a step does not run. Gone: ChunkProgress, ChunkCursor, advance_chunk, drop_chunk_lookahead, ChunkLoops.end_finished/advance, mark/clear_do_not_speculate_loops and ChunkingPolicy.max_chunks. Kept: the loop-back tick, which is what makes the next chunk a speculation target, and the start order the FIFO builder serves prompts in, now a (rid, node) -> sequence dict. Both runtimes change alike and the parity tests cover the new calls. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] The sampler keeps rows mid-prefill out of its commit A non-final chunk's sampled token is not a real sample, so its RNG offset and seen-token mask must not commit. Every chunking model wrapped its SamplerStep in keep_final_chunk_samples for that, and a model that forgot it erred nowhere: its outputs just changed once its prompts were chunked. The sampler now does it itself, the way the KV and recurrent resources apply a step's forks only to fork_rids: it drops the rows of the step context's chunk_tokens_after (the rows a chunk runs after this step) from what its draw commits, eagerly and under capture. keep_final_chunk_samples goes. The models' postprocess guards that dropped a non-final chunk's new_token go too: the chunk Loop carries only the final chunk's outputs (the body's copies route to EMPTY_DESTINATION with no flags), and the worker takes no stop or prefix key from a row mid-prefill. Their tests become the sampler's own and the runtimes' (test_only_the_final_chunks_token_leaves_the_loop). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Model wraps chunkable nodes wherever the walk graphs are read Each chunking model had to pass its walks through apply_chunking at the end of get_graph_walk_graphs, and a load-time check caught the ones that did not. Model.__init_subclass__ now wraps every subclass's own get_graph_walk_graphs so its result goes through apply_chunking, which makes the ~20 call sites that read the walks see the chunk Loops with no accessor to route them through. apply_chunking leaves a node already in its chunk Loop alone, so a subclass building on its parent's walks is wrapped once. The four models drop the call, and the check that a chunkable walk is wrapped goes, since it can no longer fail. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] The oldest prompt mid-prefill always gets the chunk floor Under start order, prompts that can finish this step go ahead of the ones already mid-prefill, and nothing bounded how long those could wait. With 16 decode rows at a 2048 budget and four new 500-token prompts each step, a 20k-token prompt that had started got 0 tokens for as long as the stream lasted. The oldest prompt mid-prefill now keeps min_chunk_tokens out of the finishing pass, so it advances every step; the step stays within its budget (one new prompt waits a step instead). Unchanged when nothing is mid-prefill. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: track the chunk Loop's stride end and main's dtype fake EngineManager.build reads the model's chunking policies on this branch, so the fake that stands in for what build reads (main's #810) has to answer it too. A cache-served row's chunk Loop now ends by its route stride rather than a stop, so the dropped-row test checks the stride. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Seed a chunk Loop on the ranks hosting its node The conductor added the chunk ticks after splitting a walk's inputs, so they reached only the ranks that already got an input. Qwen3-Omni's prefill_audio sends every edge to the audio encoder, so with the encoders and the Thinker on different ranks the Thinker's Loop was never seeded. The ticks now join the inputs before the split, which broadcasts a signal-only edge to the node's own ranks, and only where a walk is entered: a client stream chunk carries none. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Validators read a chunked node's edges off its Loop After apply_chunking a chunked node's body routes only to EMPTY_DESTINATION and its tick, so streamed_partitions lost Qwen3-Omni's Thinker->Talker edges and the conductor's prefix-cache guard for the Talker never fired. Both validators now read the enclosing chunk Loop's outputs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A hit that lands after a row was measured is declined A row measured with nothing cached can find a hit by its first chunk, when an earlier request with the same prefix commits while it waits. The first chunk treated that as the policy declining and applied no prefix, which ends the row's chain. It now applies a zero-token prefix, keeping the chain, and the recurrent pool gives back a snapshot a zero-token prefix does not restore. A walk whose p…

* Admission by row: leave out a row whose KV does not fit - NodeSubmodule.get_input_sequence_len(walk, fwd_info, inputs) reports a row's token count from input shapes, before prepare_inputs; None (the default) keeps today's behaviour. InputSeqLenInfo can give a resource a different segment length than the row's. - Resource.has_room(rid, node, walk, segment_len) answers whether a step adding that many tokens could be admitted now; KVManager counts the pages each needed label grows by against free pages plus the pages the prefix index alone holds. - StepRunner.admit_retrieve / Engine.check_ready take the optional InputSeqLenInfo; no room is reported as not ready, not as a failure. - MicroScheduler measures popped rows of nodes that report lengths: a fresh row without room goes back to its ready queue, a backlogged one stays parked, and the rest of the step runs. The check is per row: a step of rows that each fit can still fail admission, which the existing hold path handles. - Orpheus (prefill, decode) and BAGEL (prefill_text, decode) report lengths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: report Thinker row lengths for admission by row Every Thinker walk's KV span is known from its inputs before prepare_inputs: text from its token ids, audio and vision from the encoder's embeds plus the two sentinels around them, decode one token. The embeds arrive on the Thinker's input edges, so admission can leave out a row whose KV does not fit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: report LLM row lengths for admission by row Text prefill rows are their token count and decode rows one token, known before prepare_inputs, so admission can leave out a row whose KV does not fit. A vision row's span depends on its image grids; it is admitted whole. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Backlog: measure a parked row once and check it once per pass A backlogged row was checked for readiness, then measured and checked again with its length, on every pass it waited -- every speculation step under a split decode set. Its inputs do not change while it waits, so the batch keeps each row's length, and the one readiness check takes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Update docs * Drop the has_room admission gate from scheduling The gate refused a row whose KV did not fit, intending to let it sit a step out rather than refuse the whole step. Two reports show it does not hold up: - Decode deadlocks. Every decode row needs a page for its next token, so a pool held by decode rows refuses all of them; no step is built, nothing runs, and nothing is ever evicted or offloaded to break the cycle. The gate runs before the admit-error path, which is the only thing that triggers an offload. - Every row of a pass sees the same free count, with no running sum, which is what let eight chunked prompts in at once on a 48-page pool. has_room does not exist on main; it arrives with this PR. So dropping it returns to main's behaviour -- the admit-error path refuses the step, OOM backoff holds rows, and the retry is smaller -- and cannot regress against it. Real per-row admission control belongs in #400, which reserves a request's whole uncached prompt up front; this gate was a stop-gap for it. Keeps the declaration itself: get_input_sequence_len, InputSeqLenInfo and engine.input_sequence_len stay, since chunked prefill sizes its chunks off them. Only the readiness gate goes. Co-Authored-By: Claude <noreply@anthropic.com> * Scheduling: try the next node when a step comes out empty Only the first round-robin pick was tried, so a node whose step came out empty left the worker idle for the pass -- and because such a node is never marked scheduled, round robin picked it again on every later pass too. Under a colocated model that starves the node which would unblock it: on orpheus_colocated, snac_decoder waits behind the LLM, and the LLM's pages only come back once snac_decoder finishes the request. _select_node_rr becomes _node_rr_order, returning the whole order rather than its minimum. A key already tried off its backlog above is skipped; a key that declines takes no rows, so moving on costs nothing. Co-Authored-By: Claude <noreply@anthropic.com> * docs: tighten get_input_sequence_len, add per_request_info to declare_step The length a model reports is the whole input; the engine is what subtracts the cached prefix, so say so rather than leaving the reader to guess which side cuts it. None covers a walk that declares no length, not only absent inputs. Drops the claim that the scheduler uses it to leave out a row whose KV does not fit -- that gate is gone. declare_step grew per_request_info when the non-final-chunk sampling fix landed (#353); the signature here never caught up. Co-Authored-By: Claude <noreply@anthropic.com> * BAGEL: an absent text input is undeclared, not a KeyError get_input_sequence_len's contract is to return None when the walk declares no length, and every other model that implements it -- orpheus, qwen3_5, qwen3_omni -- guards the key before reading it. BAGEL was the one that indexed blind, so a prefill_text row that reached here without its text input took down the measuring pass instead of simply going unmeasured. Co-Authored-By: Claude <noreply@anthropic.com> * Pre-plan a speculative packed step from its rows' spans A packed step (a mixed or prompt step) had no token count before prepare_inputs, so it could neither lease a capture nor pre-plan, and planned inline on the gpu thread behind the step in flight. Its rows' spans are now measured when the speculation is armed: a row in hand from its inputs (get_input_sequence_len), a row continuing from the step in flight from the submodule's prediction (speculative_input_seq_len, None by default). The step leases its packed bucket on those spans and declares over host-only rows of them (declare_inputs_from_spans, None by default), on the plan thread like a batched step. The resources promote a stage by its rows and slot, not its spans, so after prepare_inputs the engine compares the prepared spans with the planned ones; a mismatch (a cached prefix cut a prompt, a row left) drops the stage and the packed lease, and the step leases and plans inline. MSTAR_PRE_PLAN_PACKED=0 turns this off. Qwen3.5 opts in for its mixed and text prefill walks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus and BAGEL: pre-plan a speculative packed step from spans Both declare a text step from each row's span and walk; BAGEL also from its guidance flag, so declare_inputs_from_spans now takes each row's forward info. A step declared from spans whose capture key is not its lease's is not pre-planned, so a hook that misreports guidance cannot stage a plan for the wrong capture. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Pre-plan: estimate a continuing row's span from its last one speculative_input_seq_len takes the row's span in the step in flight (-1 if it ran on another node), and by default repeats it, the rule chunked prefill's merge budget uses; a walk that fixes its span says so (decode: one token, which a prompt's last chunk continuing as decode would otherwise get wrong). ExecutingBatch.seq_len_of reads a prepared step's row spans, for the worker to hand over. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: run a long input across several steps - NodeSubmodule.supports_chunked_prefill(walk) / max_batch_tokens(walk); NodeInputs.chunk_start / chunk_total / is_final_chunk. A chunked walk must report get_input_sequence_len; chunking under TP/SP is refused at load for now. - MicroScheduler keeps chunk_progress per (rid, node). FIFOBatchBuilder divides max_batch_tokens max-min fairly (fill_token_budget): rows under an equal share go in whole, chunkable rows split what is left, others sit out; a step is never empty. Chunk ranges ride on ScheduledBatch, through the speculation merge too; a chunkable row's room check is for one step. - Engine prepares a chunked row once, cuts each chunk with split_inputs, and keeps the whole input until the final chunk lands, so a refused admit retries without re-reading advanced counters. A non-final chunk's row is left out of the sampler's kept rows. Chunked rows skip the prefix cache for now. - Worker leaves a non-final chunk's row out of input cleanup, stop checks and routing: its outputs wait in ChunkOutputAccumulator (its own file) and its node goes back to the ready queue once the chunk lands, so chunk k+1 never runs before chunk k; the final chunk routes the concatenation. A mid-prefill row never continues a speculation chain. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus: chunk long prefills prefill can be split; max_batch_tokens is 512 for prefill and mixed steps. postprocess drops a non-final chunk's sampled token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: chunk long Thinker text and audio prefills prefill_text and prefill_audio can be split, with a 1024-token budget for them and thinker_mixed steps. split_inputs also cuts masks_for_talker on its sequence dim; row lengths come from the Thinker's get_input_sequence_len (cp/2); only the final chunk of the last prefill keeps its token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Refine estimated sequence length of continuing requests; add output accumulation policy * Qwen3-Omni: chunk long Thinker vision prefills prefill_vision can be split like text and audio, with its own 2048-token budget (vision steps run alone, at bs=1). split_inputs also cuts the per-layer deepstack features and the talker masks with the chunk's tokens, and refuses tensor inputs it does not know how to cut. A vision span advances the MRoPE counter by its 3D-grid extent, not its token count, and its chunks' positions are cut from the whole span's, so only the final chunk declares the advance. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a final-chunk-only output mode ChunkedPrefillOutputMode.FINAL keeps an edge's final-chunk tensors alone; earlier chunks' are dropped at hold time rather than held and discarded. It does not replace a model dropping a non-final chunk's sampled token in postprocess, since check_stop reads outputs before they are held. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: chunk long text prefills prefill_text can be split, with a 1024-token budget for it and mixed_text (placeholder). A chunk keeps its row's requires_cfg flag. A guided prompt's pre-forward main -> cfg_text fork (cfg_text keeps the pre-text context) happens on its first chunk only, via KVStep.fork_rids, and only the final chunk's sampled token is kept. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: start a chunked row past its cached prefix The scheduler's first measurement of a keyed, chunkable row resolves its cached prefix (the KV lease, under the manager's lock, as admit_retrieve does) and reports only the rest; the row's chunks then start past it, and its first chunk applies the prefix with the whole prepared input. A model can refuse a row before it is prepared (reuses_cached_prefix; BAGEL refuses a guided prompt, which writes two labels). Also gives test_kv_prefix_stage's stub batch the wrapped request info. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Worker: read the speculated step's row lengths off its executing batch The speculation merge read them off the scheduled batch, which has no node_batch, so every speculative step raised and failed its requests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: a prefill step planned as a decode keeps every token An eager step with one token per row gets FlashInfer's decode plan, which sets no qo_indptr; select_last_hidden then failed, and BAGEL's copy of it asserted. Chunking makes such steps common (a one-token final chunk). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a serving config can set token budgets max_batch_tokens: per node, one budget for every walk the model budgets, or per walk; null turns chunking off. Refuses unknown nodes and walks, and walks the model sets no budget for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: correct the FINAL mode's note on sampled tokens Held rows leave the batch before the stop check, and the engine takes non-final chunks out of the sampler's kept rows, so neither reads a non-final chunk's token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * ruff * Qwen3.5: chunk long text prefills prefill_text can be split, with a 1024-token budget for it and mixed steps, under the largest prefill capture with a mixed step's decode rows. Chunks resume the GDN state from the request's slot and take positions from the counter, as successive prefill steps already do; only the final chunk of the last prefill keeps its token. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: one chunking policy per walk NodeSubmodule.get_chunking_policy(walk) returns a ChunkingPolicy (chunkable, max_batch_tokens, output_policies, reuses_cached_prefix) in place of four separate hooks, so a submodule declares a walk's chunking in one place. The engine asks once per (node, walk) and keeps its accessors. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a row's progress is a ChunkProgress(done, total) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: final-chunk sampling is the submodule's call keep_final_chunk_samples(sampler_step, request_ids, inputs) drops the rows mid-prefill from a SamplerStep's kept rows. An autoregressive submodule that chunks calls it in declare_step; the engine no longer rewrites every step, since not every chunked node samples. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Recurrent state: scope a step's forks to some of its rows RecurrentStep.fork_rids, as KVStep's: the rows the step's forks apply to, None for every row. Nothing forks a recurrent pool yet; without it, a chunked prompt's later chunks would copy the source over the fork target again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: the runtime never speculates into a chunkable walk A row entering a chunkable (node, walk) is measured and cut when it is scheduled. A continuing row's inputs are the in-flight step's outputs, so scheduled ahead its first prefill would run whole, past the token budget. The worker loads its submodules first (EngineManager.load_submodules), derives the chunkable (node, walk) pairs from their chunking policies, and hands them to the graph runtime's constructor as disable_spec_node_walks; speculate_node, Python and Rust, never returns one as a target. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: the chunking engine fake follows _SizedEngine's rename Dropping the has_room gate one branch down renamed _SizedEngine's `too_big` parameter to `unready`, and the stack rebase replayed this branch's commits on top without updating the call site -- the two lines never overlap textually, so git merged it clean and every _ChunkingEngine test died in construction with a TypeError rather than on anything it meant to assert. Co-Authored-By: Claude <noreply@anthropic.com> * Chunked prefill over combined walks: tests follow the own-walk merge and the uniform path A speculation merge takes only its own walk, so the budget test's continuing row is a prompt. The uniform-inputs fakes carry chunk ranges and row lengths, and a uniform step records no row lengths: one token a row is what an unrecorded row is taken for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: main's new engine fakes carry the chunking fields test_dropped_row_stream_stop, test_kokoro_chunk_loop and the Cosmos3 step guard build an Engine by hand; chunked prefill reads its per-walk policy cache and each batch's chunk ranges, so they get both. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the packed pre-plan and the row lengths admission never read Pre-planning a packed speculative step from row spans measured ahead measured slower than planning it inline: at c32 on Qwen3.5-0.8B it saved ~196 us/iter of inline planning and added ~278 us/iter of plan-thread wait, with no gain at c64 or c128. This removes it (MSTAR_PRE_PLAN_PACKED, declare_inputs_from_spans, measure_ahead_spans and the packed-lease bookkeeping), and the per-row InputSeqLenInfo passed through check_ready to admit_retrieve, which ignored it. Kept for chunking: speculative_input_seq_len (it sizes the token budget left beside continuing rows), ExecutingBatch.seq_len_of and the scheduler's measured row lengths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Speculation merge: estimate the continuing rows' spans only when sizing a budget The merge estimated every continuing row's span (a submodule call per row, plus a dict over the in-flight batch) before scheduling, though the token budget only needs it once fresh rows are merged in -- which a steady decode merge rarely has. get_next_batch now takes the estimate as a callable and asks for it only then. About 8 us/step at c16 and 50 us/step at c128 off the speculate phase on Qwen3.5-0.8B. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: skip the per-row chunk bookkeeping on steps without chunks A split copies row_lens only when it holds some; the speculation candidates and input cleanup check incomplete_node_rids only when a row is mid-prefill; the batch builder collects its backlogged rids only when the token budget dropped a row; and a speculation copies chunk_ranges and incomplete_node_rids only when they are set. All per-row work on every decode step that never chunks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: admit a chunked prompt's whole walk on its first chunk A chunk admitted only its own tokens, so on a tight pool several long prompts got in at once, each held the pages of the chunks it had written, and all waited on pages the others held (the 48-page wedge in review). StepContext.chunk_tokens_after carries the tokens a chunked row runs after this step, and KV reserves them with the step: the first chunk admits the whole walk, as an unchunked prompt does on main, and later chunks take no more. A prompt that does not fit is refused before it holds anything; admission control proper is #400's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: fail a row whose prepared input is not the length it was cut by The scheduler cuts a row by the length get_input_sequence_len reported, and nothing checked that prepare_inputs agreed (PR #365 review): a mismatch would drop tokens or run a chunk past the end. chunk_ranges now carries the measured total with each [start, end), and the first chunk's prepare compares it with the prepared input; a mismatch raises, which fails that row and lets the rest of the step run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Positions: fork only the rows KV forked KVStep.fork_rids scopes a step's forks to some of its rows (a chunked guided prompt forks on its first chunk and its last), and KV honoured it, but positions applied every fork to every row: a guided row on a later chunk, sharing a step with another prompt's first chunk, had its cfg_text counter re-forked to main's while its cfg_text KV stayed put (PR #365 review). KV's plan output now carries fork_rids, and positions fork the same rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a step of held chunks still routes the rows the cache served _settle_chunks ended postprocess when no row that ran was left, but #768 routes a row the prefix cache served whole, which is not among them: a step holding only non-final chunks beside a served row would never complete that row's walk. It now counts cached_rids, as main's own check after the stop check does. Adds the _settle_chunks tests the review asked for: holding a non-final chunk, joining and releasing on the final one, and the early end. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: join a final chunk's outputs on the side stream, after step N only The join ran on the current stream, behind step N+1 that the gpu thread had already launched, while every copy of its result -- the check-stop D2H and the store's -- waited only on N's completion event and could read it before it was written (17 of 64 Qwen3-Omni answers with broken audio over SHM at c16, PR #365 review). The join now runs on the worker's side stream after N's event; the current stream waits on it, the joined tensors are recorded on it, and the step's completion event moves past the join, so those copies wait for it without waiting for N+1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: per-walk budgets inside a combined walk, behind can_override_mixed_max_tokens A walk inside a combined walk runs in the combined walk's steps, and the scheduler only ever read the combined walk's budget, so a per-walk override did nothing -- and the overrun warning then compared each step against it (PR #365 review). Now: - By default a walk's own budget applies only to a walk outside any combined walk, and a per-walk override of one inside is refused at load. The combined walk can be named directly (LLM: {mixed: 512}). - A node listed in the serving config's can_override_mixed_max_tokens caps each member walk's rows within a step by that walk's budget, its own or the model's, under the step's budget; null keeps that walk's prompts whole. - The overrun warning compares a step against the budget it was scheduled under, the combined walk's for a step relabelled to one walk, and names both causes of an overrun. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs: chunked prefill in adding_models Step 4 The hooks a walk opts in with (get_chunking_policy, get_input_sequence_len, split_inputs), how a chunked row's outputs are held and joined, and the two things a chunking AR node does itself that nothing checks at load: wrap its sampler step in keep_final_chunk_samples, and drop a non-final chunk's token or mark its edge FINAL (PR #365 review). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: a chunked row refused, cancelled or failed between chunks The paths the review found untested: an OOM refusal holds the row and its retry runs the same range, since nothing landed; and a request removed mid-prompt leaves no progress, cached prefix, held outputs or prepared input behind in the scheduler or the engine. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: a 2048-token default chunking budget The 1024-token budget was sized to fit a 2048-token prefill capture, but the capture ladder now runs to 8192. Swept on the stack (long mix of 4 ~200-token prompts per ~3k-token one, two restarts each), 1024 cost 3-10% closed-loop throughput at c48 (0.8B the most) in extra steps and fewer speculated ones; 4096 recovered it but gave back most of the short-prompt TTFT gain and doubled decode ITL p99 at c48 on 4B and 9B TP2. 2048 keeps throughput within about 1% of chunking off with most of the TTFT gain. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Engine: chunked prefill's state in ChunkingManagement The engine held seven attributes for chunking and combined-walk budgets: the per-walk policy cache, the serving config's budget overrides, the combined walk of each walk, the nodes whose walks cap their rows in a mixed step, the walk-cap cache, the overrun counter, and the whole prepared input of each row mid-prefill. They and the methods that only touch them (policy lookup, budgets, walk caps, the overrun warning, releasing a row's input) move to ChunkingManagement in mstar/engine/chunking.py, as SubmoduleManagement holds a node's. The engine keeps one-line delegates for what the scheduler and worker call. An engine fake now sets one attribute instead of a list that grew with every change to the engine's chunking. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a cut row runs at least a floor of tokens When decode rows used up a step's budget, max-min fair share cut a prompt to a token or two a step (or, after a speculation merge's continuing rows spent it all, to the one-token head rule). A cut row now runs at least min(128, the node's configured budget) tokens, or all it has left if that is fewer. When the share is under the floor, only as many rows as what is left holds are cut, sharing it; if it holds none, one row runs over budget by less than the floor. Rows that fit whole -- decode rows, short prompts, a last chunk -- are untouched, and check_budget allows the floor's slack. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: vision prefills take the Thinker's one chunking budget A Thinker step of any walk runs in thinker_mixed and so under its budget; the separate vision budget was read only as a can_override_mixed_max_tokens cap, where it sat above the step budget and did nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Chunked prefill: a row mid-prefill takes its next chunk in the speculative step The speculation merge builds step N+1 while N runs, and a row whose chunk is in N has its node off the ready queue until that chunk lands, so beside decode rows a prompt ran a chunk every other step. The merge is now handed N's rows that need another chunk (MidPrefillRow, with where the next chunk starts: the end of N's) and the builder takes them first, without a pop. When N's chunk lands, a row whose next chunk is in the step launched after N stays off the ready queue; any other goes back on it, so a speculation that is cleared, or a step refused at admit (which re-queues all its rows), needs no undo. A row the merge leaves out is neither parked nor returned. Only a merge that may take the row's walk does this: with speculation_mixes_walks off, a decode chain keeps to decode rows. From Gaurav Shah's version on #430, with the next chunk's start passed in rather than kept as a second progress cursor. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Respond to bot review comments: TP/SP refusal, stream block reuse, documentation * [Graph][Test] Parity: a Loop output read off an EMPTY_DESTINATION body edge Chunked prefill as a Loop gives the chunked node no loop-back edge and no destination of its own, so the Loop's declared output is the only thing that carries a chunk's result on. Nothing covered that shape: the existing loop cases all cache from a loop-back edge. The two runtimes cache loop outputs from different sources -- Python from the node's declared `outputs` list, Rust from the routed-edge slice -- and an edge that routes nowhere is where those sets could differ. They don't, but an empty `Loop.outputs` would drop the final chunk silently rather than raise, so it is worth a test rather than a reading of both. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Declare chunking on the model, and run a chunkable node in a Loop Chunked prefill's state lives in six places and each model implements four to six hooks for it. Expressing a chunk as one iteration of a graph Loop removes most of that: the Loop already owns per-request iteration state, re-queues its node, accumulates or drops per-iteration outputs, and is speculated like decode. This is the first step -- the declaration and the wrapping -- with no behaviour change. The policy moves from the submodule to the model because the graph build needs it: `get_graph_walk_graphs` is read long before any submodule loads, so a submodule-level answer is unreachable at the point the node has to be wrapped. A policy now names the walks it covers, and the model passes its walk dict through `apply_chunking`. `apply_chunking` is called by the model rather than applied to its result from outside, since `get_graph_walk_graphs` is what declares the graph and some twenty call sites read it directly. Forgetting the call would report a prompt done after its first chunk, so `_check_chunkable_nodes_are_looped` refuses it at startup instead. Policies and budgets are set separately: the policy is the model's and fixed, the budget is the deployment's and layered on top. The Loop holds `max_iters=1` for now. The worker still drives chunking and holds a non-final chunk's row back without completing its node, so the loop is only ever completed once, by the final chunk. It becomes a real bound when the engine takes over ending the loop. Nested chunk loops are refused at graph build: a chunkable node inside an existing Loop orders its completion cascade differently in the two runtimes. No model declares one. Tested on 2 x H100 (CPU-only paths; the box's GPUs were in use): - `ruff check .` clean. - `test/modular` against this branch's base (f28ce39a), both CPU-only: 6106 passed / 11 failed / 1 error, against 6104 / 11 / 1. Identical failure sets, all pre-existing. The two extra tests are an untracked shell script in the working tree that `test_launch_script_socket_prefix` parametrises over. - `test_startup_graph_validation`: 95 shipped configs pass, which is what checks that `apply_chunking` wrapped the right node in each of the four chunkable models. - Graph-runtime cost of the added Loop entity, interleaved A/B in one process on the Rust runtime, 4000 requests x 9 rounds per arm: -0.07 us/request against a 4-5 us per-arm spread, on ~21 us of graph work per request -- below the measurement floor. End-to-end still to run on a quiet machine. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Move a chunked node's outputs onto its Loop, and seed the loop-back A chunk is one iteration of the node's Loop, so the Loop is what should decide what a chunk's output means. The node's real edges move onto it: an edge the policy gathers collects every chunk's tensors, and any other carries the final chunk's alone, earlier ones dropped as the loop advances. That is what the worker's accumulator and side-stream join do by hand today. The body keeps a bare same-named edge to EMPTY_DESTINATION, because the Loop caches by edge name and so needs the name emitted. The copy drops every flag: persisting a per-chunk tick would stage a copy of each chunk for a remote reader, and counting one would bill the conductor a token per chunk. It also routes `chunk_tick_name` back to itself. Speculation offers a next-iteration target only for a node that routes an edge back to itself, so without it every chunk-only step is a pipeline bubble. A loop-back has to be declared as an input too, which leaves the node unready for its FIRST chunk -- nothing has produced it yet -- so the conductor seeds one on every entry into the walk. Re-entry included: a walk entered twice (BAGEL think-mode) takes its inputs through `_send_partition_inputs`, not `initial_inputs`, and a seed on the first entry alone would hang the second. The tick is free: the submodule produces no tensor for it, and a declared signal with no tensor still satisfies readiness. `max_iters` is now the policy's `max_chunks` rather than 1. It is a runaway guard, not the real count -- the engine ends the loop on the final chunk in the next commit -- and a prompt cannot have more chunks than tokens, so the model's longest prompt is a safe value. Tested (CPU-only; the box's GPUs are in use by other tenants): - `ruff check .` clean. - New parity cases over the real `_wrap_node_in_loop` output, not a copy of it, on both runtimes: first-chunk readiness, that the seed is what provides it, next-chunk readiness and speculability (`is_new_loop_iter=True`), that only the final chunk's token leaves the loop, and walk re-entry. - `test/rust/`: 39 failed / 426 passed, against 39 / 416 on the branch base -- same failures, all pre-existing, plus the 10 new tests. - `test_startup_graph_validation`: 95 shipped configs pass. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Let the Loop drive chunking, and drop the worker's parallel machinery A chunk is one iteration of the node's Loop, so the Loop decides when the prompt is done and what a chunk's output means. The worker's own version of that goes: the output accumulator, the side-stream join, the push-back onto the ready queue, and the surgery that took a non-final chunk's row out of the batch before postprocess could see it. `Worker._end_finished_chunk_loops` ends the loop for a row that just ran its last chunk. Keyed on the walk having a Loop rather than on the row having a chunk range: a row with NO range ran whole, which is also a final chunk -- a walk with no token budget is never measured or cut, and without this its Loop ran to max_iters while the request held its KV, which looks like pool exhaustion rather than a scheduling bug. What the Loop does not absorb is ~10 lines: until its last chunk a row's sampled token is not a real sample, so it decides no stop (a spurious one would cut the prompt short mid-prefill) and keys no prefix. That was enforced before by removing the row from the batch; now it is stated where the reason is. A streaming consumer needs the pieces joined, which is `Connection.join_dim`. Its ChunkPolicy counts ITEMS and a chunked producer sends one per chunk, so to Qwen3-Omni's FixedChunkPolicy(1) Talker the pieces looked like separate items and it read only the first -- which was empty, the opening chunk being all system prompt, every position of which the Thinker drops. Joined on arrival at the consumer, which owns those tensors outright: the producer-side join this replaces cat'd on the default stream behind step N+1 while a reader waited only on N. A consumer that can take the pieces as they come should do that instead and set no join_dim. The join registers in the producer's units, one per piece, because that is what `signal_done` reports. Counting a group as one left the buffer permanently short of the producer's total, so it never agreed the producer had finished and the Talker's speech came out truncated. Tested on 2 x H100 (slurm), CPU suite on a quiet box: - `ruff check .` clean. - `test/modular test/rust` against the branch base: 6536 passed, 46 failed, 1 error -- NO new failures; all 39 test/rust and the rest are pre-existing. Four base failures pass here: timeout-based tests that fail only under load (one reads 18s quiet, 180s under a concurrent GPU job). - Qwen3-Omni t2s, 192 requests per arm, Thinker prefill unchunked vs cut at 16 tokens so every prompt takes several chunks. Per-answer WER: each answer's audio transcribed with whisper-large-v3-turbo and scored against that same answer's text, over the spoken prefix -- the thinker writes about twice what the talker speaks, so scoring the whole text reads 0.54 on correct audio. | arm | n | WER | p50 | p90 | >0.5 | spoken/text | |---|---|---|---|---|---|---| | unchunked | 192 | 0.110 | 0.078 | 0.298 | 1 | 0.48 | | 16-token | 192 | 0.102 | 0.061 | 0.266 | 2 | 0.47 | Fisher exact p = 1.000 on the garbled counts; the unchunked arm's 1/192 is the model's own floor (#365 measured ~1/64 with chunking off). 192/192 wavs both arms, no failed requests. A budget that leaves prompts in one chunk is a false pass: the join takes a single-piece fast path and never runs. Confirm chunking happened first. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Key the chunk Loop set on the graph, and restore the prefix chain Two defects from review. The worker took ``_chunk_loop_walks`` from the model's chunking POLICY, but ``apply_chunking`` builds the Loop from that policy at graph-build time while the policy-derived set excludes nodes this deployment chunks nowhere. Under TP/SP the Loop is in the walk's graph and the worker's set omits it, so nothing ends it and it spins to ``max_iters`` holding the request's KV -- a hang on every TP deployment of a chunkable model, which the shipped-config check cannot see because it only verifies the Loop exists. Read off the graphs the worker was given instead (``_chunk_loops_in``). The policy says where a Loop may be built; the graph says where one is, and cannot disagree with itself whatever rule excluded a node. ``extend_prefix_chains`` had been dropped from postprocess entirely, so nothing keyed what a request generated and cross-request reuse silently stopped learning. Restored, with the rows mid-prefill skipped: a non-final chunk's sampled token is not a real sample, so keying the chain on it keys a token the request never generated. Tested on 2 x H100, CPU only: `ruff check .` clean; 386 passed over the chunk Loop, scheduler, chunking, admit-failure, prefix-stage and runtime-parity suites. New cases: the Loop ended for a walk whose deployment chunks nothing, the set read off wrapped and unwrapped graphs, a decode Loop not mistaken for a chunk Loop, and a guard that postprocess hands the mid-prefill rows to the prefix chain as skips -- that call was dropped once and nothing noticed. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Speculate into a chunk Loop, cutting the next chunk from a lookahead A chunk-only step had nothing queued behind it: the runtime never offered a chunkable (node, walk) as a speculation target, so the decode rows waiting on it paid the whole step. The filter was there because a speculated chunk k+1 needs where chunk k ENDS, and chunk k has not landed -- ``_assign_chunks`` would fall back to the end of k-1 and re-run k's tokens into the KV. ``ChunkProgress`` now carries that end. ``done`` is unchanged: the ground truth, advanced in postprocess when a chunk lands. Beside it a ``ChunkCursor`` records where the chunk just scheduled begins and ends, kept apart for a speculative schedule so a cleared speculation leaves no trace. A chunk is cut from the lookahead when there is one and from ``done`` otherwise, and landing promotes the speculative cursor -- the step it was cut for is the one now in flight, so it is no longer a guess. The lookahead is only ever ahead of ``done``, so every abandon path just discards it. Three reach it: the worker's push-back (admit refusal, OOM), the scheduler's own hold backoff, and a speculation cleared before it runs. The hold is the one a test caught -- it never touches the worker, so a held row retried past tokens it had not run. ``MidPrefillRow`` needed no change: its construction was never gated on walk mixing, only on the spec target being the chunked node, which the filter made impossible. TP is unaffected -- chunking is off for TP/SP nodes until #367 -- so the lockstep question does not arise here. Tested on 2 x H100, CPU only: `ruff check .` clean; 454 passed over the runtime-parity, scheduler, chunking, chunk-Loop, speculation, admit-failure, resident-delta and worker-drain suites. New cases: a speculative chunk cut after the one in flight, a cleared speculation leaving the next chunk where it was, and a landed chunk promoting the speculative cursor. NOT YET MEASURED. The point of this is closing the bubble, and that claim needs a `benchmark/worker_phases` breakdown showing it; the machine this was written on could not give a reliable allocation. Until then, treat the win as unverified -- correctness is tested, the performance argument is not. Note for bisecting: this removes a Rust constructor argument, so it and the commit before it need different `maturin develop --release` builds. A stale extension fails at request time with a TypeError, not at import. Co-Authored-By: Claude <noreply@anthropic.com> * [Chunked Prefill] Mark a chunk Loop's last iteration so speculation skips it A prompt short enough to run in one chunk still goes through a one-iteration chunk Loop, whose loop-back makes the prefill node a speculation target after every prefill step. Every row then dropped out of that speculated step (its loop-back never arrives), so decode lost its overlapped step: on Qwen3.5-4B short prompts the branch was -1.5% tok/s at c16 and -3.8% at c48 against the main it merged, ITL p95 11.1 -> 14.2 ms, with ~500 "Speculation: dropped rids" warnings per run. The runtime now takes last-iteration marks: `mark_last_loop_iters(rids, walks, loops)` says a row's scheduled iteration is its last, and `_can_continue_loop` honours it. Only speculation reads a mark; the loop still ends by its stop, so a stale mark can cost at most one unspeculated step. A mark goes when the loop's stop is applied, on `clear_last_loop_iters`, and when the request is removed. `speculate_node` takes every row of the step, grouped by walk, rather than one sample rid per walk: a loop-back target is kept only if some row can continue, so a walk whose first row is on its last iteration still speculates for the rest. The scheduler marks rows running their final chunk (or no range) when it builds a step, and `drop_chunk_lookahead` clears them, since a refused step's retry is marked again. Python and Rust runtimes both carry it. Measured on 2 x H200, Qwen3.5-4B, short prompts: c16 3576 tok/s vs 3576 on main, c48 8027 vs 8043, and no dropped-rids warnings. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Don't publish a stopped loop for a cache-served row A chunkable walk the prefix cache serves whole runs no iteration, but its chunk Loop is still stopped. Publishing after that stop looked up the row's forward-pass info, which a row that did not run does not have: a KeyError that failed every request in the step. Seen on Qwen3-Omni with prefix caching. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Default max_chunks to effectively no cap `ChunkingPolicy.max_chunks` becomes the chunk Loop's `max_iters`, and it defaulted to 1. BAGEL and Orpheus never set it, so any prompt cut into chunks ended its prefill walk after the first chunk; the conductor then started decode with the rest of the prompt still to run, and decode steps carried it ("plan label 'main' carries 929 tokens but its captured bucket holds 16"). The count is only a runaway guard: the engine ends the Loop on the final chunk, and the number of chunks is not known until the budget cuts each step. So the default is now 2**31 - 1, which fits the runtimes' u32 `max_iters`. Tested on 2 x H200, prefix cache off: Orpheus with a 64-token budget went from 93/287/188 failed requests (c16/c48/Poisson) to 0; BAGEL short, mixed, long and image workloads, closed loop and Poisson, ran with 0 failed requests and no server tracebacks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Measure a walk the cache holds whole at 0 tokens A chunkable keyed row is measured without its cached prefix, but only a partial match was taken off: a walk the cache holds whole was measured at its full length. The budget then cut it into chunks from 0, and a chunk cut from 0 writes over the lease the walks before it were served from, which fails the request whenever those walks end inside a page ("the walks the cache served end inside a page (43 slots), so their KV cannot be kept"). Qwen3-Omni image requests hit it whenever an image repeated: the text walk is served (43 slots), then the vision walk, held whole, was chunked. Measured at 0 it takes no budget and no range, and prepare serves it from the cache like any unchunked row. Tested on 1 x H200, Qwen3-Omni colocated, prefix cache on, images repeating, c8: 8 of 48 image requests failed before, 0 after, with no server tracebacks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Refuse a Thinker input split_inputs cannot cut The merge took main's `split_inputs`, which cuts the inputs it knows and silently drops any other, where this branch's refused one. A dropped key would leave a chunk missing an input it was prepared with, so refuse it again. `test_qwen3omni_vision_chunks.py::test_an_input_it_cannot_cut_is_refused` failed since the merge and passes again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Drop the explicit max_chunks from Qwen3.5 and Qwen3-Omni The default no longer caps a prompt's chunks, so the per-model value (each model's max_position_embeddings) is redundant. The docs described max_chunks as something to size to the longest prompt; they now say to leave it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Split a step's prompt budget in start order The token budget was a max-min fair share across every waiting prompt, so with several long prompts each advanced a little every step: all of them finished late, and their KV pages were held the whole time. The default builder now fills the budget with the prompts that fit whole (shortest first), then gives what is left to prompts in the order their first chunk was cut, cutting at most one. A prompt mid-prefill keeps its start position, so it finishes before a later one begins. The equal share moves to FairShareBatchBuilder, selected with `batch_builder: fair_share`. On its own this trades ITL for TTFT (H200, Qwen3.5-4B, against the equal share): long prompts c32 TTFT p50 1.3-1.4 s -> 0.37 s but ITL p95 43 -> 80 ms; very long c16 TTFT p50 3.3-3.8 -> 3.0 s, ITL p95 86-100 -> 130 ms. A prompt left out of a step now waits as a backlog, and the step after a prompt step serves it alongside the decode rows. The next commit gives that step to the decode rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Let ready work go first after a chain that cannot continue The fairness peek (`has_ready_excluding`) yields a speculation chain to other ready work only once the chain has run a step past its first. A chain that could not continue at all skipped it, and the fallback `get_next_batch` was unrestricted, so a (node, walk) with a backlog won. A prompt-only chunk step is such a chain (rows mid-prefill are not speculated), so under the start-order builder the next step was another prompt's chunk with the decode rows riding along through the prefill kernels. The peek now applies there too: with other ready work, the next step is that work, giving two decode steps per prompt step. With `MSTAR_SPEC_PEEK_FOR_FAIRNESS=0`, or nothing else ready, nothing changes. H200, two rounds, start order with this against the previous equal share: Qwen3.5-4B very long c16 ITL p95 86-100 -> 49 ms; long c32 TTFT p50 and ITL p95 flat (1.36 s, 42 ms). Qwen3.5-4B mixed c48 ITL p99 flat, +1.2% tok/s; BAGEL mixed c48 ITL p95 54-69 -> 44-45 ms, +2.6% tok/s; Qwen3-Omni mixed c24 ITL p95 80 -> 64 ms, +1.5% tok/s. TTFT on long prompts is in the next commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Park the prompts a step's budget leaves out, in line A fresh row the token budget left out went back to its ready queue, and a backlogged one was parked behind the rows the row cap pushed out. The Rust runtime scans ready rows in handle order and the backlog is served first, so either way a prompt rejoined behind ones that arrived after it. Under start order, which leaves out most prompts every step, the first prompts of a burst started last. Every row the budget leaves out now parks ahead of the cap's leftovers, so prompts start in arrival order. H200, two rounds, against the previous commit: BAGEL long c32 TTFT p95 44.2-44.4 -> 7.2 s, Qwen3.5-4B very long c16 28.9 -> 6.3-6.5 s, with mean TTFT and E2E unchanged (BAGEL 5.8 -> 5.9 s, 7.5 -> 7.6 s; Qwen 5.2 -> 5.1 s, 7.4 s). TTFT p50 rises to match (0.7 -> 6.2 s, 1.2-1.5 -> 5.4 s): the low median came from the prompts left waiting. Against the equal share before start order, E2E mean drops 4-8% (BAGEL 7.9 -> 7.6 s, Qwen 7.9-8.4 -> 7.4 s) and TTFT p50 rises (2.8-4.2 -> 6.2 s, 3.3-4.7 -> 5.4 s). Qwen3.5-4B long c32 and mixed c48 are flat; BAGEL mixed c48 is -1.3% tok/s, TTFT p95 2.3 -> 2.9 s, p99 4.6-6.6 -> 4.2 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [GDN] Size a captured mixed step's decode half to its bucket `GDNMixedWrapper` runs its decode half on every row it is sized for, and a captured wrapper was sized to the largest capture (`_cg_max_bs`), the floor the other wrappers take for their buffers. So every mixed bucket ran the GDN decode kernel on 128 rows: on Qwen3.5-4B at c48, ~2.6 ms of a mixed step was decode work on rows addressing the null slot. The split wrapper now takes the bucket's rows; the plain decode and prefill wrappers keep the floor, since their kernels run on the step's rows whatever their buffers hold. The dropped rows address only the null slot: outputs and every real slot's state are bit-identical to the 128-row wrapper (`test_the_decode_halfs_width_changes_no_real_rows_result`). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Engine] Order a split walk's steps with its prompt rows last A resource that runs a combined walk on split kernels (CombinedSplit) reads the step's one-token rows on a decode kernel and its prompt rows on a prefill kernel. The engine records each split's prefill walks at load, refusing two resources that disagree, and the worker's one batch builder puts those walks' rows last, keeping each side's order. The kernels can then read the decode rows as the step's first tokens in place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [GDN, Attention] Read a mixed step's decode rows in place GDNMixedWrapper and FlashInferMixedWrapper gathered every token into decode-then-prompt order and scattered the outputs back, per layer (index_select / index_copy_ over the whole step). With the engine putting the decode rows first, the decode half reads the leading run of one-token rows as the step's first tokens and a bs-row torch.where merges its outputs in place. Under capture the prefill half takes every token, the run as one sink row ahead of the prompt rows; an eager step's prefill half takes the tokens after the run, so it writes prefix-cache checkpoints only for the prompt rows its checkpoint buffer is sized to. Row order only bounds capacity: a decode row behind a prompt row runs on the prefill half, which is still correct while the step holds no more prompt rows than the bucket was captured for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Attention] Keep a mixed step's padding rows out of the decode split budget FlashInfer's tensor-core decode splits long KV rows across CTAs within a budget of 2 CTAs per SM over the KV heads (66 rows on an H200 with 4 KV heads), and every row it plans costs at least one, padding rows included. A 64-row mixed bucket holding ~46 decode rows left no room to split, so its long-KV rows ran unsplit: 328 us a layer against 135 us for the same rows in a 48-row decode bucket. The split's decode half now runs the same FA2 paged kernel planned through the prefill wrapper, where a row with no query costs nothing: the run's rows take one token each and one row carries the rest. Measured on 46 rows with 15% at 4k-16k KV: 159 us at 64 rows. The CUDA-core option (decode_tensor_cores: false) keeps the decode wrapper. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] docs: cut the model author's chunking section down Keep what a model author does: declare policies, wrap the walks with apply_chunking, the policy fields, join_dim, the two length/cut hooks and the final-chunk sampling rules. Drop engine internals (how resources read reserve_whole_prefill, why join_dim joins on arrival) and the stale note that the engine holds and joins non-final outputs, which the Loop does now. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] reserve_whole_prefill as a per-row setting StepContext.reserve_whole_prefill was one bool for the step, the AND over its chunked rows, so a step mixing a walk that opts out with one that opts in turned the reservation off for both. It is now a PerReqSetting (SET, UNSET, or PARTIAL with the rows it holds for), and the KV manager reserves ahead per row. No model opts out today, so no shipped step changes. Also cuts StepContext's comment on it to three lines. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A serving config can set min_chunk_tokens The chunk floor was a module constant. It is now the serving config's min_chunk_tokens (default 128), set on the engine's ChunkingManagement with the token budgets and validated there. Both readers take it from the engine: the micro-scheduler's cut and check_budget's overrun tolerance. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Comment fixes from review - ChunkingPolicy: the multi-line field comments move into the docstring. - PythonGraphRuntime._combined_walk_of gets its own comment back, which a note about the removed disable_spec_node_walks had replaced; the same leftover goes from speculate_node's docstring and from the Rust runtime's struct and constructor, where rustdoc attached it to the next field. - MicroScheduler: drop a comment whose field went with the accumulator, and finish ChunkProgress.cursor_spec's. - Worker._advance_chunks keeps only what it does now. - Engine: a blank line before _prepare_chunk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Name a chunk Loop and its tick from (node, walk) alone _wrap_node_in_loop built the tick name from the loop name it was handed, a second spelling of chunk_tick_name that matched only while every caller passed chunk_loop_name(node, walk). If they drifted, the conductor would seed an input the node does not declare and the first chunk would never become ready. It now takes neither name and calls chunk_loop_name and chunk_tick_name itself. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Graph] Rename last-iteration marks to do-not-speculate marks mark_last_loop_iters / clear_last_loop_iters read as if they changed how a loop is scheduled, but only speculation reads them: the loop still ends by its stop. They become mark_do_not_speculate_loops / clear_do_not_speculate_loops, with the runtime state, the scheduler's _mark_final_chunks (now _do_not_speculate_final_chunks), tests and docstrings following. Python and Rust runtimes both, which needs a rebuilt extension. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Load the submodules in EngineManager.build again The worker loaded its submodules before building the graph runtime so the runtime could be told which (node, walk) pairs never to speculate into. That argument went when speculation into a chunk Loop landed, and nothing else needs the submodules before the runtime exists, so the split and its stale comment go: build loads them, as on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Move the chunk Loop bookkeeping into ChunkLoops Worker had grown four chunking handlers. They move to ChunkLoops in mstar/worker/chunk_loops.py: the (node, walk) pairs read off the graphs at load (Worker.chunk_loops.walks, and walks_by_node for the scheduler), ending a Loop on its final chunk (end_finished) and recording what each chunk ran (advance). The worker calls them at the same points of postprocess as before, on the same thread; no behaviour change. Tests drive a real ChunkLoops over a real ScheduledBatch instead of a bare Worker. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] A decode step beside a chunk Loop pays one lookup, not one per row ChunkLoops.end_finished walked every row of every step on a node with a chunk Loop, decode steps included, only to find decode is not a chunk walk. It now returns early for a single-walk step whose walk has no chunk Loop, as the scheduler's _do_not_speculate_final_chunks already does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Speculate a prompt's next chunk on its own A step whose rows were all mid-prefill never speculated: the worker took the held rows out of the candidates and built the merge that carries their next chunk only behind a continuing row, so with walk mixing off (the default) a prompt's chunk k+1 waited for chunk k to land. A speculation can now be made of those rows alone, its capture group taken from one of them, and the merge keeps to their walk. So the step after a prompt step still goes to other ready work, the fairness peek runs for a prompt chain from its first step: two prompt steps in a row only when nothing else is ready. `speculate_mid_prefill` (serving config, per node or per walk, default on) turns it off; off, the next chunk waits for the current one to land, in a mixed step too. H200, Qwen3.5-4B, two rounds, PR head / on / off: a lone very long prompt's TTFT p50 455-456 / 437-441 / 451-455 ms (p95 582-585 / 564-566 / 580-582); a lone long prompt 119-121 / 116 / 117-118 ms. Long c32, very long c16 and mixed c48 are flat in tok/s, TTFT and ITL. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Skip the chunk bookkeeping when nothing is chunked _assign_chunks returns at once when no row is budgeted, in progress or holding a cached prefix (4.9 -> 0.1 us for a 64-row decode batch), and drop_chunk_lookahead walks chunk_progress once against a set of the rids rather than once per rid (16 -> 0.9 us for 64 rows, 8 prompts in progress). A StreamBuffer entry that is a lone tensor, every item of an unjoined stream, skips the joined path (0.88 -> 0.53 us per item; main is 0.37). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: main's request-added and speculation-readiness fakes carry the chunking fields Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Worker] Key a step's prefix chains before its withheld outputs are dropped Main indexes the stop check's host copy before dropping withheld and empty-span outputs; the merge had moved the indexing after the drop, which on a step without a completion event (where the host copy is the outputs themselves) left a withheld token unkeyed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A cut row waits a step rather than run the step over its budget The row a step's budget cuts took at least the chunk floor even when less than that was left, so a step that finished one prompt could run up to the floor over its budget: 1990 + 128 tokens under 2048, say. Qwen3.5's prefill captures go 2048 -> 4096, so that step padded to the next bucket up, and the budget check allowed exactly the floor's slack, so it was never logged. A row now waits for the next step when less than the floor is left, under both start order and the equal share; the floor still sets the smallest chunk. The budget check warns on any step over its budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A chunk Loop counts the tokens its chunks ran Chunk progress lived in the scheduler (ChunkProgress and its lookahead cursors), because the Loop only counted iterations up to a runaway guard and could not tell when a prompt was done. So the worker ended the Loop by hand on the final chunk (ChunkLoops.end_finished), both runtimes carried do-not-speculate marks for the final chunk, and every abandon path had to drop the lookahead or the retry would skip tokens. The Loop now counts tokens. Each row's completion moves its enclosing loop by a stride (RouteInput.loop_strides): a chunk's ChunkRange.stride, the END stride for a chunk-walk row that ran whole, one otherwise. A prompt's first chunk caps its Loop at the prompt's length (GraphRuntime.set_loop_caps, which assigns rather than lowers, so the walk's next prompt sets its own). The Loop then ends itself on the last chunk, and speculation eligibility takes the in-flight step's strides (speculate_node, prep_spec_rids, prep_spec_rows), so nothing speculates past it. The scheduler resumes a popped or backlogged row mid-prefill from the Loop's count (GraphRuntime.loop_progress), and a speculative one from the end of the chunk in flight. Nothing is cut ahead of what landed, so nothing has to be undone when a step does not run. Gone: ChunkProgress, ChunkCursor, advance_chunk, drop_chunk_lookahead, ChunkLoops.end_finished/advance, mark/clear_do_not_speculate_loops and ChunkingPolicy.max_chunks. Kept: the loop-back tick, which is what makes the next chunk a speculation target, and the start order the FIFO builder serves prompts in, now a (rid, node) -> sequence dict. Both runtimes change alike and the parity tests cover the new calls. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] The sampler keeps rows mid-prefill out of its commit A non-final chunk's sampled token is not a real sample, so its RNG offset and seen-token mask must not commit. Every chunking model wrapped its SamplerStep in keep_final_chunk_samples for that, and a model that forgot it erred nowhere: its outputs just changed once its prompts were chunked. The sampler now does it itself, the way the KV and recurrent resources apply a step's forks only to fork_rids: it drops the rows of the step context's chunk_tokens_after (the rows a chunk runs after this step) from what its draw commits, eagerly and under capture. keep_final_chunk_samples goes. The models' postprocess guards that dropped a non-final chunk's new_token go too: the chunk Loop carries only the final chunk's outputs (the body's copies route to EMPTY_DESTINATION with no flags), and the worker takes no stop or prefix key from a row mid-prefill. Their tests become the sampler's own and the runtimes' (test_only_the_final_chunks_token_leaves_the_loop). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Model wraps chunkable nodes wherever the walk graphs are read Each chunking model had to pass its walks through apply_chunking at the end of get_graph_walk_graphs, and a load-time check caught the ones that did not. Model.__init_subclass__ now wraps every subclass's own get_graph_walk_graphs so its result goes through apply_chunking, which makes the ~20 call sites that read the walks see the chunk Loops with no accessor to route them through. apply_chunking leaves a node already in its chunk Loop alone, so a subclass building on its parent's walks is wrapped once. The four models drop the call, and the check that a chunkable walk is wrapped goes, since it can no longer fail. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] The oldest prompt mid-prefill always gets the chunk floor Under start order, prompts that can finish this step go ahead of the ones already mid-prefill, and nothing bounded how long those could wait. With 16 decode rows at a 2048 budget and four new 500-token prompts each step, a 20k-token prompt that had started got 0 tokens for as long as the stream lasted. The oldest prompt mid-prefill now keeps min_chunk_tokens out of the finishing pass, so it advances every step; the step stays within its budget (one new prompt waits a step instead). Unchanged when nothing is mid-prefill. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Test: track the chunk Loop's stride end and main's dtype fake EngineManager.build reads the model's chunking policies on this branch, so the fake that stands in for what build reads (main's #810) has to answer it too. A cache-served row's chunk Loop now ends by its route stride rather than a stop, so the dropped-row test checks the stride. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Seed a chunk Loop on the ranks hosting its node The conductor added the chunk ticks after splitting a walk's inputs, so they reached only the ranks that already got an input. Qwen3-Omni's prefill_audio sends every edge to the audio encoder, so with the encoders and the Thinker on different ranks the Thinker's Loop was never seeded. The ticks now join the inputs before the split, which broadcasts a signal-only edge to the node's own ranks, and only where a walk is entered: a client stream chunk carries none. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] Validators read a chunked node's edges off its Loop After apply_chunking a chunked node's body routes only to EMPTY_DESTINATION and its tick, so streamed_partitions lost Qwen3-Omni's Thinker->Talker edges and the conductor's prefix-cache guard for the Talker never fired. Both validators now read the enclosing chunk Loop's outputs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [Chunked Prefill] A hit that lands after a row was measured is declined A row measured with nothing cached can find a hit by its first chunk, when an earlier request with the same prefix commits while it waits. The first chunk treated that as the policy declining and applied no prefix, which ends the row's chain. It now applies a zero-token prefix, keeping the chain, and the recurrent pool gives back a snapshot a zero-token prefix does not restore. A walk whose policy declines reus…

Gaurav-Shah05 added 30 commits September 30, 2026 02:35

conductor: stamp a request's max_tokens on every kv config

b685b6d

The pool has to know how far a request can grow before it lets it in. A model that returns no config for a KV cache, as Whisper doesn't, now gets a default one to carry it.

test: every kv cache is handed the request's max_tokens

d5ffdc4

kv: count the pages the index can give back

bed58c0

Eviction takes leaves only, so a page only the index holds is out of reach while a page below it is held. Admission counts what eviction can reach, and asks what would match without re-stamping the LRU.

test: the index counts only what eviction can reach, and a peek re-st…

6b49857

…amps nothing

kv: never chain a page under one only the index holds

794d1de

A live child would keep eviction from a page the pool counts as free, so the room admission counted could vanish after it admitted into it.

test: a page only the index holds stays evictable when the stream aft…

bfe82ed

…er it fills a page

model: bagel stamps each label's prompt slots

60f48fd

An image request writes image tokens, latents and guidance copies that no prompt key describes, so the pool can only size it from the model's own count. The conductor hands each cache its share, next to the prefix keys.

test: bagel counts every token an image request writes to each label

00a19d2

test: a reservation covers the prompt's uncached pages and max_tokens…

720153c

… of growth

engine: tell a resource which of its spec's nodes this worker runs

497b090

Under CFG parallel each worker builds its own cache from a spec naming all three LLM nodes, so a request's reservation counted the guidance nodes' labels on every worker.

test: a worker reserves only the labels its own nodes open

73d0008

test: a request that doesn't fit is not ready until room frees

8d68c90

model: whisper stamps its caches' prompt slots

25c412c

The decoder's stream is the forced prompt plus what it decodes, and the audio context is always the whole window, so both caches are counted exactly and admitted by reservation.

test: whisper hands each cache the forced prompt and the audio window

b2ad9f5

test: a reload waits for its reservation, and nothing takes its room …

6eef6ea

…meanwhile

api: a client never sends the cache its own page counts

99ac542

A model that stamps no prompt_slots let a client's own reach the pool, which admits by them: a small count takes room others were promised, and a malformed one raises in the gate.

test: a client's page counts never reach the conductor

d8723a8

test: a replay's padding rows take no page of their own

d9e8c2d

kv: cap what decode adds, not the model's own count

483e511

Bagel counts an attachment's image tokens, which are not positions, so a request with enough images passed max_seq_len and was reserved short.

test: a model's count past max_seq_len is reserved whole

e0e1a85

test: under CFG parallel the leader's cache decides, and the guidance…

7d434ca

… caches follow

kv: drop the preplan's lease check

ecd8379

No preplanned step carries a lease: a preplan runs only on a batched capture, and every keyed walk is a packed one.

conductor: stamp a request's resource configs in one place

2321039

So the stamp every cache admits by is the one its tests exercise, not a copy of the loop.

test: the conductor's own stamp hands each cache its max_tokens and c…

ce2a761

…ounts

stephen-dwq requested changes Oct 18, 2026

View reviewed changes

dghosef mentioned this pull request Oct 20, 2026

PR labels: area, model and size from the file list, type from the review bot #738

Merged

This was referenced Oct 20, 2026

[KV cache][Bug fix] Size and capture every bucket on SINK_PAGE, so a small KV pool still boots #729

Merged

[Micro-scheduler][Bug fix] Retry a refused batch smaller when nothing can be evicted #651

Merged

github-actions Bot added area/api area/conductor area/engine area/model-shared area/resources area/worker model/many size/L labels Oct 20, 2026

dghosef added type/feature type/bug labels Oct 20, 2026

This was referenced Oct 22, 2026

kv: page footprints from one layout per request #832

Merged

prefix caching at TP > 1 #835

Merged

kv: a host index keyed by prefix #834

Merged

arun-elr reviewed Oct 24, 2026

View reviewed changes

arun-elr mentioned this pull request Oct 24, 2026

[Worker][Bug fix] On a KV OOM, try the next eviction victim when one will not move #461

Draft

2 tasks done

Gaurav-Shah05 mentioned this pull request Oct 25, 2026

[Preemption 1/7] Slim Admission Control (A1) #869

Merged

Gaurav-Shah05 closed this Oct 25, 2026

NSagan271 mentioned this pull request Oct 27, 2026

[Chunked Prefill] Chunked prefill as a graph Loop (replaces #365) #916

Merged

2 tasks done

Navigation Menu

Repository navigation

Admission control for KV resources - #400

Admission control for KV resources#400
Gaurav-Shah05 wants to merge 96 commits into
mainfrom
admission-control

Gaurav-Shah05 commented Oct 10, 2026 •

edited

Loading

Uh oh!

stephen-dwq left a comment

Uh oh!

Gaurav-Shah05 commented Oct 20, 2026 •

edited

Loading

Uh oh!

arun-elr left a comment

Uh oh!

arun-elr Oct 24, 2026

Uh oh!

Gaurav-Shah05 Oct 24, 2026

Uh oh!

arun-elr Oct 24, 2026

Uh oh!

arun-elr Oct 24, 2026

Uh oh!

Gaurav-Shah05 Oct 24, 2026

Uh oh!

Gaurav-Shah05 commented Oct 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

Gaurav-Shah05 commented Oct 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

How it works

Trade-off

Next

Testing

Checklist

Uh oh!

stephen-dwq left a comment

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

arun-elr left a comment

Choose a reason for hiding this comment

Uh oh!

arun-elr Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

arun-elr Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

arun-elr Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 commented Oct 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Gaurav-Shah05 commented Oct 10, 2026 •

edited

Loading

Gaurav-Shah05 commented Oct 20, 2026 •

edited

Loading