Repository navigation
fix(scheduler): cap the admitted ceiling instead of refusing the request - #335
Merged
Merged
Conversation
WIP checkpoint. Justin hit 'scheduler did not admit the session generate request' on the ORT backend with CUDA, right after the shared GQA KV allocation, running a reasoning model for several turns. Root cause: PR #277 changed the CLI default max_new_tokens from 128 to the remaining context window. The scheduler reserves the worst-case footprint prompt + max_tokens up front, deliberately, so that an admitted sequence can never push the shared cross-session KV budget over its ceiling (DESIGN.md 26.4/26.11). That reservation was fine when the ceiling was 128 and becomes unsatisfiable when the ceiling is the whole window. The conservative reservation is kept because it is what guarantees an admitted sequence can run to completion. Instead the admitted ceiling is capped to what the budget can actually guarantee, so the request is served with a smaller bound rather than refused outright. Confirmed the native-vs-ORT asymmetry Justin observed: native generate bypasses scheduler admission entirely, so it never meets this gate. Both load paths do set bytes_per_token, so the difference is the code path, not the config. Also adds cancel cleanup so a failed admission cannot leak queued or running entries, and rewrites the error to name the cause, the numbers, and a next step. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
Justin hit "scheduler did not admit the session generate request" on the ORT backend with CUDA, running a reasoning model for several turns, right after the shared GQA KV allocation. Root cause: PR #277 changed the CLI default max_new_tokens from 128 to the remaining context window. The scheduler reserves the worst-case footprint prompt + max_tokens up front, deliberately, so an admitted sequence can never push the shared cross-session KV budget over its ceiling (DESIGN.md 26.4/26.11). That was fine when the ceiling was 128 and becomes unsatisfiable when the ceiling is the whole window. The conservative reservation is kept, because it is what guarantees an admitted sequence can run to completion. Instead the admitted ceiling is capped to what the budget can actually serve, so the request is served with a smaller bound rather than refused outright. The invariant holds: the ceiling was reduced to fit, not the reservation weakened. The cap is observable rather than silent, which was the harder half. It is reported on stderr, exposed as GenerateResult.budget_cap, and carried into the stats and profile output in both text and JSON. Callers can therefore distinguish "the model finished" from "we shortened your budget" -- which matters because this CLI is used for benchmarking, and an invisibly shortened token budget corrupts a measurement without ever looking wrong. No new FinishReason variant: budget_cap carries the distinction, which is proportionate and keeps the enum stable. Also confirms an asymmetry Justin observed: the native backend never hits this because native generate bypasses scheduler admission entirely, while ORT generation goes through it. Both load paths do set bytes_per_token, so the difference is the code path rather than configuration. Adds cancel cleanup so a failed admission cannot leak queued or running entries, and rewrites the admission error to name the cause, the numbers, and a next step, per RULES.md rule 1. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
.msg.txt was used to pass a long commit message to git commit -F and was committed by accident. Removes it and ignores that pattern so the mistake cannot repeat. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #335 +/- ##
==========================================
- Coverage 81.19% 80.48% -0.71%
==========================================
Files 308 308
Lines 118086 118446 +360
Branches 118086 118446 +360
==========================================
- Hits 95876 95334 -542
- Misses 18276 19178 +902
Partials 3934 3934
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
justinchuby
pushed a commit
that referenced
this pull request
Jul 28, 2026
Drops the stray .msg.txt scratch file, which leaked onto main from my own commit-message workflow and is removed there by #335. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a scheduler admission regression where requests could be refused after the CLI began defaulting max_new_tokens to the remaining context window. Instead of weakening the conservative KV-byte reservation invariant, the scheduler now caps the admitted generation ceiling to what the shared KV byte budget can guarantee, and makes that cap observable end-to-end (engine result + CLI stderr + stats/profile output).
Changes:
- Scheduler: add admission-time ceiling capping (
ScheduledBudgetCap), richer admission failure reasons, and cancellation cleanup for failed/aborted requests. - Engine: route ORT generation through the scheduler admission result API, propagate
budget_capviaGenerateResult, and ensure max token ceilings reflect scheduler admission. - CLI: surface the cap on stderr (stdout remains stable) and include it in stats/profile text + JSON; remove/ignore scratch commit-message files.
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| crates/onnx-genai/src/lib.rs | Re-export GenerationBudgetCap from the top-level crate API. |
| crates/onnx-genai-scheduler/src/lib.rs | Implement admitted-ceiling capping, admission error reporting, and request cancel cleanup; add tests. |
| crates/onnx-genai-engine/src/session.rs | Track optional budget_cap in active generation state. |
| crates/onnx-genai-engine/src/pipeline/nested_autoregressive.rs | Initialize GenerateResult.budget_cap in a pipeline result constructor. |
| crates/onnx-genai-engine/src/lib.rs | Re-export GenerationBudgetCap from the engine crate API. |
| crates/onnx-genai-engine/src/engine/runtime.rs | Use scheduler admission result API for ORT generation; propagate admitted max tokens + budget_cap; add a unit test for scheduler usage behavior. |
| crates/onnx-genai-engine/src/engine/mod.rs | Wire scheduler/engine types needed for cap propagation (ScheduledBudgetCap, GenerationBudgetCap). |
| crates/onnx-genai-engine/src/decode_loop.rs | Ensure decode-loop GenerateResult initializer sets budget_cap: None. |
| crates/onnx-genai-engine/src/config.rs | Define GenerationBudgetCap and add GenerateResult.budget_cap. |
| crates/onnx-genai-cli/src/profile.rs | Add budget-cap fields to stats/profile and expose them in JSON output; add tests. |
| crates/onnx-genai-cli/src/output.rs | Emit stderr notice when the scheduler caps max_new_tokens; propagate cap into profiling; add tests. |
| .squad/agents/batty/history.md | Update agent history notes for this change set. |
| .msg.txt | Remove stray scratch file from the repo. |
| .gitignore | Ignore common scratch commit-message files (.msg.txt, .pr-body.md, .git-msg.txt). |
Comment on lines
+368
to
+370
| if scheduled.request_id != request_id || scheduled.seq_id != session_id { | ||
| self.scheduler.cancel_request(scheduled.request_id); | ||
| self.scheduler.complete(session_id); |
Comment on lines
+340
to
+364
| Err(error) => { | ||
| let available_bytes = error.available; | ||
| let Some(admitted_max_tokens) = | ||
| self.max_tokens_for_available_bytes(request.prompt_tokens, error.available) | ||
| else { | ||
| return Err(self.admission_byte_error(request, requested_bytes, error)); | ||
| }; | ||
| admitted.max_tokens = admitted_max_tokens.min(request.max_tokens); | ||
| let capped_bytes = | ||
| self.estimated_bytes(admitted.prompt_tokens, admitted.max_tokens); | ||
| match self.try_reserve_bytes(capped_bytes) { | ||
| Ok(()) => Ok(( | ||
| admitted, | ||
| capped_bytes, | ||
| Some(ScheduledBudgetCap { | ||
| requested_max_tokens: request.max_tokens, | ||
| admitted_max_tokens: admitted_max_tokens.min(request.max_tokens), | ||
| requested_bytes, | ||
| admitted_bytes: capped_bytes, | ||
| available_bytes, | ||
| }), | ||
| )), | ||
| Err(error) => Err(self.admission_byte_error(request, requested_bytes, error)), | ||
| } | ||
| } |
…on race Three follow-ups from review of the capping fix. The cap was cleared on swap-in while the capped ceiling stayed in effect, so a sequence swapped out and back in generated under a shortened budget while reporting it was not capped. That is worse than never annotating it, because it is confidently wrong rather than silent -- and it lands exactly where it hurts, since a long benchmark run is the one most likely to be swapped and the one most in need of the annotation. The cap is a property of the admitted sequence, so it now travels with it. The capped retry read available bytes and then reserved, but ByteBudget is shared across schedulers by design, so another reservation could land in between and the capped retry could fail too. Adds ByteBudget::try_reserve_at_most, which computes the capped amount and reserves it under a single hold of the shared budget. An atomic reserve-adjust rather than a retry loop, because retrying against a contended budget trades a clear failure for a stall. The mismatched-admission branch cancelled the admitted request but left the original enqueued, where it could be admitted later unexpectedly, and called complete() where it should not have. Cleanup is now symmetric across every exit from that path. Each fix has a regression test, including budget_cap_survives_swap_out_and _swap_in. Engine lib tests now run clean locally against a conda ORT following #315. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
Resolves an additive conflict in output.rs: #325 added should_emit_stats_line and its test while this branch added budget_cap_notice and its test, in the same two places. Both are kept. Verified after resolution: build, 96 CLI lib tests, clippy, fmt. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
Resolves import-list conflicts in engine/mod.rs and lib.rs: #291 added the checkpoint and rewind types while this branch added GenerationBudgetCap. Both sets are kept. Verified after resolution: engine builds, 253 engine lib tests pass, scheduler tests pass, clippy and fmt clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 023d5a3f-c551-4aa8-afbe-a329af15fbdd
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.
The failure
Justin, running a reasoning model for several turns on the ORT backend with CUDA:
INFO onnx_genai_ort::session: allocating shared GQA KV on CUDA device memory device_id=0 error: scheduler did not admit the session generate requestRoot cause — a regression we shipped earlier today
#277 changed the CLI default
max_new_tokensfrom 128 to the remaining context window.The scheduler reserves the worst-case footprint
prompt + max_tokensup front. That is deliberate, not an oversight — it is what guarantees an admitted sequence can never push the shared cross-session KV budget over its ceiling (DESIGN.md §26.4/§26.11). Fine when the ceiling was 128; unsatisfiable once the ceiling is the whole window.Fix — cap the ceiling, keep the reservation conservative
The conservative reservation stays, because it is the thing that guarantees an admitted sequence can run to completion. Instead the admitted ceiling is capped to what the budget can actually serve, so the request is served with a smaller bound rather than refused.
The invariant holds: the ceiling was reduced to fit, not the reservation weakened.
The cap is observable — this was the harder half
A silent cap would have been a fifth entry in today's list of "the system quietly did something other than what was asked":
--temperatureignored,cudadocumented as silently falling back to CPU, context exhaustion writing an empty turn,--backendreporting the request rather than the resolution.It matters more here than elsewhere, because this CLI is used for benchmarking — an invisibly shortened token budget corrupts a measurement without ever looking wrong.
So the cap is reported:
GenerateResult.budget_capCallers can distinguish "the model finished" from "we shortened your budget". No new
FinishReasonvariant —budget_capcarries the distinction, which is proportionate and keeps the enum stable.Confirms the asymmetry Justin observed
Native
generatebypasses scheduler admission entirely; ORT generation goes through it. Both load paths do setbytes_per_token, so the difference is the code path, not configuration. Worth knowing independently of this fix.Also
--vram-limit, or an explicit--max-new-tokens). The old message said only that admission failed, and cost real diagnosis time today..msg.txtthat leaked into the tree from my own scratch workflow, and ignores the pattern.Gates
Scheduler tests, CLI lib tests, build, fmt, clippy all pass. Engine lib shows the known local ORT mismatch baseline (178 passed / 64 failed), unchanged.