Repository navigation
per-node dtype, yaml doesn't override autocast, downcast warning - #810
Merged
Merged
Conversation
mstar-review
Bot
added
area/engine
area/model-shared
area/worker
model/bagel
size/M
labels
Automated review🟡 No findings on Description
What the review checked
Advisory, not a merge gate. Rules: |
vasilevklart
approved these changes
Oct 21, 2026
vasilevklart
left a comment
vasilevklart
left a comment
Collaborator
There was a problem hiding this comment.
Lgtm, I am looking more closely into the weight side: per tensor weight schemes so we can load any quantized checkpoint (fp8 block / awq / gptq / nvfp4 / gguf) without touching model code, hope we can get glm fp8 merged soon so we can exercise all this agentic side sooner
| resources=resources | ||
| resources=resources, | ||
| node_dtype=( | ||
| node_dtypes[node_name] if node_dtypes is not None and node_name in node_dtypes |
Collaborator
There was a problem hiding this comment.
node_dtypes and the fallback compute the same thing from the same inputs
|
|
||
|
|
||
| @contextmanager | ||
| def node_scope(node_name: str | None) -> Iterator[None]: |
Collaborator
There was a problem hiding this comment.
nit: capture_into_graph enters autocast_scope without a name, so node_scope(None) clears it and a downcast during capture logs <unknown>
Conflicts in mstar/engine/engine.py and mstar/worker/engine_manager.py: both sides added independent lines (main's fp32_precision field and moe_parallel config beside this branch's node_dtype field and node_dtypes dict), so each keeps both; the engine_manager import list takes main's shorter form, since main moved the refusal helpers that used NodeResourceSpec out of the module. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Engine.load_model took a node_dtypes map from the engine manager and fell back to resolve_node_dtype(submodule, self._autocast_dtype) for a missing node. Both come from the same submodule and the same model dtype (the engine manager passes its resolved autocast_dtype to Engine), so the map added a second path to the same value. load_model now always resolves. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ings capture_into_graph and capture_with_static_outputs entered autocast_scope without a node name, so node_scope(None) cleared the name the warm-up had set and an implicit downcast recorded during capture logged <unknown>. Both now take node_name, and every runner call site passes its node. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Since the merge, build reads model_config["node_groups"] (for each node's ep_config) and model.kv_labels() (for prefix caches), and no longer reads prefix_key_streams. Give the fake config one node group and the fake model kv_labels in place of prefix_key_streams. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A custom op inside a compiled forward sees no autocast, so rms_norm and RoPE took the no-autocast branch and warned for nodes that autocast in bf16 (Bagel's LLM, for one). node_scope now carries the node's resolved dtype; the warning stays silent when the op ran in that dtype, and names the node's dtype when it ran in another one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dghosef
added a commit
that referenced
this pull request
Oct 30, 2026
Conflicts with #810 (per-node dtype): NodeSubmodule.compute_dtype is deleted in favour of main's resolve_node_dtype, which the engine's _compute_dtype now calls; _autocast_dtype_for takes main's node-name form; the engine manager keeps main's resolve_model_dtype and parses the YAML string with parse_autocast_dtype. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Merged
2 tasks done
dghosef
added a commit
that referenced
this pull request
Oct 30, 2026
…split-2-speech-language Brings #858's deletion of NodeSubmodule.splits_heads in favour of RaggedAttentionConfig.shard_over_sp, and origin/main's #810 (per-node dtype and autocast), #961 and #931 (Rust loop-back fix). The compile-region warm-up now reads the node's resolved dtype by name and names the node in its autocast scope. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
NSagan271
added a commit
that referenced
this pull request
Oct 31, 2026
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>
NSagan271
added a commit
that referenced
this pull request
Nov 1, 2026
* 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…
NSagan271
added a commit
that referenced
this pull request
Nov 1, 2026
* 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…
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Preliminary step to proper per-node quantization of resources/computation.
Per-node override is no longer ignored at runtime.
Engine._autocast_dtype_foronly checkeddisable_autocast.NodeSubmodule.get_autocast_dtype()is only honored by load-time castingdtypeonce inEngineManagerand it is stored inSubmoduleManagement._autocast_dtype_fornow returns thisdtype.Various stale docstrings
YAML can no longer re-enable autocast in models that opt out of autocast.
Nodes that opt out of autocast still get
bf16norm and RoPErms_norm._prepare_inputsandRopeManager.apply_qkon CUDA still turnedfp32tobf16.Bagel no longer calls device-less autocast
How was it tested?
test/modular/test_position_autocast.pyandtest/modular/test_node_dtype_resolution.py(8 pass): per-node dtype resolution, the warning naming the node in eager and in both capture paths, and no warning for a downcast to the node's own dtype (that test fails without the fix).Checklist
ruff check .passes