server: preempt a slot instead of ending every conversation when the KV pool fills - #184
Open
danielhanchen wants to merge 4 commits into
Open
server: preempt a slot instead of ending every conversation when the KV pool fills#184danielhanchen wants to merge 4 commits into
danielhanchen wants to merge 4 commits into
Conversation
…KV pool fills
With --parallel N --kv-unified there is one pool of cells and every slot believes it owns
all of them. When the pool fills, llama_decode returns 1, the retry ladder in decode()
halves n_batch down to 1, and the server calls send_error on EVERY processing slot:
"Context size has been exceeded". Four chats sharing a 8192-cell pool on Qwen3.5-4B-MTP
die together after six seconds, none of them anywhere near its own 8192 limit. The code
already says what should happen instead: "TODO: try to terminate only the largest active
slot/sequence and continue with the rest".
Terminate nothing. Once per update_slots(), before the batch is built, compare what the
pool holds against what the next decode will ask for. If it does not fit, take the cells
back from one slot: copy its sequence out with llama_state_seq_get_data_ext, release the
cells, and park the slot in a new SLOT_STATE_PREEMPTED. When the pool has room the copy
goes back with llama_state_seq_set_data_ext and the slot carries on. The task, the
sampler, the generated text and the position the stream has reached never left the slot,
so the continuation is the one the slot would have produced without the pause, and a
streaming client sees a gap and nothing else.
The check sits before the batch is built on purpose: at that point every slot is at a
token boundary, prompt.tokens is exactly what the cache holds for it, and no draft is in
flight, so a slot can be removed without unpicking a half-decoded batch. The speculative
draft is dropped with the cells, which costs the step its speedup and nothing else.
Victim policy: keep the slot that is furthest along, since it is the closest to finishing
and to giving its cells back, and among the rest prefer one that has not been preempted
three times already, then the smallest. A prompt cached on an idle slot is cheaper than a
conversation waiting to continue, so try_clear_idle_slots() is asked first, both before
preempting anyone and before deciding a resume does not fit.
Measured on Qwen3.5-4B-UD-Q4_K_XL with an embedded MTP head, --parallel 4 --kv-unified
-c 8192, four streaming chats with 1000-token prompts at temperature 0:
base 4 of 4 chats killed by "Context size has been exceeded" after 6.7 s
with this 4 of 4 chats completed, 0 errors, 7 preemptions, 7 resumes,
22944 tokens in 66.9 s (343 tok/s aggregate)
and the retry ladder never fires at all. At -c 16384 the same load still kills all four
on the base and still completes all four here.
LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens
regardless of pressure. With one request on an idle server the batch has the same shape at
every step, so it isolates the resume from batch nondeterminism: over 91 forced
preemptions across four prompts, every continuation is byte-identical to the same prompt
run without any.
Two tests on the two-slot unified pool. The first runs one request with LLAMA_SERVER_PREEMPT_EVERY=8 and asserts the tokens match the same request without the knob. The second runs two requests that each fit alone but not together and asserts both finish with no context error. Both fail on master: the knob is unknown there, and the second request dies with Context size has been exceeded.
…, and bound the parked state with --preempt-ram A slot still processing its prompt is between two chunks of it, which is as clean a boundary as between two sampled tokens, so it is a victim too: two prompts that do not fit together no longer fail together, and a large prompt arriving beside a running chat waits for it instead of ending it. A slot that has not started yet holds at most a cached prefix and is parked the same way, which is how it waits. Restoring takes the most-preempted parked slot first, but one that does not fit yet no longer holds up a smaller one that does: the smaller one is the first to be parked again if the pool fills, so the head of the line loses nothing. --preempt-ram N (LLAMA_ARG_PREEMPT_RAM) bounds the host RAM parked sequences may hold, default 8192 MiB like --cache-ram. A slot that would not fit under the budget is not parked, and when nothing can be parked the KV-full path runs as before. --preempt-ram 0 disables preemption. The prompt batching pass skips parked slots explicitly. Speculation is only restarted on restore for a slot that was generating; one parked mid-prompt starts it when its prompt is done, as it always did. Tests: two prompts that overflow the pool together, a generating slot beside a large prompt, and --preempt-ram 0 restoring the old behaviour.
|
You have reached your Codex usage limits for security reviews. Please try again later. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Counters n_preempt_total and n_resume_total, gauges requests_preempted and preempt_ram_bytes, and is_preempted plus n_preempt on each /slots entry, so a client can tell a parked request from a slow one and an operator can see the parked host RAM. A parked slot no longer counts as busy in n_busy_slots_per_decode, since it took no part in the decode.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 32c0a77e16
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // continue, so give those cells up first - same call the KV-full path makes. | ||
| for (;;) { | ||
| for (auto * slot : parked) { | ||
| if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { |
There was a problem hiding this comment.
When a parked request needs more than n_cells - PREEMPT_N_MARGIN cells, this predicate can never select it even after every other sequence and idle cache have been removed. For example, with -c 256 and a logical batch large enough, two valid 252-token prompts cause one empty SLOT_STATE_STARTED slot to be parked; after the leader finishes, preempt_n_need() returns 252 and 252 + 8 <= 256 remains false. Because preempt_restore() is never attempted, the restore-failure timeout also never runs, leaving that request parked indefinitely. The margin must not prevent a sequence that fits the actual cache from being restored, or the requested prompt chunk must be reduced accordingly.
Useful? React with 👍 / 👎.
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
1 participant
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.
Summary
With
--kv-unified --parallel N, every slot is told it has the whole context while all of them share one pool of cells. Four chats that each fit on their own are admitted together, grow into the pool, and when the next decode does not fitupdate_slots()enters the KV-full retry ladder and ends withsend_erroron every processing slot. Four users lose four conversations at once, none of them anywhere near their own limit.This change makes the server park a slot instead. When the cells the next decode will need do not fit, one running slot is chosen, its sequence is copied to host RAM with
llama_state_seq_get_data_ext, its cells are released, and its task stays alive in a newSLOT_STATE_PREEMPTEDstate. When the pool has room again the state is copied back withllama_state_seq_set_data_extand decoding continues from the same token, with the same sampler state. The client sees a pause in its stream and nothing else. The pause is invisible over the wire, so the OpenAI-compatible API, the built-in web UI and any third-party client all get it without changes.All of it is in
tools/server/server-context.cpp, insideserver_slotandupdate_slots(), tagged[TAG_PREEMPT].Policy
update_preemption()runs beforepre_decode()on every iteration ofupdate_slots(), only whenkv_unifiedis set and there is more than one slot.try_clear_idle_slots(). A running conversation is never asked to wait while a finished one is holding cells.PREEMPT_N_STARVED(3) times is passed over while any other candidate exists. The smallest remaining slot is parked, which frees the least work per pause.n_cmpl > 1) slots are not preempted. They share cells throughseq_cp, so a per-sequence save and restore would free less than it costs to put back./metricsgains the countersn_preempt_totalandn_resume_totaland the gaugesrequests_preemptedandpreempt_ram_bytes; each/slotsentry gainsis_preemptedandn_preempt. A client can tell a parked request from a slow one, and an operator can see the parked host RAM.--preempt-ram N(envLLAMA_ARG_PREEMPT_RAM) bounds the host RAM parked sequences may hold, default 8192 MiB like--cache-ram. A slot that would not fit under the budget is not parked; when nothing can be parked the KV-full path runs as before.--preempt-ram 0disables preemption.Why the smallest slot
A discrete-step simulation of the pool (
scripts/preempt_policy_sim.pyin the Unsloth workspace, costs taken from the measured runs) compares victim choices with everything else held equal: keep the leader, restore most-preempted first, fit-first, the same anti-starvation rule. With four chats the choice barely matters (within 1 percent of makespan). With eight chats on 8192 or 16384 cells parking the smallest slot gives the shortest makespan, the shortest mean completion, the least waiting, the fewest preemptions and the fewest cells copied; parking the largest is the worst on every count (5 to 6 percent longer, 40 to 50 percent more cells copied); parking the newest arrival, which is what vLLM does, is within 1 percent of smallest. Restoring from host RAM beats recomputing the sequence by 3 to 7 percent of makespan at these sizes, which is the argument for a server-side save over a client-side resume.Results
Qwen3.5-4B UD-Q4_K_XL with the embedded MTP head,
--parallel 4 --kv-unified --spec-type draft-mtp --spec-draft-n-max 2 --flash-attn on, four concurrent streaming chats with roughly 1000-token prompts, nomax_tokens, temperature 0, seed 1234.-cOn master the retry ladder halves
n_batchand three chats die on the speculative sub-batch index assertion (ggml-org#24840, fixed by #182). With #182 all four die together onContext size has been exceeded. With this branchfailed to find free space in the KV cacheis logged zero times: the ladder is never entered.The watermark fires at 8185 to 8192 wanted cells out of 8192, so it is not preempting early. Releasing 2000 to 4000 cells takes 105 to 420 ms; restoring takes 28 to 85 ms into an empty pool and 300 to 420 ms into a nearly full one. Aggregate throughput across the four chats was 353 tok/s against 279 tok/s for running the same four chats one after another, because the server still batches whenever the pool allows. The first cut restored strictly in priority order; letting a smaller parked slot through when the head does not fit took the run from 66.9 s to 50.5 s on the same load.
Exactness
Four concurrent greedy streams do not reproduce four solo greedy streams even without this change, because the batch shape differs and the matmul reductions are not shape-invariant. A chat that was never preempted diverged from its solo run after 50 characters. So a concurrent-versus-solo comparison cannot measure the pause.
The env var
LLAMA_SERVER_PREEMPT_EVERY=Npreempts a slot every N generated tokens regardless of pressure. With one request at a time the batch shape is identical with and without it, so the pause is the only difference:91 preemptions, byte-identical output, identical token counts, MTP drafting on throughout. Two further runs of the same check on prompts 0 and 2, one on the first cut (68 preemptions) and one on the final build (50 preemptions), also matched byte for byte. The save and restore round trip is exact and the sampler survives it.
A second model: Qwen3.6-35B-A3B UD-Q4_K_XL with its MTP head
Same flags and load, a mixture-of-experts model with the draft head active:
-cForced preemption every 200 tokens on prompts 0 and 2, 69 park and restore cycles: byte-identical to the unforced run (7059 and 6915 tokens). After the run
/metricsreportedn_preempt_total 9,n_resume_total 9,requests_preempted 0,preempt_ram_bytes 0.Cost when it does not fire
update_preemption()is a loop over the slots summingprompt.n_tokens()and returns immediately when the pool has room. One chat at a time on the same server flags, two prompts, two runs each, master against this branch:Tests
tools/server/tests/unit/test_preempt.py, six tests on the two-slot unified pool with the stories260K model:LLAMA_SERVER_PREEMPT_EVERY=8produces the same tokens as the same request without the knob, with at least six park and resume cycles in the log;--preempt-ram 0parks nothing and the requests fail the old way./metricsshows the preemptions and resumes, no request still parked and no parked RAM, and/slotsshows no slot parked.All pass on this branch with and without GPU offload, three runs each. The first two fail on master, the second with
Context size has been exceeded.Limitations
These are known and are the reason this is a first cut rather than the whole feature.
--preempt-ramnothing more is parked; there is no fall back to recompute yet. The parked state is about 36 KiB per token for this model.Context size has been exceeded. That is now the only way this server ends a conversation on pool pressure.Relation to #182 and #183
#182 fixes the speculative sub-batch index once the retry ladder is entered. #183 reduces the damage on a full pool from every slot to one slot. This change stops the pool from filling in the first place. All three are independent and merge cleanly; a prefill that cannot fit can still reach the ladder, so #182 is still needed.