Repository navigation
Conversation
…registry Canonical state recovery re-reads, densely and in the background, a prompt range that a SpecPrefill sparse prefill already served, so the prefix it leaves behind becomes reusable cache state. This commit adds the pieces it is rationed and gated by; nothing calls them yet. - CanonicalRecoveryBudget: a process-global share of wall time, accounted per window, with carried overshoot capped at one allowance. Slices are uninterruptible, so the share is a budget with slice-grain overshoot, not a strict ceiling. - An execution claim on the budget, so two engines sharing one GPU never run recovery slices at once. - canonical_recovery_is_runnable: every condition a slice needs, as one predicate - idle scheduler, no inbound or foreign foreground work, no SpecPrefill RoPE wrapper installed on the shared model, two idle steps. - A foreground arrival registry, so a request that has arrived but not yet been admitted still counts as foreground work. - PrefillProgressTracker.any_active(exclude_prefix=...), so recovery's own tracker entries do not read as foreign prefill. - Per-model settings, off by default.
A sparse SpecPrefill prefill serves the request but leaves no reusable prefix state: its cache is dropped rather than stored, so the next turn of the same conversation prefills the whole range again. Canonical state recovery queues that range as a background job and re-reads it densely, in slices, only while the scheduler is idle, publishing canonical state at each cache block boundary it reaches. - The job starts from whatever canonical prefix already exists and re-reads only the rest; it publishes only at block boundaries, and a read-back probe confirms each publication is independently restorable. - Every slice is gated by canonical_recovery_is_runnable. Foreground work retires the job's live state; the job resumes from its last published boundary. - Recovery never runs MTP and never primes MTP history, so blocks it publishes carry no Lightning-MTP sidecar; a request restoring from them starts MTP unprimed. That costs decode speed on that request, not correctness. - EngineCore announces a request before handing it to the executor, and abort_all_requests cancels recovery so an unload never waits on it. - EnginePool owns one process-global budget that every scheduler adopts; the admin API updates it in place. - BlockAwarePrefixCache.release_fetched_blocks returns the references a bare fetch_cache takes, so the read-back probe does not pin blocks. - is_specprefill_rope recognises every wrapper this module installs by identity, sharing one tuple with _unwrap_rope. Off by default: the global budget defaults to 0% and each model must opt in. With either unset no job is ever served.
…ical recovery Covers the serving-safety conditions (inbound requests, installed RoPE wrappers, idle steps, foreign engines), budget windows and the shared claim, lineage and stale-ground revalidation, state retirement and block reference release, MTP isolation, and the OOM requeue path. conftest also resets the foreground arrival registry and the prefill tracker around every test. Both are process-global; a tracker entry left by a test that drives a single prefill chunk otherwise reads as foreign foreground prefill to every later recovery predicate in the process.
Only tests/conftest.py conflicted: jundot#3933 added an autouse fixture next to the two this branch adds. Both are kept.
A recovery job recorded id(block_aware_cache) of the scheduler that served the request and refused to queue or publish when it no longer matched, with a weakref to release blocks from the old instance. That state cannot occur: block_aware_cache is assigned once in Scheduler.__init__ and is only ever set to None afterwards, and the job is queued and run on the same scheduler. The tests reached it only by hand-setting _serving_prefix_cache_id. The None case is still covered by _canonical_recovery_enabled() and the early return in _canonical_recovery_publish.
…s the layers The docstring said the active request id could not be trusted because of a race around sparse_prefill and cited jundot#766. There is no race: the wrapper is installed and the id set in one synchronous call, and every failure branch there calls cleanup_rope. The case that does need the layer walk is _requeue_or_fail_prefill, which clears the id and leaves the wrapper on the model. Say that instead.
The recovery code carried about 500 comment lines and long docstrings, many of them retelling how earlier versions failed or repeating the code in prose. Each is now a line or two on the non-obvious reason, or gone. No code changes.
A second TestIdleAccounting class shadowed the first, so the first class's tests never ran. Renamed the second one. The rest is import sorting, an unused import and an unused loop variable.
…yers _specprefill_rope_installed walked every attention layer because _requeue_or_fail_prefill was said to clear the id and leave the wrapper on. That cannot happen on main: the id is set only after a sparse prefill succeeds, while it is set _schedule_waiting defers every other request, and all three _requeue_or_fail_prefill call sites sit in the dense external/chunked prefill branch that a SpecPrefill request skips. The id and the wrapper are removed together everywhere else. With the wrapper unable to outlive the id, the recovery gate is the id check, and the blocked-step deadline that dropped a job stalled behind a leftover wrapper has no trigger left: an id that is set means a live SpecPrefill request, which is foreground work that ends. Drop the deadline, is_specprefill_rope and _ROPE_WRAPPERS, and the tests that built the leftover-wrapper state by hand.
…strings The test docstrings narrated earlier versions, a prototype review and old workarounds. Say what each test checks instead.
|
Field confirmation from a different machine class. oMLX 0.7.0, The "sparse state cannot be stored as an ordinary reusable prefix" cost is measurable here, and it dominates:
TTFT is 14× worse with SpecPrefill ON, and decode does not regress when it's off. Sparse prefill runs ~375 t/s here — below the dense rate — so it pays selection cost without reuse underneath. The static half of #3908 works; only the advancing conversation prefix fails:
Caveat so the numbers aren't over-read: #4138 merged after I started collecting, so unstable tool-schema ordering is likely a co-cause of some Two questions:
|
|
It's not in any release yet, only on this branch. To try it you need both Caveat: I've only run this on Qwen3.8-27B, not on Flash-Next, so your setup would be the first test on that model. On the second question, yes, I'd keep SpecPrefill off for that kind of session for now. Recovery only gives back the reuse the sparse turns lost, and only while the server is idle, so I don't know if on+recovery actually beats off on your workload. If you do try it, off / on / on+recovery with cached, TTFT and decode per turn would answer that. Also, blocks written by recovery don't carry MTP prompt history, so a request that restores from them starts MTP cold and acceptance is lower at the start of that response. I haven't measured how much that costs on your model. |
|
The only CI failure is test_cluster_pipeline_smoke_json_is_runnable, and it isn't from this PR. MLX's launcher stops reading a rank's output as soon as the rank exits, so a rank that prints its result and exits right away can lose that line; the launcher still exits 0. This PR doesn't touch the cluster code, and the test passes 3/3 locally on this branch. Fix sent to MLX in ml-explore/mlx#4615. Could you re-run the job? |
|
The recovery code is careful work, and with the feature off it doesn't change anything on the main paths. I'm still going to close it, and I want to explain why. Recovery redoes the skipped tokens as a normal dense prefill during idle time. So SpecPrefill plus recovery does more total work than just leaving SpecPrefill off, and on most turns the best it can do is match off. Your benchmark has no off run to compare against, and the field report above shows 3.1s median TTFT with SpecPrefill off vs 43.8s with it on. You also suggested keeping it off for that kind of session. On the other side the PR adds about 7,500 lines, around 1,400 of them in the runtime and mostly in scheduler.py, which gets changed almost every day. Realistically I can't keep maintaining that much added code going forward for something that only helps when an experimental feature is turned on. It's the same reasoning I used on #4353. |
|
Thanks for taking the time to review this and explain your concerns. I understand the maintenance burden, especially with how often the scheduler changes. I'm fine with closing the PR. Just wanted to add some context on why I started this. The main problem I was trying to solve was the slow first prefill on long contexts. With SpecPrefill off, the first request can take a long time, but later turns are fast because the KV cache is reusable. SpecPrefill helps with that initial wait, but the sparse state can't be reused as a normal prefix. In longer agent sessions, the uncached part keeps growing, and later turns can end up doing the same work over and over. My idea was to use idle GPU time to fill in that missing cache state. That way, we could potentially get a faster first response without losing the cache benefits on later turns. I wasn't trying to reduce total compute, just move some of that work out of the user's waiting time. The 228.38s → 79.06s result shows that recovery helps compared to SpecPrefill alone, but you're right that I didn't test against SpecPrefill off. That's an important missing baseline. The 3.1s vs 43.8s field result also doesn't include recovery, so we still don't know how the three compare. I'll keep experimenting with this on my own branch. If I can show a clear benefit over SpecPrefill off, I'll see if there's a simpler way to implement it without adding so much code to the scheduler. Thanks again for the review. |
SpecPrefill serves a request from a sparse subset of its prompt, but that sparse state cannot be stored as an ordinary reusable prefix. On an append-heavy session the reusable prefix stops advancing while the uncached suffix keeps growing, so later turns keep prefilling work the sparse turn did not leave reusable.
This PR adds canonical state recovery: while the process is idle, the scheduler densely re-reads that already-served range in bounded slices and publishes ordinary block-aligned prefix state through the existing cache path. Later requests restore it normally; there is no new cache format or restore path.
Off by default. The server-level budget
scheduler.canonical_state_recovery_global_budget_pctdefaults to 0, and each model must also setcanonical_state_recovery_enabled. With either unset, no job is ever served. The one thing that runs regardless is the foreground arrival registry, which records every request's arrival, admission and departure (three updates to a locked dict).Supersedes #3793, which grew to include a follow-up feature and a 490-line description; the code here is #3793's recovery without the adoption-grain commit.
How it stays safe
fetch_cacheresolves it under a throwaway id, and publication floors to a whole cache block. The read-back gives back the block references it takes (BlockAwarePrefixCache.release_fetched_blocks); without that, repeated probes pinned the very blocks recovery publishes.EnginePooland adopted by every scheduler, so M loaded models do not get M times the share. It is a wall-time budget with slice-grain overshoot, not an exact cap: a slice cannot be interrupted, so it is charged after it runs. An execution claim keeps two engines from recovering at once.fail_all_requests, so background work never blocks an unload.job.tokens[:boundary], so crossing lineages is unrepresentable rather than checked.step()'s error handling.Evidence
One M4 Max (64 GB),
Qwen3.8-27B-oQ4e-mtp(hybrid, 4,096-token cache blocks), greedy at temperature 0, MTP off, one client, seeded synthetic prompts, one run per arm. Single-run numbers.Recovery repays the debt. Seven-turn synthetic coding session, every turn forced onto the sparse route:
One controlled workload, not a general speedup. Both arms took the sparse route on every turn; what changed is how much of each prompt was already canonical.
Slice size bounds what an arriving request waits behind.
Different quantities; do not combine the columns.
canonical_state_recovery_slice_tokenssets it; the default is the ordinary prefill step. The budget controls how often a request collides with a slice; the slice controls how badly.The budget is process-global. Two models loaded, 10% cap: both engines report the same shared budget (owner count 2, identical window service to six decimals), measured aggregate share 10.36%. The excess is the slice-grain overshoot above.
Limitations
shutdowntimes out; registration grants nothing, so the cost is a stale count.Relationships
Review guide
Three commits, meant to be read in order:
feat(recovery): budget, execution claim, runnable predicate, arrival registry, per-model settings. No callers; 42 unit tests.feat(scheduler): scheduler integration, engine hooks, admin API,release_fetched_blocks. Runtime total: 13 files, +2,386/−38.test(recovery): 304 tests across 15 files, plus conftest resets for the arrival registry and prefill tracker.A follow-up that separates adopting recovered state from publishing it (
canonical_state_adoption_grain_blocks, default behaviour unchanged) will be proposed once this lands.Verification
On
upstream/main4c8f634f:pytest tests/ -m "not slow and not integration"(the CI command): 15,075 passed, 0 failed, 404 skipped, 86 deselectedpytest-xdist -n 8: 15,073 passed, 0 failed