Repository navigation
helios/scheduler: deterministic solo prefill chunking by remaining prefill segment - #22
Open
tunamitom wants to merge 1 commit into
Conversation
…efill segment Shared scheduler policy change (every model's solo prefill; not a Kimi change). When no other lane competes, one request's prefill chunking is a function of its remaining prefill segment alone -- the tokens to the next message boundary or prefill limit, cut at atomic spans -- not of the prompt length: full-size chunks while the remainder is at least one chunk, then the single smallest graph covering the rest. Measured durations never enter the choice (cost::deterministic_prefill_capacity + the policy application sites), so the launch sequence is identical across boots and requests; under concurrency the measured catalog competes as before and decode interleaving is unaffected. A mixed shape may carry a solo chunk only when its prefill row is exactly the deterministic chunk graph. Liveness fix (this rework): the rule offers one capacity, and the candidate filter still drops a prefill candidate whose graph is unmeasured when a smaller catalog graph would consume the same work. An atomic span (or a message boundary plus a span) that straddles the end of the full chunk cuts it back to the span start, where the unmeasured full graph consumes exactly what a smaller measured graph consumes: the single offered candidate was dropped and the solo request stalled forever after the Init graph. The rule now bails whenever its chunk would not survive the candidate filter (no measured duration, an unmeasured chunk duplicating a smaller graph's work, or a chunk consuming no work), leaving the choice to the measured catalog, and the candidate site falls back to the full catalog if the rule's candidate is rejected anyway. The regression tests reproduce the confirmed stall (both fail against the pre-fix rule) with atomic spans and with message boundaries. TTFT evidence (scheduler harness before/after, not a live serve): simulated first-token time and prefill launches, before (cost-based) / after (rule), per prompt arm, under a clean measured catalog and one where a slow capture over-measures the full graph by 1.6x (true execution model unchanged): 3.5k clean: 3 launches 62.10 ms -> 1 launch 66.10 ms 3.5k drift: 3 launches 62.10 ms -> 1 launch 66.10 ms 14k clean: 6 launches 232.10 ms -> 4 launches 232.10 ms 14k drift: 9 launches 238.10 ms -> 4 launches 232.10 ms Read honestly: TTFT is neutral on the long prompt when the catalog is faithful and 6.1 ms better when it drifts; on the short prompt the rule pads the 432-token tail to the covering 4096 graph and costs 4.0 ms more than the cost model's tail splitting (its existing tuned behaviour) to buy one launch instead of three and a boot-invariant sequence. End-to-end serve TTFT was not re-measured on this tree. Maintainer decision requested: this replaces measured-cost selection with a fixed rule for solo prefill chunking. Accept the determinism/padding trade (numbers above), or prefer the narrower alternative of fixing the measurement drift the rule bypasses -- in that case drop this change. Validation: helios `cargo test --lib -- scheduler::` 278 passed, 0 failed, 2 ignored on this branch (release/k31-exl3-tp8-20261108 tip: 273 passed including the harness test, 0 failed, 2 ignored; no pre-existing failures on either side). The two stall regression tests fail against the pre-fix rule (reproduced before fixing). Material limitations: no live model serve was run; the shared policy is model-agnostic and only scheduler unit tests exercise it (no per-model runs -- untested on every recipe, including Kimi K3/K3.1); no CUDA code is touched. Chunk boundaries differ from the cost model's choices, so a given prompt's KV/state rounding order can differ from the previous scheduler's while becoming reproducible across boots. Numerics class: NUMERICS-NEUTRAL at the arithmetic level (chunk selection only; per-chunk compute and the exact reductions are unchanged), with the caveat above about chunk-boundary rounding order. Design-Signed-Off-By: Hermes Agent (hermes-agent); design-review=a38f
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 / why
Shared Helios scheduler policy change (every model's solo prefill). When no other lane competes, one request's prefill chunking is a function of its remaining prefill segment alone — the tokens to the next message boundary or prefill limit, cut at atomic spans, not "prompt length alone": full-size chunks while the remainder is at least one chunk, then the single smallest graph covering the rest. Measured durations never enter the choice; under concurrency the measured catalog competes as before and decode interleaving is unaffected. A mixed shape may carry a solo chunk only when its prefill row is exactly the deterministic chunk graph.
Liveness fix (solo-request stall, reproduced and fixed)
The rule offers a single capacity while the candidate filter still drops an unmeasured prefill candidate whose chunk is cut back to a smaller graph's work (atomic span or message boundary + span straddling the full chunk's end). The single offered candidate was dropped and a solo request stalled forever after the Init graph. Fixed minimally: the rule bails whenever its chunk would not survive the candidate filter (no measured duration / unmeasured duplicate of a smaller graph's work / no work), leaving the measured catalog the choice, and the candidate site falls back to the full catalog if the rule's candidate is rejected anyway. Two regression tests reproduce the stall (6,000 tokens, span 2000..5000, costs 1 ms for 64..2048 and 0 for 4096) plus a message-boundary variant: both fail against the pre-fix rule ("solo request stalled: the chunk rule left no prefill candidate") and pass after the fix.
TTFT evidence (scheduler harness before/after; not a live serve)
Simulated first-token time and prefill launches, before (cost-based) / after (rule), per prompt arm, under a clean measured catalog and one where a slow capture over-measures the full graph 1.6x (true execution model unchanged):
Read honestly: TTFT is neutral on the long prompt when the catalog is faithful and 6.1 ms better when it drifts; on the short prompt the rule pads the tail to the covering graph and costs 4.0 ms more than the cost model's tuned tail splitting to buy one launch instead of three and a boot-invariant sequence. This is a scheduler-harness measurement; end-to-end serve TTFT was not re-measured.
Maintainer decision requested
Replacing measured-cost selection with a fixed rule for solo prefill chunking needs the maintainer's agreement: accept the determinism / tail-padding trade above, or prefer fixing the measurement drift the rule bypasses (then drop this PR). The shared change is validated only by scheduler unit tests — no per-model runs were performed (untested on every recipe, including Kimi K3/K3.1); the repo's cross-model requirement is flagged here for the maintainer.
Numerics class
NUMERICS-NEUTRAL at the arithmetic level (chunk selection only; per-chunk compute and the exact reductions are unchanged), with the disclosed caveat that different chunk boundaries change the BF16 KV/state rounding order relative to the old scheduler's choices (outputs become reproducible across boots).
Tests
cargo test --lib -- scheduler::: 278 passed, 0 failed, 2 ignored on the branch (base branch tip: 273 passed including the harness test, 0 failed, 2 ignored — no pre-existing failures on either side). The two stall regression tests fail against the pre-fix rule (reproduced before fixing).--nocapture).Dependencies: none on the companion FlashInfer PRs.