server: preemption notices, asynchronous parks and exact concurrency together - #197
server: preemption notices, asynchronous parks and exact concurrency together#197danielhanchen wants to merge 503 commits into
Conversation
|
Merged #194's c6c3cb6 as ef257d7. It fixes the speculative cost the Cost section measured: under exact mode the batch splitter treated each slot's three-token verify batch as a prompt and gave it a ubatch of its own, so every decode step ran the graph once per slot. Prompts are now isolated by width and the column split runs in batch-of-one-equivalent slices. Two MTP drafts, four chats, the same full-pool shape: 0.74 of the mode-off aggregate against 0.45 before, byte-identical to the solo run on the 4B and the 35B with and without parks. The 16 server tests pass on the merged binary. The Cost section is updated. |
|
Merged #194's 379ca5d as abccc6e: on mixture-of-experts models the exact mode no longer re-enters every expert projection once per token, so the 35B decodes at the mode-off rate within noise with and without MTP drafts, byte-identical to the solo run with and without parks. The 16 server tests pass on the merged binary. |
|
Merged the second review round: #184's 2a7e277 as e8ab02c (a parked sequence that cannot fit the pool alone is failed with the context error instead of parked for ever, and the scheduling margin is waived when nothing is resident) and #194's adeed65 as db95586 (head widths other than 256, tokens carrying several sequence ids, and hybrid cross-sequence edits are refused by name under exact mode). The merged binary passes the 17 server tests, including the new near-capacity test, and the 35B GUI drive completed 4 of 4 chats with 9 parks and no errors. |
…mption-integration
…e-preemption-integration
…e-preemption-integration
|
Merged #184's 84311fd and 6dbc4e7 as a7e661c: when the KV-full retry ladder runs out, resident slots are rewound to the cache boundary and the smallest parked instead of every slot getting the context error, and a batch holding a draft is never narrowed through the group. Suite 13 + 4 + 2 passed. Exact-mode accuracy with the planner off (last resort only), 4B and 35B, with and without MTP drafts: identical to solo in every round, 0 context errors. |
…e-preemption-integration
The last resort parks through the same call as the planner but did not send the stream comment the planner sends, so a client that shows the pause from that comment showed nothing for a slot parked this way.
…de-preemption-integration
|
Merged #184 86845c1 (the draft shortcut only when there is a budget to park into) through #192 47ef8de and #194 6573242 as 1130aed, and #190 5864dae as b89302b: a park made as a last resort now sends the stream comment the planner sends. Found by a 35B GUI drive with the planner off: before, four of four chats finished with nine resumes and no pause shown; after, four of four with the pause shown on every chat, 12 parks and 12 resumes, 0 errors. Suites 13 + 4 + 2 passed. |
…e-preemption-integration
A slot on its way back in already holds its cells and starts decoding the moment its copy lands, so it needs the same runway as the slots already running. Leaving it out let two back-to-back restores land into a pool that then had to park someone again at once.
…precated the flash attention setter
…t the flag behaves as upstream
…er exact concurrency
…est-exact-pages, the OpenVINO support probe's struct return, and park assertions that match the park log lines rather than the slot JSON
…d-slot follower behind a slot that is seen busy, and give the notify resident time to decode on a loaded runner
… disconnect from raising test_a_resident_cycling_through_context_shifts_is_rotated_out_for_a_parked_head ran into the 600 s default request timeout on the ubuntu Server leg: the server was healthy and the third request came back with all 9000 tokens, the other two were still generating when the client gave up. The rotation is triggered by a 2 s wait, not by a token count, so the generation cannot be shortened without losing the rotation on a fast host; the requests get a longer timeout instead. test_a_resident_that_cannot_be_swapped_out_is_rotated_by_recompute stops the server with a request still in flight on purpose. Its thread raised the resulting ConnectionError as an unhandled thread exception, which is the noise that made the timeout above look like a crash.
…st's prompt stats A recompute restore re-enters prompt processing to put back the cells the park dropped. Those tokens are submitted with is_prompt set, so every one of them was added to n_prompt_processed and pushed t_prompt_last to the end of the replay, while n_gen deliberately carries across the park. The reported prompt length grew with every park, and the generation time covered only the tokens after the last re-prefill, so predicted_per_second was inflated by the ratio of the two. The replay is still counted in the server-wide prompt metrics, where it is real compute; it just no longer moves the slot's prompt count or the prompt/generation boundary.
… streamed /v1/responses A non-streamed /v1/responses carries preempt in the response object, next to usage. The streamed one wrote it on the SSE data beside the response object, so a client that keeps the response of the response.completed event, which is the object the OpenAI SDK hands back, never saw it. It now sits in the same place either way, and unconditionally, as the non-streamed body already did.
…two requests overlapping
|
@codex review |
Stacked on #184. This branch merges #190 (park notices on the stream), #192 (park and restore copies off the decode loop) and #194 (opt-in exact concurrency) so they can be tested together, plus the fixes that only showed once they shared a binary. It is rebuilt from the three branches' heads as they move.
Correction, 2026-09-07. The sentence that used to stand here said every fix found on this branch had been carried back to the PR it belongs to. That is no longer true and it should not have been left standing.
git log --no-mergesagainst the three merged heads returns 16 commits unique to this branch, nine of them from 07 September: f8fdf06, 951a219, 57efd21, 7e424ad, 7549c68, 4cf9baf, 00855c2, bb510f5, 09cecf4. They share the park and resume bookkeeping between the copy paths, count the KV reserve in one pass, build the async state io on the host io, and share the preemption test setup. Each one cuts across #190, #192 and #194 together, so they cannot be split back into those branches cleanly.That makes this branch the merge target for the three, not a scratch integration of them. #190, #192 and #194 stay open as the reviewable layers, and
gh api compareconfirms this branch contains all three in full (#190 and #192 report ahead=0; #194 differs only by a merge commit with no unique content). It does not contain #186, which still holds four commits of its own including 296 lines oftest_preempt_hybrid.py, so the preemption stack currently has two live heads and #186 also needs another merge from #184, which it is four commits behind.Rebuilt over
75944db80from the current heads:1c5aa4904feat/server-side-preemption-notify(#190,5a791e03c)51eca0837feat/server-side-preemption-async(#192,888603d03)70d49bfacfeat/exact-concurrency(#194,da8556d31)ef257d7b8feat/exact-concurrency(#194,c6c3cb671): speculative verify batches stay grouped, sliced column splitabccc6e95feat/exact-concurrency(#194,379ca5d42): one launch per expert projection under the knob instead of one per tokena7db6262afeat/server-side-preemption(#184,50b617ae4): parked slots come back head of line by park time3f817d880feat/server-side-preemption(#184,8057a74af) andfeat/exact-concurrency(#194,f4e45646d): resume-order knob logged at load; equal-width grouping under exact modee8ab02c94feat/server-side-preemption(#184,2a7e277ab): a parked sequence that cannot fit the pool alone is failed instead of parked for ever; margin waived when nothing is residentdb9558696feat/exact-concurrency(#194,adeed654b): head widths other than 256, tokens with several sequence ids, and hybrid cross-sequence edits refused by name under exact mode70b874c33feat/server-side-preemption-async(#192,01dfecc5e, carrying #18484311fd3dand6dbc4e7b6): when the KV-full retry ladder runs out, resident slots are rewound to the cache boundary and the smallest parked instead of every slot getting the context error; a batch holding a draft is never narrowed through the groupa7e661c14feat/exact-concurrency(#194,c89d87db5): the same, with the page-rounded margin1130aed0b47ef8de54and #1946573242a1(#18486845c15e): a batch holding a draft is narrowed the old way when there is no budget to park intob89302b25feat/server-side-preemption-notify(#190,5864dae95): a park made as a last resort sends the stream comment the planner sends, so Studio shows the pausef8ec5665afeat/server-side-preemption-async(#192,cd1cd4e6d): the lookahead margin counts a slot being restored5e0f99aedfeat/exact-concurrency(#194,0c9fc0ed8): every context reports the widest decode step it can build, so the column bound covers direct library users3fc833954301f480c0, #1920303888e8, #19408e10dabb(#184662ec2029): the planner and the last resort step aside for a context without memory7c8a6d51576baf1a68, #192cac5ee6c0, #19467063e476(#184a9b712eee): unlimited budget enables the last resort; per-slot draft bound for restores and reservations; shared prompt charged once per family; reused slot charged its prefix; context shift before the planner measures800e166ad7fb42b582and #19446e7fa742: parked prompts of a multi-prompt stream kept apart; ngram-cache verify width; DFlash refused by name; an explicit bound below the decode width refused at context creation; width replayed to backendseb428268c53c59a109, #1928d74efbb2, #194b7ace528f(#18464a5064ef): a resident cycling through context shifts takes turns with a parked head9391aa71e356978a26, #19242e536970, #194dbd82ca5e(#1846744b3d9f): a non-causal context with a cache refused at creation under exact mode; one page boundary per waiting prompt slot in the reserve cap; a stream parked before its first token starts with the notice; the rotation park announced; a reused slot trimmed to its shared prefix before it is sized or parked; a round with a context shift waits for every copy in flight; the decode width of every context follows the token figure645d406dd3306d4b50, #192ad3856ae9, #194d4e3fc8a6(#184a1c34dad2), on top of6af0d4e99: a reused slot's need counts from its shared prefix and its trim is safe for memories without partial removal; a waiting child is charged its own cache; equal-count grouping stays on for recurrent and hybrid memories; a width the explicit column bound cannot cover is refused; a restore landing inside the context-shift wait sends the resume notice276dd25296fb0b91c5:LLAMA_SERVER_PREEMPT_POLICY, a test knob to compare victim choices on one workload (off unless set)ecdb0f5317efef20fa, #1922ba0dcf5f, #194bf00ac38a(#184a7a04c2a9): the rotation parks the resident that lets the head in (page-rounded here); soft-capped attention refused at load; width reports serialised and monotonic in the backend; the token figure never lowered089fedd923d179fd0b, #19202a3e11d4, #19498fe86dfc(#18464a3f6e7d): a parked slot, or one whose copy is in flight, survives an aborted round; a rotation counts the head's bytes as leaving; the exact setup runs before any context exists; one lock for the token figure, the sequence count and the widthd72e2a40da06419d00, #192e432b15c2, #19449c66c5bd(#18455f04bbd6): the leader measured by what a reused slot keeps; no rotation while a park is copying and an asynchronous rotation park re-examined when it lands; the rotation budget under asynchronous copies and buffers returned over budget; the cache-reuse shift waits for copies in flight; an isolated ubatch takes only sets that finish in it; page-aware margin at every step testb4a5f18d82b4a69122, #1942f2258dc0: a whole-context restore under exact mode is refused before it can clear the cache; the width is published once construction succeeds; the page table is refused on every backend that ignores it; DSpark refused with DFlash; a transfer that fails part way posts no copies; pageable host memory parks synchronously; no transfer for state that is not on a device; the runway rounding explained48f7d7c0b1d528856f, #192faa3dd3f7, #19423c9dc2ef(#184270fdd6f3): test that a parent and child alone in the pool get the context error and the server lives0611fd99ba4b62f2be, #194918a8bf4f: a park buffer that comes back pageable parks synchronously from then on; the graphs after a restore wait for its copies on the device; the exact-mode cache asks the device whether it runs the paged attention for the layer rather than trusting the backend's namef737e6d49294d2a912: idle parked RAM is given back when another slot needs to park, so a budget that holds one sequence is not spent for good by the first restore; the copies of a park or restore wait for the compute stream on the device instead of draining the hostbf0f731273800ddee9through #190a158200deand #192dffa102e3: a pure recurrent cache is served without preemption; a rotation holds both states at once, so the RAM cap is asked for the resident's state in full and a budget that holds one sequence but not two does not rotate, said once per parka140e80c881eec0bb3through #190ce02834c4and #192a5749d11b: the copies wait for a fence the context records after every decode, so restores issued in one pass run independently on the device; a started slot's reservation counts from the prefix it keeps (page-rounded here, the step starting from that prefix)e36a4b6ddaf560907bthrough #1905ecdba815and #192a56d49e1f: the context shift and the planner run inside the guarded part of the step, so a failure there ends the affected slots rather than the loopb0fae413a00b27d2df, #1904918294ff, #1928cb58dfdfplus 4b91744 here: what a started slot keeps is decided by the batch builder's rule (no caching, aLoRA cutoff); a parked stream keeps a shorter ping it asked for; staging counted by what the buffer charges; fences after the layout check; transfers only where a park can happen; a decode that fits goes ahead beside a park in flight instead of waiting for it; a slot restoring into the prompt phase counts in the page-boundary cap0d39b88baad89538d2, #190801e29ad1, #192114e230ab: the resume order read on every load; the shift wait runs before the draft is asked for, so the draft cache shifts after the copies have landed; the last transfer takes the fences with it7fdb58bb3301caa31d, #190929809ffa, #192b1a31b42d: the resume order belongs to the context and the recurrent flag is assigned on every load; the shift wait runs before the decode as well as before the draft; a context freed with live transfers drains and disowns them1bf8519163c99fafdf, #190a33c56397, #19258c6e3329: the pool is measured by the cells each slot physically holds (page-rounded here), and when nothing fits every started slot is trimmed to the prefix its request keeps before idle slots are cleared, so a resume is not attempted against cells the batch builder has not yet released474569ca390a5094d3and #194b4b0f9bd7: a started slot's prompt is validated before it can be parked, so a request the STARTED block would reject is never given a notice ahead of its non-stream error; shorter comments from #194, the conflicting hunks keeping this branch's textf8fdf063a951a21931tobb510f556preempt_park()serves the four park sites and one pair of helpers the three copy completions, the kv reserve is counted in one pass, the two test files share their server setup and stream readers (29 tests and every assertion unchanged), the async host io derives from the sync one, one-reader wrappers and the target-and-draft transfer repeats go. 1050 lines fewer. Verified at each commit by the suite, at the end by the seq-copy and fragmented-restore ctests on the 4B, the exact-mode soak identical 6 of 6 with forced parks, and test-backend-ops under the mode09cecf4f0The four fixes below now come in with those merges. Two conflicts had to be resolved by intent rather than by text: at the KV-full park site the notice and the new host-memory report both belong, in the order the forced-park site already uses; and #194's synchronous planner meets this branch's asynchronous one, where the resolution keeps the asynchronous margin and its
RESTORINGreservation and takes #194's compile-time-checked rounding helpers, so the two planners share one set of arithmetic and the synchronous branch of the margin rounds as well.Summary
Three merges and four fixes. Merging #190 into #192 needed intent rather than text: git placed two of the three notice sites after the asynchronous early-outs, where an asynchronous park or restore would have announced nothing. The rule applied is that a park is announced when the slot stops taking part in a decode, on entering the copy-out state, and a restore when it starts again, when the copy back lands; the keepalive therefore spans all three in-flight states. #194's page allocator and #192's park compose without change: a park removes whole pages, both in-flight states hold live cells, and a restore allocates through the page branch of the slot finder.
Fixes found here
llama_memory_i::alloc_granularity()now reports the page size and the server rounds its four planner figures by it. After: 4 of 4, two parks per round, the parked chat's bytes unchanged.nabove one reached a cache assert under exact mode and killed the server. It is refused with a 400 at task build, and the cache logs and returns instead of asserting.--sse-ping.Results
Qwen3.5-4B, two MTP drafts,
-c 8192, four slots, seed 0, temperature 0, forced parks every 64 tokens, all three features on, re-run on the rebuilt branch: 4 of 4 in three rounds, 28 parks and 28 restores per round, seven park and seven resume notices on every stream, and the observed chat byte-identical to its solo run in every round, at the same sha2563922db5817b706aa...the branch gave before the rebuild. Survivor stall at a park 52 to 67 ms and at a restore 39 to 85 ms with exact mode on, against 76 to 78 and 57 to 64 with it off, on a shared GPU.Qwen3.6-35B-A3B through Unsloth Studio on an unshared GPU: exact mode loads with Studio's own launch args; a chat's bytes are unchanged by three different neighbours across three rounds and by seven forced parks per stream, while the same cell with the mode off diverges at byte 488; four API chats and eight GUI chats in Chromium and WebKit all completed with zero errors and every one of 21 natural parks lasting up to 96 s resumed with the paused label shown.
Cost
Re-measured on this tree, interleaved: three on/off pairs, four chats, prompts of 937 tokens as the server counts them, 2048 tokens each with
ignore_eos,-c 8192, four slots, unforced, medians of three.70d49bfac70d49bfacef257d7b8ef257d7b8All six runs finished 4 of 4 with parks and restores and no context errors. The 37 to 42 percent this paragraph used to quote was measured on #194's reviewed head and is gone with speculation off: #194's own fixes stopped one prompt serialising every concurrent decode for the whole prefill and dropped the default column bound from a fixed 16 to the width a decode step can reach. What is left there is the CUDA column policy rather than the page bookkeeping, which #194 measures at 0.79 for the kernel half alone against 0.92 for the whole mode.
The speculative cost at
70d49bfacwas the batch splitter, not the column policy: under exact mode it isolated every sequence set with more than one token left as a prompt, and a slot's three-token verify batch is such a set, so each decode step ran the whole graph once per slot. #194'sc6c3cb671isolates by width (llama_set_exact_decode_tokens, one plus the draft length) and slices the column split into batch-of-one-equivalent launches, which is the 0.45 to 0.74 in the table, measured on the #194 tree and carried here byef257d7b8. What remains is the sliced column policy on a twelve-wide verify batch plus paged attention. One B200 shared with another tenant throughout, so the ratios are the result and the absolute figures are not.Tests
Server harness 31 of 31 (27 preemption, 4 notify, 22 parametrised cases), the fragmented-restore and sequence-copy ctests on CUDA with the 4B, and
test-backend-ops5234 of 5234 in exact mode.Follow-up: seven fixes and fewer tests
Nine fixes, each checked against the previous head.
llama_memory_update()applies the recorded shift on the target and the draft before the planner runs, so every park site sees applied state; the reviewer's harness goes from 26621 of 71864 checkpoint bytes differing to 0, and the token-identity test passes./slotsreportsis_preemptedonly for a slot whose cells are gone and a newis_transferringfor a copy in flight, so a reader counting residency keeps counting a transferring slot.--preempt-ramcannot hold a victim, the park drops the victim's cells and keeps its tokens, sampler and counters, and the resume re-prefills the tokens through the ordinary prompt path and carries on from the already-sampled token. Both the planner and the last resort use it, so the stock all-slot context error is unreachable while another slot can be parked. Two 2000-token streams under--preempt-ram 1go from both erroring to both finishing; a forced park at 1800 tokens with 1819 re-prefilled produces output bit-identical to the unparked run.-c 400 -b 64go fromfailed to process mtmd chunkto both completing after a park.Fewer tests. The two server test files go from 874 lines and 29 tests to 731 lines and 31 tests (22 parametrised cases plus the nine tests above), with the setup, prompts and stream readers shared and the near-duplicates folded into parametrisations. The survivors were checked against four scratch mutations (no victim ever picked, the pre-fix cell count, the started-slot validation gate removed, the park notice removed); each was caught by the intended test. Two behaviours that had no test before are written rather than kept: the physical cell count of a started slot, and the parked keepalive.
One thing found on the way, not of this PR: a 1900-token prompt processed in 32-token batches after a 2000-token decode hangs the CUDA build with graphs enabled, on master as on this branch, with
GGML_CUDA_DISABLE_GRAPHS=1or a larger batch it completes. The recompute test re-prefills in one batch for that reason.Verified on the merged tree: the two files 31 of 31 on CUDA with the recurrent, hybrid and M-RoPE models supplied and nothing skipped; the upstream completion, context-shift, vision, basic, stream and kv-keep files 73 passed and 1 skipped with the one pre-existing failure, a wording regex on the tiny model that fails on master too; the CPU-only build serves the recompute park and re-prefills 1910 tokens in 32-token batches; the exact-mode soak on Qwen3.5-4B with two MTP drafts and forced parks every 64 steps identical in 6 of 6 rounds against the solo run at 295 to 300 tok/s aggregate, 164 parks and 182 resumes with zero asynchronous markers and no context error. Live through Unsloth Studio: four GUI tool chats at 8k finished 4 of 4 with 9 parks and 9 resumes, and
unsloth runat 16k with two MTP drafts and exact concurrency finished 4 of 4 at 191 tok/s aggregate with 5 parks and 5 resumes, no context error in either.llama_decodeas 1, the code callers retry on, when it can never succeed. It is rejected during batch validation with -1 now, before the memory is asked.Round 23
/propscarriesexact_concurrency--preempt-ramto the poolRound 24: audit fixes and the exact-concurrency follow-ups
Six fixes from an isolated review of this head, each with a test, plus the three code fixes #194 gained since this branch took it (cherry-picked: the weight placement check at context creation, the allocation-granularity comment, and per-page live counts so a removal that empties nothing leaves the page index alone, with
tests/test-exact-pages.cpp).keep_firstdropped or kept the wrong tokensrewind_to_cacheconverts withsize_up_to_pos, steps back to the start of a media chunk it would cut, and asserts the result;tests/test-server-tokens.cppbuilds a 256-token, 16-position image chunk and checks the conversionLLAMA_SERVER_PREEMPT_FAIL_SAVE=Nfails the Nth save for the test)n_ubatchchunks unless the rest fits, a chunk a neighbour would cut short is deferred, and a batch too small to hold a decode step plus one ubatch warns once; measured identical ubatch grids alone and beside three decoders: recomputedfollows the: resumedof such a resume,llamacpp:preempt_recompute_totalin/metrics,n_recomputeper/slotsentry, andpreempt: {parks, recomputes}on every final completion object; documented in the server READMEmin(params.n_max, dp.n_max); with one token of allowance left the unfixed build logged the truncation, the fixed build none, and a 4-token draft got three times faster/propsstill reporting the modecommon_exact_batch_geometryrefuses the geometry at init, naming the batch size that would work (-b 514for a 512-token ubatch and two slots); the warn path is gone, a prompt that does not fit whole waits for the next iteration;tests/test-exact-geometryand a server test--preempt-asynchelp said the slots that keep running do not wait for a park: preemptedat park time on every streaming surface (320 ms before the resume in the test scenario); now stated in the README and pinned by a test parametrised over/completion, chat, responses and messages that the notice arrives before the first data linesse_ping_interval; ten consecutive runs passVerified on one GPU with Qwen3.5-4B:
test_preempt.pyandtest_preempt_notify.py39 passed, 5 skipped;test_metrics.py10 passed;test-server-tokens,test-exact-pages(also underLLAMA_KV_CACHE_DEBUG=1) andtest-state-seq-copypass;-ot attn_q=CPUunderLLAMA_EXACT_CONCURRENCY=1is refused namingblk.0.attn_qkv.weightinCUDA_Host; four exact-mode sequences of 256 greedy tokens with a forced park every 64 are identical to their solo references (12 parks, no recompute). The exact and MTP tests skip unlessLLAMA_SERVER_TEST_EXACT_MODELandLLAMA_SERVER_TEST_MTP_MODELname a model. Not covered end to end: the multimodal rewind, since the only local VLM has media positions equal to its token counts.