From 6009998a74ba95a6a7edcd8f6c61ed77104e6c8a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 00:12:09 +0000 Subject: [PATCH 01/29] server: preempt a slot instead of ending every conversation when the 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. --- tools/server/server-context.cpp | 401 ++++++++++++++++++++++++++++++++ 1 file changed, 401 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..85e9494563d1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -59,8 +59,27 @@ enum slot_state { SLOT_STATE_PROCESSING_PROMPT, SLOT_STATE_DONE_PROMPT, SLOT_STATE_GENERATING, + SLOT_STATE_PREEMPTED, // [TAG_PREEMPT] cells released, everything needed to resume is in host RAM }; +// [TAG_PREEMPT] server-side request preemption +// +// With --kv-unified the cells are one pool shared by every slot, and each slot believes it +// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves +// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size +// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the +// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and +// continue with the rest". +// +// Nothing is terminated here. The cells of one slot are taken back and given to it again +// later: its sequence is copied to host RAM, its cells are released, and when the pool has +// room the copy goes back and the slot carries on with the same sampler, the same generated +// text and the same open stream. A streaming client sees a pause, not an error. +constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation +constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected +constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on +constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked + struct server_slot; // forward declaration struct server_batch { @@ -293,6 +312,115 @@ struct server_slot { prompt.clear(); } + // [TAG_PREEMPT] state of a slot whose cells were taken back + // + // Only the KV cells leave. The task, the sampler, the generated text and the position + // the stream has reached stay on the slot, so a resume is a memcpy and not a new + // request: no retokenisation, no replayed prompt, no seam in the output. + slot_state state_before_preempt = SLOT_STATE_IDLE; + std::vector preempt_state_tgt; + std::vector preempt_state_dft; + int32_t n_preempt = 0; // times the CURRENT task has been preempted + int32_t n_preempt_fail = 0; // consecutive failed restores + int64_t t_preempt_us = 0; // when it was parked + + size_t preempt_state_size() const { + return preempt_state_tgt.size() + preempt_state_dft.size(); + } + + void preempt_state_free() { + preempt_state_tgt.clear(); + preempt_state_tgt.shrink_to_fit(); + preempt_state_dft.clear(); + preempt_state_dft.shrink_to_fit(); + } + + // copy the sequence out of the cache and release its cells + bool preempt_save() { + const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + + try { + preempt_state_tgt.resize(size_tgt); + preempt_state_dft.resize(size_dft); + } catch (const std::bad_alloc & e) { + SLT_ERR(*this, "failed to allocate %.3f MiB for the preemption state: %s\n", + (size_tgt + size_dft) / (1024.0 * 1024.0), e.what()); + preempt_state_free(); + return false; + } + + if (llama_state_seq_get_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to copy the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_get_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to copy the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + // The draft is a prediction, not a result, so it goes with the cells. Preemption + // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already + // holds exactly the tokens the state above covers -- including the rollback done by + // the checkpoint path when a draft was only partially accepted. + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + + // note: prompt.tokens is deliberately kept. It is the mirror of the state just + // copied out, and the resume needs it to know how many cells to ask for. + mem.seq_rm(id, -1, -1); + + state_before_preempt = state; + state = SLOT_STATE_PREEMPTED; + t_preempt_us = ggml_time_us(); + + n_preempt++; + + return true; + } + + // put the sequence back; the slot then continues from the token it was about to decode + bool preempt_restore() { + const size_t size_tgt = preempt_state_tgt.size(); + const size_t size_dft = preempt_state_dft.size(); + + if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + // no room after all: drop the half-written sequence and stay parked + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + if (size_dft > 0 && + llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + preempt_state_free(); + + n_preempt_fail = 0; + + state = state_before_preempt; + + // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks + // that the draft context is where it should be, which the restore above ensures + if (can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + + return true; + } + std::vector lora; int32_t alora_invocation_start = -1; @@ -351,6 +479,13 @@ struct server_slot { n_predict_max = -1; + // [TAG_PREEMPT] + preempt_state_free(); + state_before_preempt = SLOT_STATE_IDLE; + n_preempt = 0; + n_preempt_fail = 0; + t_preempt_us = 0; + llama_set_sampler(ctx_tgt, id, nullptr); // clear alora start @@ -503,6 +638,14 @@ struct server_slot { t_last_used = ggml_time_us(); + // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be + // released while parked), so the mirror of them must not outlive them: the next + // task on this slot would otherwise take a prefix match against an empty cache + if (state == SLOT_STATE_PREEMPTED) { + preempt_state_free(); + prompt_clear(); + } + state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough @@ -1249,6 +1392,16 @@ struct server_context_impl { } } + { + const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); + preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; + + if (preempt_test_every > 0) { + SRV_WRN("LLAMA_SERVER_PREEMPT_EVERY = %d (test knob: preempting every slot every %d tokens)\n", + preempt_test_every, preempt_test_every); + } + } + { const char * LLAMA_SERVER_SLOTS_N_DIFF = getenv("LLAMA_SERVER_SLOTS_N_DIFF"); slots_n_diff = LLAMA_SERVER_SLOTS_N_DIFF ? atoi(LLAMA_SERVER_SLOTS_N_DIFF) : 0; @@ -2674,6 +2827,251 @@ struct server_context_impl { }; #endif + // + // [TAG_PREEMPT] server-side request preemption + // + + int64_t n_preempt_total = 0; + int64_t n_resume_total = 0; + + // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, + // whether or not the pool is under pressure. It exists to answer the only question that + // matters about a resume: with one request on an idle server the batch has the same + // shape at every step, so a preempted continuation that is not byte-identical to an + // uninterrupted one is the preemption's fault and nothing else's. + int32_t preempt_test_every = 0; + + int32_t preempt_n_spec_max() const { + return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; + } + + // Cells the pool is holding right now. A released slot keeps its prompt in the cache + // for the next request to reuse as a prefix, so idle slots count too: the first version + // of this counted only the running ones, decided a pool holding 8185 cached cells was + // empty, and every resume failed against a cache that was actually full. + int32_t preempt_kv_used() const { + int32_t res = 0; + + for (const auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTED) { + continue; // parked: its cells are in host RAM, not in the pool + } + + res += slot.prompt.n_tokens(); + } + + return res; + } + + // cells those slots are about to ask for on the next decode + int32_t preempt_kv_reserve() const { + const int32_t n_spec = preempt_n_spec_max(); + const int32_t n_batch = llama_n_batch(ctx_tgt); + + int32_t res = 0; + int32_t res_pmt = 0; + + for (const auto & slot : slots) { + switch (slot.state) { + case SLOT_STATE_GENERATING: + case SLOT_STATE_DONE_PROMPT: + { + res += 1 + n_spec; + } break; + case SLOT_STATE_STARTED: + case SLOT_STATE_PROCESSING_PROMPT: + { + const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + + res_pmt += std::max(1, std::min(n_batch, n_left)); + } break; + default: + break; + } + } + + // one batch is all the prompt slots get between them, however many are waiting + return res + std::min(res_pmt, n_batch); + } + + // Keep the slot that is furthest along -- it is the closest to finishing and to giving + // its cells back -- and among the rest prefer one that has not been preempted + // PREEMPT_N_STARVED times already, then the smallest. + server_slot * preempt_pick_victim() { + server_slot * leader = nullptr; + int32_t n_running = 0; + + for (auto & slot : slots) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + n_running++; + + if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) { + leader = &slot; + } + } + } + + if (n_running < 2) { + // a single conversation that does not fit the pool on its own is a real context + // overflow and not a scheduling problem - leave it to the existing error path + return nullptr; + } + + server_slot * victim = nullptr; + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_GENERATING) { + continue; // a generating slot is the one with a clean point to stop at + } + + if (&slot == leader || slot.prompt.n_tokens() == 0) { + continue; + } + + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + continue; // n_cmpl > 1 slots share one sequence, out of scope here + } + + const bool starved = slot.n_preempt >= PREEMPT_N_STARVED; + const bool starved_cur = victim && victim->n_preempt >= PREEMPT_N_STARVED; + + if (!victim || + (starved_cur && !starved) || + (starved_cur == starved && slot.prompt.n_tokens() < victim->prompt.n_tokens())) { + victim = &slot; + } + } + + return victim; + } + + // called once per update_slots(), before the batch is built: 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 from the picture without unpicking a + // half-decoded batch + void update_preemption() { + if (!params_base.kv_unified || slots.size() < 2) { + return; // with a cache per slot, no slot can take another one's cells + } + + const int32_t n_cells = n_ctx; + + // put back what fits, the most-preempted slot first + for (;;) { + server_slot * best = nullptr; + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_PREEMPTED) { + continue; + } + + if (!best || + slot.n_preempt > best->n_preempt || + (slot.n_preempt == best->n_preempt && slot.t_preempt_us < best->t_preempt_us)) { + best = &slot; + } + } + + if (!best) { + break; + } + + const int32_t n_need = best->prompt.n_tokens() + 1 + preempt_n_spec_max(); + + // room for the sequence AND for the next step of everything already running, + // so that a resume cannot immediately trigger the preemption of someone else. + // A cached prompt on an idle slot is worth less than a conversation waiting to + // continue, so give those cells up first - same call the KV-full path makes. + while (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells && + try_clear_idle_slots()) { + } + + if (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells) { + break; + } + + const int64_t t_start = ggml_time_us(); + + if (!best->preempt_restore()) { + // update_slots() runs in a tight loop while tasks are pending, so a counter + // alone burns its whole budget in a couple of milliseconds. Give up only on + // a slot that has been failing for a while, and keep the log quiet. + if (best->n_preempt_fail % 64 == 1) { + SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", + best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); + } + + if (best->n_preempt_fail >= PREEMPT_N_FAIL_MAX && + ggml_time_us() - best->t_preempt_us > PREEMPT_FAIL_US) { + send_error(*best, "failed to restore the preempted sequence"); + best->release(); + } + + break; + } + + n_resume_total++; + + SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", + (ggml_time_us() - best->t_preempt_us) / 1e6, + best->prompt.n_tokens(), + (ggml_time_us() - t_start) / 1e3, + preempt_kv_used(), n_cells, + best->n_preempt); + } + + // forced preemption, for the determinism test only + if (preempt_test_every > 0) { + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_GENERATING && + (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && + slot.preempt_save()) { + n_preempt_total++; + + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", + (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + } + } + } + + // and take cells back until the next decode fits + for (;;) { + const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); + + if (n_used + PREEMPT_N_MARGIN <= n_cells) { + break; + } + + // a prompt cached on an idle slot is the cheapest thing in the pool to give up + if (try_clear_idle_slots()) { + continue; + } + + server_slot * victim = preempt_pick_victim(); + + if (!victim) { + SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted\n", n_used, n_cells); + break; + } + + const int32_t n_tokens = victim->prompt.n_tokens(); + const int64_t t_start = ggml_time_us(); + + if (!victim->preempt_save()) { + break; // could not park it; the existing retry ladder is still behind us + } + + n_preempt_total++; + + SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + } + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -2715,6 +3113,9 @@ struct server_context_impl { } } + // [TAG_PREEMPT] make the pool fit the step that is about to be built + update_preemption(); + try { scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); From 41cbff49d02c21998c66215ad473b1686c96c2e7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 00:41:11 +0000 Subject: [PATCH 02/29] server: test that preemption keeps the output and finishes every slot 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. --- tools/server/tests/unit/test_preempt.py | 110 ++++++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 tools/server/tests/unit/test_preempt.py diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py new file mode 100644 index 000000000000..479b62b91ef5 --- /dev/null +++ b/tools/server/tests/unit/test_preempt.py @@ -0,0 +1,110 @@ +import os +import tempfile +import pytest +from utils import * + +# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked +# (its sequence copied to host RAM, its cells released) instead of every slot being +# terminated. Both tests need more than one slot and --kv-unified, which is the only +# configuration where one slot can take another one's cells. + +server = ServerPreset.tinyllama2() + + +class LogReader: + def __init__(self, path): + self.path = path + self.pos = 0 + + def drain(self): + with open(self.path) as f: + f.seek(self.pos) + content = f.read() + self.pos = f.tell() + return content + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + server.server_slots = True + server.temperature = 0.0 + server.seed = 42 + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + + +def _complete(n_predict: int, prompt: str = "Hi how are you"): + res = server.make_request("POST", "/completion", data={ + "n_predict": n_predict, + "prompt": prompt, + "ignore_eos": True, + "return_tokens": True, + "temperature": 0.0, + "seed": 42, + }) + return res + + +def test_forced_preemption_does_not_change_the_output(): + # Park and restore the only running slot every 8 tokens. With one request the batch + # has the same shape at every step whether or not the slot was parked in between, so + # any difference in the output is the preemption's fault and nothing else's. + global server + server.n_ctx = 512 + server.start() + reference = _complete(64) + assert reference.status_code == 200 + assert reference.body["timings"]["predicted_n"] == 64 + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_EVERY = 8" in log.drain() + + preempted = _complete(64) + assert preempted.status_code == 200 + assert preempted.body["timings"]["predicted_n"] == 64 + + text = log.drain() + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 + + assert preempted.body["content"] == reference.body["content"] + assert preempted.body["tokens"] == reference.body["tokens"] + + +def test_two_slots_that_overflow_the_pool_together_both_finish(): + # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well + # under 256. Together they do not, 336 against 256. Without preemption the retry + # ladder ends with "Context size has been exceeded" on every processing slot; with it + # the smaller slot is parked until the leader finishes and its cells are purged, and + # then it resumes from the token it was parked on. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + From 32c0a77e16ff62630a335f4caeb64263f9b57484 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 01:03:22 +0000 Subject: [PATCH 03/29] server: park prompt-processing slots too, restore whatever fits first, 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. --- common/arg.cpp | 8 ++ common/common.h | 1 + tools/server/README.md | 1 + tools/server/server-context.cpp | 124 +++++++++++++++++++----- tools/server/tests/unit/test_preempt.py | 119 +++++++++++++++++++++++ 5 files changed, 230 insertions(+), 23 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d0..5bfa4adcdf0d 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1709,6 +1709,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.cache_ram_mib = value; } ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--preempt-ram"}, "N", + string_format("with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; " + "N is the maximum host RAM for parked sequences in MiB (default: %d, -1 - no limit, 0 - disable)", params.preempt_ram_mib), + [](common_params & params, int value) { + params.preempt_ram_mib = value; + } + ).set_env("LLAMA_ARG_PREEMPT_RAM").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index de49dac9f63a..c99269f9a967 100644 --- a/common/common.h +++ b/common/common.h @@ -614,6 +614,7 @@ struct common_params { int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = disable preemption std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa9..52e5c1b8c7db 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -164,6 +164,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | +| `--preempt-ram N` | with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; N is the maximum host RAM for parked sequences in MiB (default: 8192, -1 - no limit, 0 - disable)
(env: LLAMA_ARG_PREEMPT_RAM) | | `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)
(env: LLAMA_ARG_KV_UNIFIED) | | `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)
(env: LLAMA_ARG_CACHE_IDLE_SLOTS) | | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 85e9494563d1..a41e5feda5d6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -335,6 +335,12 @@ struct server_slot { preempt_state_dft.shrink_to_fit(); } + // bytes preempt_save() would need for this slot right now + size_t preempt_state_required() const { + return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) + + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0); + } + // copy the sequence out of the cache and release its cells bool preempt_save() { const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); @@ -413,8 +419,10 @@ struct server_slot { state = state_before_preempt; // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks - // that the draft context is where it should be, which the restore above ensures - if (can_speculate()) { + // that the draft context is where it should be, which the restore above ensures. + // A slot parked while still processing its prompt makes that transition itself + // once the prompt is done. + if (state == SLOT_STATE_GENERATING && can_speculate()) { common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } @@ -2845,6 +2853,43 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } + // host RAM the parked sequences hold right now + size_t preempt_ram_used() const { + size_t res = 0; + + for (const auto & slot : slots) { + res += slot.preempt_state_size(); + } + + return res; + } + + // whether parking this slot stays under --preempt-ram + bool preempt_fits_budget(const server_slot & slot) const { + if (params_base.preempt_ram_mib < 0) { + return true; + } + + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + + return preempt_ram_used() + slot.preempt_state_required() <= budget; + } + + // cells the slot will ask for on its next step once it is back in the pool + int32_t preempt_n_need(const server_slot & slot) const { + int32_t res = slot.prompt.n_tokens(); + + if (slot.state_before_preempt == SLOT_STATE_GENERATING) { + res += 1 + preempt_n_spec_max(); + } else { + const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + + res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); + } + + return res; + } + // Cells the pool is holding right now. A released slot keeps its prompt in the cache // for the next request to reuse as a prefix, so idle slots count too: the first version // of this counted only the running ones, decided a pool holding 8185 cached cells was @@ -2920,11 +2965,18 @@ struct server_context_impl { server_slot * victim = nullptr; for (auto & slot : slots) { - if (slot.state != SLOT_STATE_GENERATING) { - continue; // a generating slot is the one with a clean point to stop at + // Before the batch is built every one of these is at a token boundary: a + // generating slot between two sampled tokens, a prompt-processing slot between + // two chunks of its prompt, a started slot with only a cached prefix (or + // nothing) in the pool. A slot holding no cells is still worth parking - it + // is about to ask for a whole batch of them. + if (slot.state != SLOT_STATE_GENERATING && + slot.state != SLOT_STATE_PROCESSING_PROMPT && + slot.state != SLOT_STATE_STARTED) { + continue; } - if (&slot == leader || slot.prompt.n_tokens() == 0) { + if (&slot == leader) { continue; } @@ -2932,6 +2984,10 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + if (!preempt_fits_budget(slot)) { + continue; + } + const bool starved = slot.n_preempt >= PREEMPT_N_STARVED; const bool starved_cur = victim && victim->n_preempt >= PREEMPT_N_STARVED; @@ -2954,39 +3010,57 @@ struct server_context_impl { return; // with a cache per slot, no slot can take another one's cells } + if (params_base.preempt_ram_mib == 0) { + return; // --preempt-ram 0: the KV-full retry ladder, as before + } + const int32_t n_cells = n_ctx; - // put back what fits, the most-preempted slot first + // Put back what fits: the most-preempted slot first, then the one parked longest. + // A slot that does not fit yet must not hold up a smaller one that does: it keeps + // its place at the head of the line, and the smaller one is the first to be parked + // again if the pool fills, so letting it through costs the head nothing. for (;;) { - server_slot * best = nullptr; + std::vector parked; for (auto & slot : slots) { - if (slot.state != SLOT_STATE_PREEMPTED) { - continue; - } - - if (!best || - slot.n_preempt > best->n_preempt || - (slot.n_preempt == best->n_preempt && slot.t_preempt_us < best->t_preempt_us)) { - best = &slot; + if (slot.state == SLOT_STATE_PREEMPTED) { + parked.push_back(&slot); } } - if (!best) { + if (parked.empty()) { break; } - const int32_t n_need = best->prompt.n_tokens() + 1 + preempt_n_spec_max(); + std::sort(parked.begin(), parked.end(), [](const server_slot * a, const server_slot * b) { + if (a->n_preempt != b->n_preempt) { + return a->n_preempt > b->n_preempt; + } - // room for the sequence AND for the next step of everything already running, + return a->t_preempt_us < b->t_preempt_us; + }); + + server_slot * best = nullptr; + + // Room for the sequence AND for the next step of everything already running, // so that a resume cannot immediately trigger the preemption of someone else. // A cached prompt on an idle slot is worth less than a conversation waiting to // continue, so give those cells up first - same call the KV-full path makes. - while (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells && - try_clear_idle_slots()) { + for (;;) { + for (auto * slot : parked) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + best = slot; + break; + } + } + + if (best || !try_clear_idle_slots()) { + break; + } } - if (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells) { + if (!best) { break; } @@ -3025,6 +3099,7 @@ struct server_context_impl { for (auto & slot : slots) { if (slot.state == SLOT_STATE_GENERATING && (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && + preempt_fits_budget(slot) && slot.preempt_save()) { n_preempt_total++; @@ -3050,7 +3125,8 @@ struct server_context_impl { server_slot * victim = preempt_pick_victim(); if (!victim) { - SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted\n", n_used, n_cells); + SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted (parked %.1f MiB of the %d MiB --preempt-ram budget)\n", + n_used, n_cells, preempt_ram_used() / (1024.0 * 1024.0), params_base.preempt_ram_mib); break; } @@ -3401,7 +3477,9 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - if (!slot.is_processing()) { + // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to + // batch; it takes no part in this pass until it is restored + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { return; } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 479b62b91ef5..2e57ebbdfd97 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -1,4 +1,5 @@ import os +import time import tempfile import pytest from utils import * @@ -37,6 +38,7 @@ def create_server(): os.close(fd) yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) def _complete(n_predict: int, prompt: str = "Hi how are you"): @@ -108,3 +110,120 @@ def test_two_slots_that_overflow_the_pool_together_both_finish(): assert res.body["truncated"] is False assert len(res.body["tokens"]) == n_predict + + +_WORDS = ( + "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor " + "incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud " + "exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure " + "dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. " + "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt " + "mollit anim id est laborum. " +) * 4 + + +def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: + """A prompt whose token count is in [n_tokens - 12, n_tokens], measured on the server.""" + words = (salt + " " + _WORDS).split() + while words: + text = " ".join(words) + res = server.make_request("POST", "/tokenize", data={"content": text}) + assert res.status_code == 200 + n = len(res.body["tokens"]) + if n <= n_tokens: + assert n >= n_tokens - 12, f"could not land near {n_tokens} tokens, got {n}" + return text, n + # about four tokens per word on this model's vocabulary + words = words[: len(words) - max(1, (n - n_tokens) // 8)] + raise AssertionError("empty prompt") + + +def test_two_prompts_that_overflow_the_pool_together_both_finish(): + # Neither slot ever generates before the pool is full: both are still processing their + # prompts. A prompt-processing slot is between two chunks of its prompt, which is as + # clean a boundary as the one between two sampled tokens, so it is parked the same way. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + prompt_a, n_a = _prompt_of_about(150, "Alpha") + prompt_b, n_b = _prompt_of_about(150, "Bravo") + n_predict = 16 + assert n_a + n_predict <= 256 and n_b + n_predict <= 256 + assert n_a + n_b + 2 * n_predict > 256 + + results = parallel_function_calls([ + (_complete, (n_predict, prompt_a)), + (_complete, (n_predict, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert len(res.body["tokens"]) == n_predict + + +def test_a_generating_slot_and_a_large_prompt_both_finish(): + # One slot is generating a long answer to a short prompt when a large prompt arrives + # beside it. Together they need far more than the pool has. The prompt is admitted + # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish. + # This model produces a thousand tokens a second, so the second request is sent right + # behind the first rather than after a delay: its prompt takes several batches to + # process, which is enough for the two to overlap however fast the first one runs. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + prompt_b, n_b = _prompt_of_about(150, "Charlie") + # b lives long enough for the two to collide: the first run of this used 16 tokens + # and b was finished and purged before a had grown into it + n_predict_a = 230 + n_predict_b = 90 + assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 + assert 8 + n_predict_a + n_b + n_predict_b > 256 + + def _late(n_predict, prompt): + time.sleep(0.02) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_a, "Hi how are you")), + (_late, (n_predict_b, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + + assert results[0].status_code == 200 + assert results[0].body["timings"]["predicted_n"] == n_predict_a + assert results[1].status_code == 200 + assert results[1].body["timings"]["predicted_n"] == n_predict_b + + +def test_preempt_ram_zero_disables_preemption(): + # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the + # KV-full path ends the requests the way it always did. + global server + server.n_ctx = 256 + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "preempted:" not in text + assert "Context size has been exceeded" in text + assert any(res.status_code != 200 for res in results) From 63cdac4d7d9c7679d08e2dd3ee37120c75774b5c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 01:13:28 +0000 Subject: [PATCH 04/29] server: report preemption through /metrics and /slots 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. --- tools/server/README.md | 4 +++ tools/server/server-common.h | 4 +++ tools/server/server-context.cpp | 19 ++++++++---- tools/server/server-task.cpp | 16 ++++++++++ tools/server/server-task.h | 2 ++ tools/server/tests/unit/test_preempt.py | 40 +++++++++++++++++++++++++ 6 files changed, 79 insertions(+), 6 deletions(-) diff --git a/tools/server/README.md b/tools/server/README.md index 52e5c1b8c7db..7b4a0330340f 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -1139,6 +1139,10 @@ In *router mode* the query param `?model={model_id}` has to be set. This endpoin | `llamacpp:spec_decode_num_accepted_tokens_total` | Counter | Total draft tokens accepted by the target model (0 when spec-decode is off). | | `llamacpp:spec_decode_num_drafts_total` | Counter | Total speculative decoding verification steps (0 when spec-decode is off). | | `llamacpp:spec_decode_num_accepted_tokens_per_pos_total` | Counter | Accepted tokens per draft position (labeled `position="N"`; absent when spec-decode is off or before the first completed speculative request). | +| `llamacpp:n_preempt_total` | Counter | Slots parked to make room in the unified KV cache (0 unless `--kv-unified` with more than one slot). | +| `llamacpp:n_resume_total` | Counter | Parked slots put back. | +| `llamacpp:requests_preempted` | Gauge | Requests currently parked, waiting for room in the unified KV cache. | +| `llamacpp:preempt_ram_bytes` | Gauge | Host RAM held by parked sequences. | ### POST `/slots/{id_slot}?action=save`: Save the prompt cache of the specified slot to a file. diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f8ea82ef4cf5..f0cf76b8c501 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,6 +467,10 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; + // [TAG_PREEMPT] slots parked to make room in the unified KV pool, and put back + uint64_t n_preempt = 0; + uint64_t n_resume = 0; + uint64_t n_draft_tokens = 0; // Total draft tokens generated uint64_t n_draft_accepted = 0; // Draft tokens actually accepted uint64_t n_draft_verif_steps = 0; // Total draft token verification steps by the target model diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a41e5feda5d6..6723c51397ed 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -796,6 +796,8 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, + {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"n_preempt", n_preempt}, }; const auto & ptask = task ? task : task_prev; @@ -2546,11 +2548,15 @@ struct server_context_impl { case SERVER_TASK_TYPE_METRICS: { int n_processing_slots = 0; + int n_preempted_slots = 0; for (server_slot & slot : slots) { if (slot.is_processing()) { n_processing_slots++; } + if (slot.state == SLOT_STATE_PREEMPTED) { + n_preempted_slots++; + } } SRV_DBG("n_processing_slots = %d\n", n_processing_slots); @@ -2558,6 +2564,8 @@ struct server_context_impl { res->id = task.id; res->n_processing_slots = n_processing_slots; res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size(); + res->n_preempted_slots = n_preempted_slots; + res->preempt_ram_bytes = preempt_ram_used(); res->metrics = metrics; if (task.metrics_reset_bucket) { @@ -2839,8 +2847,6 @@ struct server_context_impl { // [TAG_PREEMPT] server-side request preemption // - int64_t n_preempt_total = 0; - int64_t n_resume_total = 0; // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, // whether or not the pool is under pressure. It exists to answer the only question that @@ -3084,7 +3090,7 @@ struct server_context_impl { break; } - n_resume_total++; + metrics.n_resume++; SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", (ggml_time_us() - best->t_preempt_us) / 1e6, @@ -3101,7 +3107,7 @@ struct server_context_impl { (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && preempt_fits_budget(slot) && slot.preempt_save()) { - n_preempt_total++; + metrics.n_preempt++; SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); @@ -3137,7 +3143,7 @@ struct server_context_impl { break; // could not park it; the existing retry ladder is still behind us } - n_preempt_total++; + metrics.n_preempt++; SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, @@ -4422,7 +4428,8 @@ struct server_context_impl { void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { metrics.n_decode++; for (const auto & slot : slots) { - if (slot.is_processing()) { + // [TAG_PREEMPT] a parked slot is processing but took no part in this decode + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { metrics.n_busy_slots++; } metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens()); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..9afe3c7f06a8 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1562,6 +1562,14 @@ std::string server_task_result_metrics::to_metrics() { "spec_decode_num_drafts_total", "Speculative: Total speculative decoding verification steps", (double) metrics.n_draft_verif_steps + }, { + "n_preempt_total", + "Preemption: Total slots parked to make room in the unified KV cache", + (double) metrics.n_preempt + }, { + "n_resume_total", + "Preemption: Total parked slots put back", + (double) metrics.n_resume }, }; @@ -1586,6 +1594,14 @@ std::string server_task_result_metrics::to_metrics() { "n_busy_slots_per_decode", "Average number of busy slots per llama_decode() call", (double) metrics.n_busy_slots / std::max((double) metrics.n_decode, 1.0) + }, { + "requests_preempted", + "Preemption: Number of requests currently parked, waiting for room in the unified KV cache", + (double) n_preempted_slots + }, { + "preempt_ram_bytes", + "Preemption: Host RAM held by parked sequences", + (double) preempt_ram_bytes }, }; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..00734924bc63 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -494,6 +494,8 @@ struct server_task_result_metrics : server_task_result { // these are immediate stats, not accumulated (server_metrics is cumulative) int n_processing_slots = 0; int n_tasks_deferred = 0; + int n_preempted_slots = 0; // [TAG_PREEMPT] processing slots currently parked + size_t preempt_ram_bytes = 0; // [TAG_PREEMPT] host RAM their parked sequences hold server_metrics metrics; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 2e57ebbdfd97..0da885bcafd5 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -227,3 +227,43 @@ def test_preempt_ram_zero_disables_preemption(): assert "preempted:" not in text assert "Context size has been exceeded" in text assert any(res.status_code != 200 for res in results) + + +def test_metrics_and_slots_report_the_parked_state(): + # A client that wants to tell a parked chat from a slow one reads /slots, and an + # operator reads /metrics. Both must show the preemption happening, and the counters + # must survive the requests finishing. + global server + server.n_ctx = 256 + server.server_metrics = True + server.start() + + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + for slot in res.body: + assert slot["is_preempted"] is False + assert slot["n_preempt"] == 0 + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + for res in results: + assert res.status_code == 200 + + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + metrics = {} + for line in res.body.splitlines(): + if line.startswith("llamacpp:"): + name, value = line.split(" ", 1) + metrics[name[len("llamacpp:"):]] = float(value) + assert metrics["n_preempt_total"] >= 1 + assert metrics["n_resume_total"] == metrics["n_preempt_total"] + assert metrics["requests_preempted"] == 0 + assert metrics["preempt_ram_bytes"] == 0 + + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot" From 75944db803a028a2401dc479cea98e433fcf4298 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 11:05:33 +0000 Subject: [PATCH 05/29] Batch fragmented sequence restores through bounded host staging --- src/llama-context.cpp | 38 +++++++++++++++++++++++-- tests/test-state-restore-fragmented.cpp | 15 ++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b7..66940d4fc61a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2617,8 +2617,42 @@ class llama_io_read_host : public llama_io_read_i { ~llama_io_read_host() { // flush the reads - for (const auto & rinfo : rinfos) { - ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); + for (size_t i = 0; i < rinfos.size();) { + auto * tensor = rinfos[i].tensor; + size_t end = i + 1; + while (end < rinfos.size() && rinfos[end].tensor == tensor) { + end++; + } + const size_t tensor_bytes = ggml_nbytes(tensor); + auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + // A fragmented sequence can require thousands of synchronous device + // transfers per layer. For bounded tensors, stage the tensor once and + // preserve every byte belonging to other sequences. Bound scratch RAM + // and leave ordinary contiguous transfers on their original fast path. + if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && + !ggml_backend_buffer_is_host(buffer)) { + std::vector staging; + try { + staging.resize(tensor_bytes); + } catch (const std::bad_alloc &) { + // Fall back to the individual transfers below. + } + if (!staging.empty()) { + ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); + for (size_t j = i; j < end; ++j) { + const auto & rinfo = rinfos[j]; + GGML_ASSERT(rinfo.offset <= tensor_bytes && rinfo.size <= tensor_bytes - rinfo.offset); + memcpy(staging.data() + rinfo.offset, rinfo.ptr, rinfo.size); + } + ggml_backend_tensor_set(tensor, staging.data(), 0, tensor_bytes); + i = end; + continue; + } + } + for (; i < end; ++i) { + const auto & rinfo = rinfos[i]; + ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); + } } } diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index d5548afba179..428a92529811 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,6 +73,14 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); + // A fragmented restore may stage a whole device tensor. Check every + // sequence byte-for-byte, including the neighbours that must be preserved. + std::vector> before(params.n_parallel); + for (int s = 0; s < params.n_parallel; ++s) { + before[s].resize(llama_state_seq_get_size(ctx, s)); + GGML_ASSERT(llama_state_seq_get_data(ctx, before[s].data(), before[s].size(), s) == before[s].size()); + } + // clear seq 1 to create a "hole" in the KV cache (fragmentation) // 0.20.20.20.2.... llama_memory_t mem = llama_get_memory(ctx); @@ -96,6 +104,13 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : restored state into seq 1, %zu bytes\n", __func__, nset); + for (int s = 0; s < params.n_parallel; ++s) { + std::vector after(llama_state_seq_get_size(ctx, s)); + GGML_ASSERT(llama_state_seq_get_data(ctx, after.data(), after.size(), s) == after.size()); + GGML_ASSERT(before[s] == after); + } + fprintf(stderr, "%s : all %d sequence snapshots are byte-identical after restore\n", __func__, params.n_parallel); + // Verify we can decode with the restored state // Generate one token to verify the restored state is usable auto sparams = llama_sampler_chain_default_params(); From 50b617ae460e148e5130fdd3d7d6499ab4cf575f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 12:27:06 +0000 Subject: [PATCH 06/29] server: restore parked slots head of line by park time Parked slots came back most-preempted first, then longest parked, and a smaller slot could pass a head that did not fit yet, on the reasoning that the smaller slot is the first to be parked again and so costs the head nothing. It does cost the head: the small slot squeezes in, grows, and is parked again with one more preemption to its name, which puts it ahead of the head next time as well. In a four-chat Studio run on an 8192-cell pool one chat parked at 4160 tokens waited 7 minutes 22 seconds while two smaller chats were restored and parked three times each. The order is now the line itself: the slot parked longest comes back first, and nobody else while it does not fit, which is what a first come first served queue of preempted requests amounts to. Simulated with the smallest-slot victim policy over 60 seeds, this cuts the longest single wait at eight chats by 2.5 to 3x (60 to 78 s down to 24 to 28 s) for 0 to 3 percent of makespan at 8192 cells and 3 to 6 percent at 16384, and parks less often and copies fewer cells at 8192. At four chats every order is within 2 percent on everything. LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order. The victim side is unchanged, including the starvation guard. --- tools/server/server-context.cpp | 35 +++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6723c51397ed..b8247e9c00fb 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -77,6 +78,22 @@ enum slot_state { // text and the same open stream. A streaming client sees a pause, not an error. constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected + +// [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody +// passes a head that does not fit yet: the head keeps the room the pool frees until it fits, so +// its wait is bounded by the slots ahead of it and not by how often a smaller slot can squeeze +// in, grow, and be parked again. Simulated over 60 seeds at eight chats this cuts the longest +// single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often. +// LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest +// parked, and a smaller slot may pass a head that does not fit. +static bool preempt_resume_head_of_line() { + static const bool head_of_line = []() { + const char * val = getenv("LLAMA_SERVER_PREEMPT_RESUME"); + return !(val && strcmp(val, "pass") == 0); + }(); + + return head_of_line; +} constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked @@ -3022,10 +3039,12 @@ struct server_context_impl { const int32_t n_cells = n_ctx; - // Put back what fits: the most-preempted slot first, then the one parked longest. - // A slot that does not fit yet must not hold up a smaller one that does: it keeps - // its place at the head of the line, and the smaller one is the first to be parked - // again if the pool fills, so letting it through costs the head nothing. + // Put back what fits, in the order preempt_resume_head_of_line() describes: by default + // the slot parked longest, and only that one until it fits; under + // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked + // longest, and a smaller slot may pass a head that does not fit. + const bool head_of_line = preempt_resume_head_of_line(); + for (;;) { std::vector parked; @@ -3039,14 +3058,18 @@ struct server_context_impl { break; } - std::sort(parked.begin(), parked.end(), [](const server_slot * a, const server_slot * b) { - if (a->n_preempt != b->n_preempt) { + std::sort(parked.begin(), parked.end(), [head_of_line](const server_slot * a, const server_slot * b) { + if (!head_of_line && a->n_preempt != b->n_preempt) { return a->n_preempt > b->n_preempt; } return a->t_preempt_us < b->t_preempt_us; }); + if (head_of_line) { + parked.resize(1); + } + server_slot * best = nullptr; // Room for the sequence AND for the next step of everything already running, From 8057a74afd056bf78896c50b4359a9bef72a0173 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 12:54:27 +0000 Subject: [PATCH 07/29] server: read the resume order once at load, log it, and refuse a value that is neither head nor pass LLAMA_SERVER_PREEMPT_RESUME was read lazily on the first resume and never announced, unlike the other LLAMA_SERVER_PREEMPT_* knobs, and any value other than pass silently meant head of line. It is now read in load_model() next to LLAMA_SERVER_PREEMPT_EVERY: head is the default, pass is logged as a warning, and anything else fails the load with a message. --- tools/server/server-context.cpp | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b8247e9c00fb..5dc3c3277662 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -86,13 +86,11 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is // single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often. // LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest // parked, and a smaller slot may pass a head that does not fit. -static bool preempt_resume_head_of_line() { - static const bool head_of_line = []() { - const char * val = getenv("LLAMA_SERVER_PREEMPT_RESUME"); - return !(val && strcmp(val, "pass") == 0); - }(); +// LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. +static bool g_preempt_resume_head_of_line = true; - return head_of_line; +static bool preempt_resume_head_of_line() { + return g_preempt_resume_head_of_line; } constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked @@ -1420,6 +1418,17 @@ struct server_context_impl { } { + const char * LLAMA_SERVER_PREEMPT_RESUME = getenv("LLAMA_SERVER_PREEMPT_RESUME"); + if (LLAMA_SERVER_PREEMPT_RESUME && strcmp(LLAMA_SERVER_PREEMPT_RESUME, "head") != 0) { + if (strcmp(LLAMA_SERVER_PREEMPT_RESUME, "pass") != 0) { + SRV_ERR("LLAMA_SERVER_PREEMPT_RESUME = %s is not a resume order; use head (the default) or pass\n", + LLAMA_SERVER_PREEMPT_RESUME); + return false; + } + g_preempt_resume_head_of_line = false; + SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_RESUME = pass (parked slots come back most-preempted first, and a smaller slot may pass a head that does not fit)\n"); + } + const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; From 2a7e277abef3a17ba7708f3b19f897f0dac03e5b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 15:37:09 +0000 Subject: [PATCH 08/29] server: fail a parked sequence that cannot fit the pool alone instead of parking it for ever A parked slot was restored only when its sequence, its next step and the scheduling margin all fit, even with the pool empty. A prompt within n_ctx that was parked before it took any cells, but too close to n_ctx to leave room for the margin, therefore never fit, and since a restore was never attempted it never reached the restore-failure path either: it stayed parked for ever. With parked slots restored head of line, such a head would have held every slot behind it as well. Two changes. The margin is headroom for the other residents, so with nothing resident there is nobody to keep it for and a sequence that fits the pool exactly is let back in. A parked sequence whose next step would not fit an empty pool at all is the single-conversation overflow the KV-full path reports, so it is reported the same way, "Context size has been exceeded", and released, and the line is rescanned without it. New test: two prompts of 240 tokens on a 256-cell pool, sent together, both complete. Server preemption tests 7 of 7. --- tools/server/server-context.cpp | 32 ++++++++++++++++++++++++- tools/server/tests/unit/test_preempt.py | 27 +++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5dc3c3277662..b4381c9eba93 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3081,13 +3081,43 @@ struct server_context_impl { server_slot * best = nullptr; + // A parked slot whose sequence plus its next step would not fit an empty pool can + // never be restored, and would otherwise sit at the head of the line for ever + // without a restore ever being attempted: a prompt within n_ctx that was parked + // before it took any cells, but too close to n_ctx to leave room for its first + // batch. That is the single-conversation overflow the KV-full path reports, so + // report it the same way and rescan the line without it. + { + server_slot * impossible = nullptr; + + for (auto * slot : parked) { + if (preempt_n_need(*slot) > n_cells) { + impossible = slot; + break; + } + } + + if (impossible) { + SLT_WRN(*impossible, "parked sequence of %d tokens cannot fit the pool of %d cells even alone, failing it\n", + preempt_n_need(*impossible), n_cells); + send_error(*impossible, "Context size has been exceeded."); + impossible->release(); + continue; + } + } + // Room for the sequence AND for the next step of everything already running, // so that a resume cannot immediately trigger the preemption of someone else. + // The margin is headroom for the others; with nothing resident there is nobody + // to keep it for, so a sequence that fits the pool exactly is let back in. // A cached prompt on an idle slot is worth less than a conversation waiting to // continue, so give those cells up first - same call the KV-full path makes. for (;;) { + const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); + const int32_t margin = occupied == 0 ? 0 : PREEMPT_N_MARGIN; + for (auto * slot : parked) { - if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + if (occupied + preempt_n_need(*slot) + margin <= n_cells) { best = slot; break; } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 0da885bcafd5..080503971950 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -267,3 +267,30 @@ def test_metrics_and_slots_report_the_parked_state(): res = server.make_request("GET", "/slots") assert res.status_code == 200 assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot" + + +def test_two_prompts_near_the_context_size_both_complete(): + # Two prompts that each fit the context alone but not together. The second one is + # parked before it takes any cells, and it is close enough to n_ctx that its sequence + # plus its first batch would not leave the usual scheduling margin. It must still be + # restored once the first one finishes: with nothing resident there is nobody to keep + # the margin for. Before the fix it was parked for ever, with no restore ever tried. + global server + server.n_ctx = 256 + # the whole prompt in one batch, so the parked slot's first step is the whole prompt + server.n_batch = 256 + server.start() + log = LogReader(server.log_path) + + # sized in tokens, not words: the prompt is the token ids of a short sentence repeated + base = server.make_request("POST", "/tokenize", data={"content": "Once upon a time there was a little girl"}).body["tokens"] + long_prompt = (base * 64)[:250] + n_predict = 4 + together = parallel_function_calls([(_complete, (n_predict, long_prompt)) for _ in range(2)]) + + text = log.drain() + assert "cannot fit the pool" not in text + + for res in together: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict From 84311fd3d01f4455c4a6e92e60db5f36dc246e48 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:06:13 +0000 Subject: [PATCH 09/29] server: park instead of ending everyone when the KV-full retry ladder runs out The planner parks ahead of the decode, so the retry ladder is only reached when its estimate was wrong. When it was, one token finding no cell still ended every processing slot with the context error, parked slots included, which is the failure this branch exists to remove. The ladder now has a last resort. Every resident slot is rewound to the token boundary the cache is at: a batch is applied one chunk at a time and the chunk that failed left nothing behind, so the tokens added for a slot that were never decoded come off, the sampled token stays in the slot and goes into the next batch the way it went into this one, and a prompt whose last chunk was in the batch is marked not done again. The smallest slots are then parked until the planner's own bound holds, the batch is given up, and the next pass rebuilds it from the survivors; the planner brings the parked ones back as cells free up. One conversation that does not fit the pool alone, or a multimodal prompt, keep the old path. Parked slots are no longer part of the error sweep either way: they hold nothing in the batch and nothing in the cache. LLAMA_SERVER_PREEMPT_PLANNER=off is a test knob that leaves the pool to the ladder, so the last resort can be reached on purpose. Two tests use it: two generations that fit alone and not together, and a prompt in flight when the pool runs out, whose failed chunk is processed once after the rewind. --- tools/server/server-context.cpp | 144 +++++++++++++++++++++++- tools/server/tests/unit/test_preempt.py | 73 ++++++++++++ 2 files changed, 216 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b4381c9eba93..031cd817e9ef 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -444,6 +444,33 @@ struct server_slot { return true; } + // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds for this sequence. + // For a batch that is given up after it was built: the tokens added for this slot + // that were never decoded come off, the sampled token stays in `sampled` and goes into + // the next batch the way it went into this one, and a draft is a prediction that goes + // with them. A chunk that failed to decode left nothing in the cache, so the cache is + // the boundary. + void rewind_to_cache() { + const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; + + if (n_cached < prompt.n_tokens()) { + prompt.tokens.keep_first(n_cached); + } + + // a prompt whose last chunk was in the batch was marked done when the chunk was + // built; the chunk never ran, so the prompt is not done + if (state == SLOT_STATE_DONE_PROMPT && task && prompt.n_tokens() < task->n_tokens()) { + state = SLOT_STATE_PROCESSING_PROMPT; + } + + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + } + std::vector lora; int32_t alora_invocation_start = -1; @@ -1436,6 +1463,13 @@ struct server_context_impl { SRV_WRN("LLAMA_SERVER_PREEMPT_EVERY = %d (test knob: preempting every slot every %d tokens)\n", preempt_test_every, preempt_test_every); } + + const char * LLAMA_SERVER_PREEMPT_PLANNER = getenv("LLAMA_SERVER_PREEMPT_PLANNER"); + preempt_planner_off = LLAMA_SERVER_PREEMPT_PLANNER && strcmp(LLAMA_SERVER_PREEMPT_PLANNER, "off") == 0; + + if (preempt_planner_off) { + SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_PLANNER = off (test knob: nothing is parked ahead of the decode, only as a last resort)\n"); + } } { @@ -2881,6 +2915,14 @@ struct server_context_impl { // uninterrupted one is the preemption's fault and nothing else's. int32_t preempt_test_every = 0; + // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so + // the KV-full retry ladder and its last resort are the only thing between a full pool + // and the context error + bool preempt_planner_off = false; + + // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop + bool preempt_batch_abandoned = false; + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -3177,6 +3219,10 @@ struct server_context_impl { } } + if (preempt_planner_off) { + return; // test knob: leave the pool to the retry ladder and its last resort + } + // and take cells back until the next decode fits for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); @@ -3308,6 +3354,13 @@ struct server_context_impl { llama_synchronize(ctx_tgt); #endif + if (preempt_batch_abandoned) { + // [TAG_PREEMPT] the rest of this batch was never decoded and the slots no + // longer describe it; the next pass builds a new one + preempt_batch_abandoned = false; + break; + } + if (ok) { // move the head of the batch forward with the number of tokens we just processed off_next = off + n_tokens; @@ -4073,6 +4126,87 @@ struct server_context_impl { } } + // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is + // the context error for every slot in the batch. With a park budget the batch is given + // up instead: every resident slot is rewound to the token boundary the cache is at (a + // batch is applied one chunk at a time, and the chunk that failed left nothing behind), + // the smallest are parked until the planner's own bound holds again, and the next + // update_slots() rebuilds the batch from the survivors. The planner brings the parked + // ones back as cells free up. A multimodal prompt has no boundary the cache can name, + // so it keeps the old path. + bool preempt_last_resort(int32_t off) { + if (!params_base.kv_unified || params_base.preempt_ram_mib == 0 || slots.size() < 2) { + return false; + } + + int32_t n_running = 0; + + for (auto & slot : slots) { + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { + continue; + } + + if (slot.prompt.tokens.has_mtmd) { + return false; + } + + n_running++; + } + + if (n_running < 2) { + return false; // one conversation that does not fit alone is a real overflow + } + + for (auto & slot : slots) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && slot.state != SLOT_STATE_WAIT_OTHER) { + slot.rewind_to_cache(); + } + } + + const int32_t n_cells = n_ctx; + int32_t n_parked = 0; + + for (;;) { + const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); + + if (n_parked > 0 && n_used + PREEMPT_N_MARGIN <= n_cells) { + break; + } + + server_slot * victim = preempt_pick_victim(); + + if (!victim) { + break; + } + + const int32_t n_tokens = victim->prompt.n_tokens(); + const int64_t t_start = ggml_time_us(); + + if (!victim->preempt_save()) { + break; + } + + metrics.n_preempt++; + n_parked++; + + SLT_WRN(*victim, "preempted as a last resort: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + } + + if (n_parked == 0) { + return false; // nothing could be parked: the error path clears what the rewind left + } + + SRV_WRN("last resort: batch given up at off = %d, %d slot(s) parked, kv %d/%d resident\n", + off, n_parked, preempt_kv_used(), n_cells); + + return true; + } + // returns true = success ; false = retry with smaller batch size // throw std::runtime_error on fatal error bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) { @@ -4121,6 +4255,12 @@ struct server_context_impl { std::string err; if (n_batch == 1 && ret == 1) { + // [TAG_PREEMPT] park instead of ending everyone, when there is a budget to park into + if (preempt_last_resort(off)) { + preempt_batch_abandoned = true; + return true; + } + // TODO: try to terminate only the largest active slot/sequence and continue with the rest // need to remove the tokens from the current batch too err = "Context size has been exceeded."; @@ -4141,7 +4281,9 @@ struct server_context_impl { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); for (auto & slot : slots) { - if (slot.is_processing()) { + // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the + // cache; it is not part of this failure and comes back when there is room + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, err); slot.release(); diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 080503971950..932c1323305b 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -38,6 +38,7 @@ def create_server(): os.close(fd) yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_SERVER_PREEMPT_PLANNER", None) os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) @@ -294,3 +295,75 @@ def test_two_prompts_near_the_context_size_both_complete(): for res in together: assert res.status_code == 200 assert res.body["timings"]["predicted_n"] == n_predict + + +def test_the_last_resort_parks_instead_of_ending_everyone(): + # With the planner off nothing is parked ahead of the decode, so two generations that + # fit alone but not together fill the pool until a single token finds no cell. That + # is where upstream ends every slot with the context error. Instead the batch is + # given up, the smaller slot is parked, the larger one finishes, and the parked one + # comes back and finishes too. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_PLANNER = off" in log.drain() + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" not in text, "the planner was off, nothing may be parked ahead of the decode" + assert "preempted as a last resort" in text + assert "last resort: batch given up" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + + +def test_the_last_resort_rewinds_a_prompt_in_flight(): + # Same, with a prompt being processed when the pool runs out: the chunk that failed + # is taken back off the slot's tokens and processed again after the resume, so the + # prompt is neither skipped nor fed twice. The prompt is far longer than a batch, so + # the failing chunk is a chunk of it, not its last token. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + server.start() + log = LogReader(server.log_path) + + prompt_b, n_b = _prompt_of_about(150, "Charlie") + n_predict_a = 230 + n_predict_b = 90 + assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 + assert 8 + n_predict_a + n_b + n_predict_b > 256 + + def _late(n_predict, prompt): + time.sleep(0.02) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_a, "Hi how are you")), + (_late, (n_predict_b, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted as a last resort" in text + + assert results[0].status_code == 200 + assert results[0].body["timings"]["predicted_n"] == n_predict_a + assert results[1].status_code == 200 + assert results[1].body["timings"]["predicted_n"] == n_predict_b + # the chunk that was in the batch given up is processed once, after the rewind, and + # the count is the prompt plus the BOS the server adds + assert results[1].body["timings"]["prompt_n"] == n_b + 1 From 6dbc4e7b65593e790385f99b3db7c6598a607e6a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:13:27 +0000 Subject: [PATCH 10/29] server: keep a slot's draft in one view when the retry ladder narrows the batch With speculation on, a slot's sampled token and its draft go into the batch as one group and the verify step expects them in one view. The retry ladder halved the batch width without regard for that, so a narrower view could cut through a group and the verify step threw for the slot whose tokens straddled it, before the ladder ever reached its last resort. Halving cannot help there. After purging an idle slot, the ladder now goes to its last resort straight away when the batch holds a draft. --- tools/server/server-context.cpp | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 031cd817e9ef..911256206110 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4207,6 +4207,17 @@ struct server_context_impl { return true; } + // [TAG_PREEMPT] whether a slot in the batch has its sampled token and a draft in it + bool batch_has_spec_groups() const { + for (const auto & slot : slots) { + if (!slot.spec_i_batch.empty()) { + return true; + } + } + + return false; + } + // returns true = success ; false = retry with smaller batch size // throw std::runtime_error on fatal error bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) { @@ -4254,6 +4265,19 @@ struct server_context_impl { { std::string err; + // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have + // to stay in one view: a narrower view splits the group and the verify step + // throws for the slot whose tokens straddle it. Halving is no help there, so + // after the idle slots the ladder goes to its last resort straight away. + if (ret == 1 && n_batch > 1 && batch_has_spec_groups()) { + if (try_clear_idle_slots()) { + SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n"); + return false; // retry at the same width + } + + n_batch = 1; + } + if (n_batch == 1 && ret == 1) { // [TAG_PREEMPT] park instead of ending everyone, when there is a budget to park into if (preempt_last_resort(off)) { From 86845c15e780c56ae0b034adc8fbc01d8fc776d0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:23:09 +0000 Subject: [PATCH 11/29] server: narrow a batch holding a draft the old way when there is no budget to park into The straight route to the last resort only helps when the last resort can park something. Under --preempt-ram 0 the ladder is what it always was, so a sub-batch retry there still stands a chance of fitting. --- tools/server/server-context.cpp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 911256206110..7a3bdb23615d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4134,8 +4134,12 @@ struct server_context_impl { // update_slots() rebuilds the batch from the survivors. The planner brings the parked // ones back as cells free up. A multimodal prompt has no boundary the cache can name, // so it keeps the old path. + bool preempt_last_resort_possible() const { + return params_base.kv_unified && params_base.preempt_ram_mib > 0 && slots.size() >= 2; + } + bool preempt_last_resort(int32_t off) { - if (!params_base.kv_unified || params_base.preempt_ram_mib == 0 || slots.size() < 2) { + if (!preempt_last_resort_possible()) { return false; } @@ -4268,8 +4272,9 @@ struct server_context_impl { // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have // to stay in one view: a narrower view splits the group and the verify step // throws for the slot whose tokens straddle it. Halving is no help there, so - // after the idle slots the ladder goes to its last resort straight away. - if (ret == 1 && n_batch > 1 && batch_has_spec_groups()) { + // after the idle slots the ladder goes to its last resort straight away. With + // no budget to park into the ladder is what it always was. + if (ret == 1 && n_batch > 1 && preempt_last_resort_possible() && batch_has_spec_groups()) { if (try_clear_idle_slots()) { SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n"); return false; // retry at the same width From 662ec2029b06e82ba6110e0c7a4fd5858f075548 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:54:00 +0000 Subject: [PATCH 12/29] server: leave a context without memory to itself An embedding model has no cache to run out of and nothing a park could release, so the planner and the last resort step aside for it. In practice they never fired there anyway: a memory-less model takes its whole prompt in one ubatch, so the planner never sees a slot holding part of one, and 64 parallel 502-token embeddings into a 2012-cell pool completed in the same time with and without a park budget. The guard makes that a rule rather than a coincidence of batch sizes. --- tools/server/server-context.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7a3bdb23615d..b6089f14dfce 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3084,6 +3084,10 @@ struct server_context_impl { return; // with a cache per slot, no slot can take another one's cells } + if (!llama_get_memory(ctx_tgt)) { + return; // no cache at all (an embedding model): nothing to run out of, nothing to park + } + if (params_base.preempt_ram_mib == 0) { return; // --preempt-ram 0: the KV-full retry ladder, as before } @@ -4135,7 +4139,7 @@ struct server_context_impl { // ones back as cells free up. A multimodal prompt has no boundary the cache can name, // so it keeps the old path. bool preempt_last_resort_possible() const { - return params_base.kv_unified && params_base.preempt_ram_mib > 0 && slots.size() >= 2; + return params_base.kv_unified && params_base.preempt_ram_mib > 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt); } bool preempt_last_resort(int32_t off) { From a9b712eee6370f0c243a62399297ea415f442235 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 17:26:54 +0000 Subject: [PATCH 13/29] server: five planner accounting fixes from review The unlimited park budget, --preempt-ram -1, is documented and it turned the last resort off: the gate asked for a positive budget. It now asks for a non-zero one, with a test. A parked generating slot was charged the configured maximum draft when it came back, while its next step carries at most what its context and its prediction budget leave, the way get_n_draft_max() cuts it. Near the end of the context that over-charge could make a slot that fits look like one that never will, which since the never-fitting fix means the context error. The reservation for running slots was charged the same way. Both now use the slot's own bound. With n_cmpl > 1 the parent and its children share the prompt's cells through seq_cp, and the pool count charged the prompt once per slot. It is charged once per family now, to whichever resident member comes first, and the others only for what they generated on top of it. A slot just given a task still mirrors the previous request's prompt until the batch builder keeps the shared prefix and drops the rest; it was charged the whole stale prompt for that pass. It is charged the prefix. The context shift ran inside pre_decode(), after the planner had measured the pool with the cells the shift was about to give back. It runs before the planner now, as its own step. --- tools/server/server-context.cpp | 63 +++++++++++++++++++++---- tools/server/tests/unit/test_preempt.py | 25 ++++++++++ 2 files changed, 80 insertions(+), 8 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b6089f14dfce..d7e753a40e33 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2927,6 +2927,24 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } + // draft tokens this slot's next step can actually carry: the configured maximum, cut to + // what its context and its prediction budget leave, the way get_n_draft_max() cuts it + int32_t preempt_n_spec(const server_slot & slot) const { + int32_t res = preempt_n_spec_max(); + + if (res == 0 || !slot.task || !slot.can_speculate()) { + return 0; + } + + res = std::min(res, slot.n_ctx - slot.prompt.n_tokens() - 2); + + if (slot.n_remaining() > 0) { + res = std::min(res, slot.n_remaining() - 1); + } + + return std::max(0, res); + } + // host RAM the parked sequences hold right now size_t preempt_ram_used() const { size_t res = 0; @@ -2954,7 +2972,7 @@ struct server_context_impl { int32_t res = slot.prompt.n_tokens(); if (slot.state_before_preempt == SLOT_STATE_GENERATING) { - res += 1 + preempt_n_spec_max(); + res += 1 + preempt_n_spec(slot); } else { const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; @@ -2971,11 +2989,35 @@ struct server_context_impl { int32_t preempt_kv_used() const { int32_t res = 0; + // n_cmpl > 1: the parent and its children share the prompt's cells through seq_cp, so + // the prompt is charged once per family, to whichever resident member comes first; + // the others are charged only what they generated on top of it + std::vector charged; + for (const auto & slot : slots) { if (slot.state == SLOT_STATE_PREEMPTED) { continue; // parked: its cells are in host RAM, not in the pool } + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + const int family = slot.task->is_parent() ? slot.task->id : slot.task->id_parent; + + if (std::find(charged.begin(), charged.end(), family) != charged.end()) { + res += std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens()); + continue; + } + + charged.push_back(family); + } + + // a slot just given a task still mirrors the previous request's prompt until the + // batch builder keeps the prefix the two share and drops the rest; what stays is + // the prefix, so that is what the pool holds for it + if (slot.state == SLOT_STATE_STARTED && slot.task) { + res += (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); + continue; + } + res += slot.prompt.n_tokens(); } @@ -2984,7 +3026,6 @@ struct server_context_impl { // cells those slots are about to ask for on the next decode int32_t preempt_kv_reserve() const { - const int32_t n_spec = preempt_n_spec_max(); const int32_t n_batch = llama_n_batch(ctx_tgt); int32_t res = 0; @@ -2995,7 +3036,7 @@ struct server_context_impl { case SLOT_STATE_GENERATING: case SLOT_STATE_DONE_PROMPT: { - res += 1 + n_spec; + res += 1 + preempt_n_spec(slot); } break; case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: @@ -3307,7 +3348,9 @@ struct server_context_impl { } } - // [TAG_PREEMPT] make the pool fit the step that is about to be built + // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured after + // any context shift + pre_decode_shift(); update_preemption(); try { @@ -3392,9 +3435,11 @@ struct server_context_impl { } } - void pre_decode() { - // apply context-shift if needed - // TODO: simplify and improve + // apply context-shift if needed + // TODO: simplify and improve + // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift, + // not with the cells the shift is about to give back + void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { if (!params_base.ctx_shift) { @@ -3456,7 +3501,9 @@ struct server_context_impl { slot.truncated = true; } }); + } + void pre_decode() { // start populating the batch for this iteration batch.clear(); @@ -4139,7 +4186,7 @@ struct server_context_impl { // ones back as cells free up. A multimodal prompt has no boundary the cache can name, // so it keeps the old path. bool preempt_last_resort_possible() const { - return params_base.kv_unified && params_base.preempt_ram_mib > 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt); + return params_base.kv_unified && params_base.preempt_ram_mib != 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt); } bool preempt_last_resort(int32_t off) { diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 932c1323305b..e0d5f0c63395 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -330,6 +330,31 @@ def test_the_last_resort_parks_instead_of_ending_everyone(): assert len(res.body["tokens"]) == n_predict +def test_the_last_resort_works_with_an_unlimited_budget(): + # --preempt-ram -1 is the documented unlimited setting; it must enable the last resort + # the same as any positive budget does + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "-1" + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted as a last resort" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + + def test_the_last_resort_rewinds_a_prompt_in_flight(): # Same, with a prompt being processed when the pool runs out: the chunk that failed # is taken back off the slot's tokens and processed again after the resume, so the From 64a5064ef9a9c952c2de9dae44e9719e57d4b1ec Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 17:51:28 +0000 Subject: [PATCH 14/29] server: a resident cycling through context shifts takes turns with a parked head With context shift on and a generation with no end, a resident that has reached the pool's limit shifts, keeps about half the pool, and carries on for as long as it has tokens to make. A parked head that needs more than what is left never fits beside it, and the scheduler never parked the resident merely to let the head run, so the head waited for ever. Once the head has waited its turn (two seconds), a resident that has shifted at least once is parked in its place. It is at a token boundary like any other park, and when it comes back it is the one waiting, so the two take turns until both finish. Test: two generations of 12000 tokens in a 256-cell pool with context shift on, rotations every two seconds, both complete. --- tools/server/server-context.cpp | 44 +++++++++++++++++++++++++ tools/server/tests/unit/test_preempt.py | 29 ++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d7e753a40e33..fdaeb45a1ec5 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -94,6 +94,7 @@ static bool preempt_resume_head_of_line() { } constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked +constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long struct server_slot; // forward declaration @@ -336,6 +337,7 @@ struct server_slot { std::vector preempt_state_tgt; std::vector preempt_state_dft; int32_t n_preempt = 0; // times the CURRENT task has been preempted + int32_t n_ctx_shift = 0; // context shifts the CURRENT task has made: it is at the pool's limit and cycling int32_t n_preempt_fail = 0; // consecutive failed restores int64_t t_preempt_us = 0; // when it was parked @@ -534,6 +536,7 @@ struct server_slot { state_before_preempt = SLOT_STATE_IDLE; n_preempt = 0; n_preempt_fail = 0; + n_ctx_shift = 0; t_preempt_us = 0; llama_set_sampler(ctx_tgt, id, nullptr); @@ -3215,7 +3218,46 @@ struct server_context_impl { } } + // Nothing fits. A resident that has reached the pool's limit and is cycling + // through context shifts holds the room for as long as it likes to generate, + // and the head behind it would wait for ever. After the head has waited its + // turn, that resident is parked in its place: it is at a token boundary like + // any other park, and when it comes back it is the one waiting, so the two + // take turns instead of one taking everything. if (!best) { + server_slot * head = parked.front(); + + if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_GENERATING || slot.n_ctx_shift == 0) { + continue; + } + + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + continue; + } + + if (!preempt_fits_budget(slot) || !slot.preempt_save()) { + continue; + } + + metrics.n_preempt++; + + SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn, preemptions %d\n", + slot.n_ctx_shift, slot.prompt.n_tokens(), + slot.preempt_state_size() / (1024.0 * 1024.0), + (ggml_time_us() - head->t_preempt_us) / 1e6, + slot.n_preempt); + + best = head; // re-examined by the loop, which sees the room it just got + break; + } + } + + if (best) { + continue; + } + break; } @@ -3479,6 +3521,8 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); + slot.n_ctx_shift++; + slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index e0d5f0c63395..7a077067e214 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -392,3 +392,32 @@ def _late(n_predict, prompt): # the chunk that was in the batch given up is processed once, after the rewind, and # the count is the prompt plus the BOS the server adds assert results[1].body["timings"]["prompt_n"] == n_b + 1 + + +def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): + # Two generations that each outgrow the pool on their own, with context shift on. The + # resident reaches the limit, shifts, keeps about half the pool and would keep going + # for as long as it has tokens to make, while the parked one never fits beside it. + # After the head has waited its turn the resident is parked in its place, and the two + # take turns until both finish. Long enough that the resident is still going when the + # head's turn comes: this model makes a couple of thousand tokens a second. + global server + server.n_ctx = 256 + server.enable_ctx_shift = True + server.start() + log = LogReader(server.log_path) + + n_predict = 12000 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "slot context shift" in text + assert "rotated out after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict From 6744b3d9f2a841cdeaee1d23b7ffe7aab764ff9c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:24:12 +0000 Subject: [PATCH 15/29] server: keep only the shared prefix of a reused slot before it is sized or parked A slot just given a task still mirrors the previous request's prompt until the batch builder keeps the prefix the two share and drops the rest. The planner could size, charge and park such a slot by the old prompt: a short unrelated request could exceed the park budget or stay parked for room it would never use. The victim picker now trims a STARTED slot to the shared prefix first, which is what the batch builder does anyway. --- tools/server/server-context.cpp | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index fdaeb45a1ec5..5391272dc8d1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3060,6 +3060,26 @@ struct server_context_impl { // Keep the slot that is furthest along -- it is the closest to finishing and to giving // its cells back -- and among the rest prefer one that has not been preempted // PREEMPT_N_STARVED times already, then the smallest. + // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt + // until the batch builder keeps the prefix the two share and drops the rest (see the + // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out, + // charged and sized by the old prompt, and a short unrelated request could exceed the + // budget or stay parked for room it will never use. Keeping only the shared prefix now + // is what the batch builder does anyway; the chunk reuse it can add on top is given up + // for a slot the planner has to touch, which is rare. + void preempt_normalize_started(server_slot & slot) { + if (slot.state != SLOT_STATE_STARTED || !slot.task) { + return; + } + + const size_t n_keep = slot.prompt.tokens.get_common_prefix(slot.task->tokens); + + if (n_keep < slot.prompt.tokens.size()) { + slot.prompt.tokens.keep_first(n_keep); + slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); + } + } + server_slot * preempt_pick_victim() { server_slot * leader = nullptr; int32_t n_running = 0; @@ -3102,6 +3122,8 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + preempt_normalize_started(slot); + if (!preempt_fits_budget(slot)) { continue; } From a1c34dad206ba76db667e3a97a100dc4b58d012d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:57:05 +0000 Subject: [PATCH 16/29] server: the planner counts a reused slot from its retained prefix, charges a waiting child its own cache, and trims safely A slot just given a task still mirrors the previous request's prompt. Its next-step need subtracted the old prompt length from the new prompt length, so a shorter new prompt looked like it needed one cell while the batch was about to take a whole chunk of it; the need now counts from the prefix the two share, as the used-cell figure already did. A child of an n_cmpl > 1 request waiting for its parent's prompt does not share anything yet: until copy_state_to() runs it holds whatever the previous request left in its cells. It is charged that on its own, outside the family. The trim of a reused slot to its shared prefix is a partial removal, which a memory without room to roll back the stale suffix aborts on; for such a memory the whole stale sequence goes instead, and the prompt is processed from the start on resume. --- tools/server/server-context.cpp | 38 +++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5391272dc8d1..1b8b8375efaf 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2977,7 +2977,14 @@ struct server_context_impl { if (slot.state_before_preempt == SLOT_STATE_GENERATING) { res += 1 + preempt_n_spec(slot); } else { - const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + // a slot just given a task still mirrors the previous request's prompt; the batch + // builder keeps the prefix the two share and drops the rest, so what it holds and + // what it is about to ask for both count from that prefix, not from the old prompt + if (slot.state == SLOT_STATE_STARTED && slot.task) { + res = (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); + } + + const int32_t n_left = slot.task ? slot.task->n_tokens() - res : 0; res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); } @@ -3002,6 +3009,14 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } + // a child waiting for its parent's prompt does not share anything yet: until + // copy_state_to() runs it still holds whatever the previous request left in its + // cells, so it is charged that on its own, outside the family + if (slot.state == SLOT_STATE_WAIT_OTHER) { + res += slot.prompt.n_tokens(); + continue; + } + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { const int family = slot.task->is_parent() ? slot.task->id : slot.task->id_parent; @@ -3074,10 +3089,25 @@ struct server_context_impl { const size_t n_keep = slot.prompt.tokens.get_common_prefix(slot.task->tokens); - if (n_keep < slot.prompt.tokens.size()) { - slot.prompt.tokens.keep_first(n_keep); - slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); + if (n_keep >= slot.prompt.tokens.size()) { + return; + } + + // a memory that cannot remove part of a sequence (a recurrent state without rollback + // room for the stale suffix) aborts on a partial removal; for it the whole stale + // sequence goes, and the prompt is processed from the start on resume, as it would be + // without a usable checkpoint + const bool partial_ok = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART && + (!ctx_dft || ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART); + + if (!partial_ok) { + slot.prompt.tokens.clear(); + slot.mem.seq_rm(slot.id, -1, -1); + return; } + + slot.prompt.tokens.keep_first(n_keep); + slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); } server_slot * preempt_pick_victim() { From 6fb0b91c55c93ce999b860b3118a084e29926c7d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:13:03 +0000 Subject: [PATCH 17/29] server: LLAMA_SERVER_PREEMPT_POLICY, a test knob to compare victim choices on one workload smallest (the shipped choice), largest, youngest (the most recent task, as vLLM's scheduler preempts) and oldest. The leader is kept and the starvation guard applies under all of them; the knob exists for the comparison runs and changes nothing unless set. --- tools/server/server-context.cpp | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1b8b8375efaf..6553965dad20 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1462,6 +1462,17 @@ struct server_context_impl { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; + // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks, for comparing + // policies against each other on the same workload. smallest (the default and the + // shipped one), largest, youngest (the most recent task, as vLLM's scheduler + // preempts), oldest. The leader is kept and the starvation guard applies under all. + const char * LLAMA_SERVER_PREEMPT_POLICY = getenv("LLAMA_SERVER_PREEMPT_POLICY"); + preempt_test_policy = LLAMA_SERVER_PREEMPT_POLICY ? LLAMA_SERVER_PREEMPT_POLICY : "smallest"; + + if (preempt_test_policy != "smallest") { + SRV_WRN("LLAMA_SERVER_PREEMPT_POLICY = %s (test knob: victim choice for comparison only)\n", preempt_test_policy.c_str()); + } + if (preempt_test_every > 0) { SRV_WRN("LLAMA_SERVER_PREEMPT_EVERY = %d (test knob: preempting every slot every %d tokens)\n", preempt_test_every, preempt_test_every); @@ -2917,6 +2928,7 @@ struct server_context_impl { // shape at every step, so a preempted continuation that is not byte-identical to an // uninterrupted one is the preemption's fault and nothing else's. int32_t preempt_test_every = 0; + std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so // the KV-full retry ladder and its last resort are the only thing between a full pool @@ -3163,7 +3175,7 @@ struct server_context_impl { if (!victim || (starved_cur && !starved) || - (starved_cur == starved && slot.prompt.n_tokens() < victim->prompt.n_tokens())) { + (starved_cur == starved && preempt_better_victim(slot, *victim))) { victim = &slot; } } @@ -3171,6 +3183,25 @@ struct server_context_impl { return victim; } + // is a the better victim of the two? the smallest slot under the shipped policy: it + // gives up the least work and its restore is the cheapest (see the PR's simulation); + // the other choices exist for the comparison runs behind LLAMA_SERVER_PREEMPT_POLICY + bool preempt_better_victim(const server_slot & a, const server_slot & b) const { + if (preempt_test_policy == "largest") { + return a.prompt.n_tokens() > b.prompt.n_tokens(); + } + + if (preempt_test_policy == "youngest") { + return a.task->id > b.task->id; + } + + if (preempt_test_policy == "oldest") { + return a.task->id < b.task->id; + } + + return a.prompt.n_tokens() < b.prompt.n_tokens(); + } + // called once per update_slots(), before the batch is built: 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 from the picture without unpicking a From a7a04c2a961427b6386744dfe6f63ab120940d50 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:28:59 +0000 Subject: [PATCH 18/29] server: the rotation parks the resident that lets the head in With several shifting residents the rotation took the first in slot order, which could be one too small to make room for the head; its park could spend the budget and leave the head waiting anyway. The rotation now picks the smallest resident whose cells let the head in and, failing one that does so alone, the largest, since it makes the most room. Test: three generations with no end in a pool that holds two, all finishing their tokens. --- tools/server/server-context.cpp | 30 ++++++++++++++++++++++--- tools/server/tests/unit/test_preempt.py | 28 +++++++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6553965dad20..5b7c8731746e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3311,6 +3311,16 @@ struct server_context_impl { server_slot * head = parked.front(); if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { + // the resident whose cells let the head in, the smallest of those; failing + // one that does so alone, the largest, since it makes the most room. Taking + // the first shifting resident in slot order could park one too small to + // matter, spend the park budget on it, and leave the head waiting anyway. + const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); + const int32_t need = preempt_n_need(*head) + PREEMPT_N_MARGIN; + + server_slot * pick = nullptr; + bool pick_enough = false; + for (auto & slot : slots) { if (slot.state != SLOT_STATE_GENERATING || slot.n_ctx_shift == 0) { continue; @@ -3320,20 +3330,34 @@ struct server_context_impl { continue; } - if (!preempt_fits_budget(slot) || !slot.preempt_save()) { + if (!preempt_fits_budget(slot)) { continue; } + const bool enough = occupied - slot.prompt.n_tokens() + need <= n_cells; + + if (!pick || + (enough && !pick_enough) || + (enough == pick_enough && (enough ? slot.prompt.n_tokens() < pick->prompt.n_tokens() + : slot.prompt.n_tokens() > pick->prompt.n_tokens()))) { + pick = &slot; + pick_enough = enough; + } + } + + if (pick && pick->preempt_save()) { + server_slot & slot = *pick; + metrics.n_preempt++; - SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn, preemptions %d\n", + SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn%s, preemptions %d\n", slot.n_ctx_shift, slot.prompt.n_tokens(), slot.preempt_state_size() / (1024.0 * 1024.0), (ggml_time_us() - head->t_preempt_us) / 1e6, + pick_enough ? "" : " (not enough room by itself)", slot.n_preempt); best = head; // re-examined by the loop, which sees the room it just got - break; } } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 7a077067e214..a2fcd750f18c 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -421,3 +421,31 @@ def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_hea for res in results: assert res.status_code == 200 assert res.body["timings"]["predicted_n"] == n_predict + + +def test_the_rotation_parks_a_resident_that_lets_the_head_in(): + # Three generations with no end in a 256-cell pool with context shift on: two residents + # cycle through shifts while the third waits parked. Every rotation must let the head + # in, so all three keep finishing their tokens and no stream ends short. + global server + server.n_slots = 3 + server.n_ctx = 384 + server.enable_ctx_shift = True + server.start() + n_predict = 9000 + prompts = [ + "Once upon a time there was a brave knight who", + "The quick brown fox jumps over the lazy dog and", + "In a small village by the sea there lived a fisherman who", + ] + results = parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for p in prompts + ]) + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == n_predict + text = open(server.log_path).read() + assert "rotated out after" in text + assert "Context size has been exceeded" not in text From 64a3f6e7de3ce6779e8dcc23c8b1e04007260339 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:56:17 +0000 Subject: [PATCH 19/29] server: a parked slot survives an aborted round; a rotation counts the head's bytes as leaving abort_all_slots(), reached when pre_decode, decode or post_decode throws, released every processing slot, parked ones included, although a parked slot took no part in what failed: its sequence is in host RAM, not in the cache. It is left alone now, as the decode error sweep already did. The rotation's budget test counted the parked head's bytes against the resident about to be parked, so a budget that holds one sequence but not two refused every rotation and left the head parked for as long as the resident cared to generate. The head is restored on the same pass, so its bytes are on their way out and are no longer held against the resident. --- tools/server/server-context.cpp | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5b7c8731746e..8412d8c9ff33 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2880,7 +2880,10 @@ struct server_context_impl { void abort_all_slots(const std::string & reason) { for (auto & slot : slots) { - if (slot.is_processing()) { + // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in + // host RAM, not in the cache, and it comes back when there is room, the same as + // in the decode error sweep + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); } @@ -2982,6 +2985,22 @@ struct server_context_impl { return preempt_ram_used() + slot.preempt_state_required() <= budget; } + // the same for a rotation: the parked head is restored on the pass that parks the + // resident, so its bytes are on their way out and are not held against the resident. + // A budget that holds one sequence but not two would otherwise refuse every rotation + // and leave the head parked for as long as the resident cares to generate. + bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { + if (params_base.preempt_ram_mib < 0) { + return true; + } + + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + const size_t used = preempt_ram_used(); + const size_t leaving = std::min(used, head.preempt_state_size()); + + return used - leaving + slot.preempt_state_required() <= budget; + } + // cells the slot will ask for on its next step once it is back in the pool int32_t preempt_n_need(const server_slot & slot) const { int32_t res = slot.prompt.n_tokens(); @@ -3330,7 +3349,7 @@ struct server_context_impl { continue; } - if (!preempt_fits_budget(slot)) { + if (!preempt_fits_budget_for_rotation(slot, *head)) { continue; } From 55f04bbd689535b7809ce0b8ee3b4593256f349c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 20:58:24 +0000 Subject: [PATCH 20/29] server: the leader is measured by what a reused slot keeps, not by the prompt it still mirrors preempt_pick_victim() chose the leader by prompt.n_tokens() before the STARTED slots were trimmed, so a short request just handed a slot with a large stale cache could become the never-parked leader while the longest live conversation was parked in its place. Every STARTED slot is trimmed to its shared prefix first now; the leader and the victim see the same sizes. --- tools/server/server-context.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8412d8c9ff33..97f55b048b75 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3145,6 +3145,14 @@ struct server_context_impl { server_slot * leader = nullptr; int32_t n_running = 0; + // a slot just given a task is measured by the prefix it keeps, not by the previous + // request's prompt it still mirrors: measured by the mirror, a short request over a + // large stale cache would be the never-parked leader while the longest live + // conversation was parked in its place + for (auto & slot : slots) { + preempt_normalize_started(slot); + } + for (auto & slot : slots) { if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { n_running++; @@ -3183,8 +3191,6 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } - preempt_normalize_started(slot); - if (!preempt_fits_budget(slot)) { continue; } From 270fdd6f32f8e93d38345eefc7bb8a21ffa3617a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:54:25 +0000 Subject: [PATCH 21/29] tests: a parent and child that do not fit alone get the context error and the server lives --- tools/server/tests/unit/test_preempt.py | 33 +++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index a2fcd750f18c..a9b19fcb050a 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -449,3 +449,36 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in(): text = open(server.log_path).read() assert "rotated out after" in text assert "Context size has been exceeded" not in text + + +def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): + # One request asking for two completions is one conversation in two slots: a parent + # and a child sharing the prompt. When the two together do not fit the pool there is + # nobody else to park, since the family is charged once and a member of it is not a + # victim for the other, so the request gets the context error it would get alone, and + # the server carries on serving. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + server.start() + log = LogReader(server.log_path) + + res = server.make_request("POST", "/completion", data={ + "n_predict": 160, + "n_cmpl": 2, + "prompt": "Once upon a time there was a brave knight who", + "ignore_eos": True, + "return_tokens": True, + "temperature": 0.0, + "seed": 42, + }) + assert res.status_code == 500 + assert "Context size has been exceeded" in res.body["error"]["message"] + + text = log.drain() + assert "preempted as a last resort" not in text, "a family alone in the pool has no victim" + assert "GGML_ASSERT" not in text + + after = _complete(8) + assert after.status_code == 200 + assert after.body["timings"]["predicted_n"] == 8 From 3800ddee942169c15f9098dcaae5f37658882352 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 00:40:15 +0000 Subject: [PATCH 22/29] server: preemption is off for a recurrent cache; a rotation holds both states at once and the RAM cap says so A recurrent cache holds one state per sequence whatever its length, so the token count the planner measures says nothing about it: two long conversations were parked in turn on a cache with room to spare. Preemption is off for a pure recurrent model (llama_model_is_recurrent), said so at load, planner and last resort both; a hybrid keeps its attention cache and stays on. The rotation credited the parked head's bytes as leaving, since the head is restored on the same pass, but the resident is parked before the head is restored and freed, so both states are held at once and --preempt-ram was not a cap on what is held. The rotation now asks the plain budget check: a budget that holds one sequence but not two does not rotate, and the head says so once per park and waits for a resident to finish, or to shrink after a shift. Tests: mamba-130m under --kv-unified with the forced-park knob set, two completions finish and nothing is parked; three generations in a pool one of them fills under a budget that holds the parked heads but not a head and the resident together, the refusal logged, every stream finishes, no context error. --- tools/server/server-context.cpp | 54 ++++++++++++--------- tools/server/tests/unit/test_preempt.py | 62 +++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 22 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 97f55b048b75..a2bbc5347544 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -340,6 +340,7 @@ struct server_slot { int32_t n_ctx_shift = 0; // context shifts the CURRENT task has made: it is at the pool's limit and cycling int32_t n_preempt_fail = 0; // consecutive failed restores int64_t t_preempt_us = 0; // when it was parked + bool preempt_rotation_refused = false; // this park has logged a rotation refused for budget size_t preempt_state_size() const { return preempt_state_tgt.size() + preempt_state_dft.size(); @@ -404,6 +405,7 @@ struct server_slot { state_before_preempt = state; state = SLOT_STATE_PREEMPTED; t_preempt_us = ggml_time_us(); + preempt_rotation_refused = false; n_preempt++; @@ -1484,6 +1486,11 @@ struct server_context_impl { if (preempt_planner_off) { SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_PLANNER = off (test knob: nothing is parked ahead of the decode, only as a last resort)\n"); } + + if (llama_model_is_recurrent(model_tgt)) { + preempt_recurrent = true; + SRV_WRN("%s", "preemption: off, the recurrent cache holds one state per sequence whatever its length, so there is no cell pool to run out of\n"); + } } { @@ -2938,6 +2945,11 @@ struct server_context_impl { // and the context error bool preempt_planner_off = false; + // a recurrent cache holds one state per sequence whatever its length: no cell pool, + // nothing to run out of, and the token count the planner measures says nothing about + // it. Preemption is off for those models; a hybrid keeps its attention cache and stays on. + bool preempt_recurrent = false; + // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop bool preempt_batch_abandoned = false; @@ -2985,22 +2997,6 @@ struct server_context_impl { return preempt_ram_used() + slot.preempt_state_required() <= budget; } - // the same for a rotation: the parked head is restored on the pass that parks the - // resident, so its bytes are on their way out and are not held against the resident. - // A budget that holds one sequence but not two would otherwise refuse every rotation - // and leave the head parked for as long as the resident cares to generate. - bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { - if (params_base.preempt_ram_mib < 0) { - return true; - } - - const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - const size_t used = preempt_ram_used(); - const size_t leaving = std::min(used, head.preempt_state_size()); - - return used - leaving + slot.preempt_state_required() <= budget; - } - // cells the slot will ask for on its next step once it is back in the pool int32_t preempt_n_need(const server_slot & slot) const { int32_t res = slot.prompt.n_tokens(); @@ -3240,8 +3236,8 @@ struct server_context_impl { return; // no cache at all (an embedding model): nothing to run out of, nothing to park } - if (params_base.preempt_ram_mib == 0) { - return; // --preempt-ram 0: the KV-full retry ladder, as before + if (params_base.preempt_ram_mib == 0 || preempt_recurrent) { + return; // --preempt-ram 0, or a recurrent cache: the KV-full retry ladder, as before } const int32_t n_cells = n_ctx; @@ -3343,8 +3339,9 @@ struct server_context_impl { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); const int32_t need = preempt_n_need(*head) + PREEMPT_N_MARGIN; - server_slot * pick = nullptr; - bool pick_enough = false; + server_slot * pick = nullptr; + bool pick_enough = false; + bool budget_refused = false; for (auto & slot : slots) { if (slot.state != SLOT_STATE_GENERATING || slot.n_ctx_shift == 0) { @@ -3355,7 +3352,13 @@ struct server_context_impl { continue; } - if (!preempt_fits_budget_for_rotation(slot, *head)) { + // The head's own bytes are not credited as leaving: the resident is + // parked before the head is restored and freed, so both states are + // held at once, and the cap is a cap on what is held. A budget that + // holds one sequence but not two does not rotate, and the head waits + // for a resident to finish, which is said once per park below. + if (!preempt_fits_budget(slot)) { + budget_refused = true; continue; } @@ -3370,6 +3373,13 @@ struct server_context_impl { } } + if (!pick && budget_refused && !head->preempt_rotation_refused) { + head->preempt_rotation_refused = true; + + SLT_WRN(*head, "no rotation: --preempt-ram %d MiB does not hold this parked state and a resident's at once, and the two are held together while the resident is parked and the head restored; the head waits for a resident to finish\n", + params_base.preempt_ram_mib); + } + if (pick && pick->preempt_save()) { server_slot & slot = *pick; @@ -4362,7 +4372,7 @@ struct server_context_impl { // ones back as cells free up. A multimodal prompt has no boundary the cache can name, // so it keeps the old path. bool preempt_last_resort_possible() const { - return params_base.kv_unified && params_base.preempt_ram_mib != 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt); + return params_base.kv_unified && params_base.preempt_ram_mib != 0 && !preempt_recurrent && slots.size() >= 2 && llama_get_memory(ctx_tgt); } bool preempt_last_resort(int32_t off) { diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index a9b19fcb050a..674bf0a9d461 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -482,3 +482,65 @@ def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_ after = _complete(8) assert after.status_code == 200 assert after.body["timings"]["predicted_n"] == 8 + + +def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_when_a_resident_finishes(): + # Three generations with no end in a pool one of them fills, with context shift on, + # under a --preempt-ram that holds the two parked heads but not a head and the resident + # at once. The resident is parked before the head is restored and freed, so a rotation + # holds both states together: under this budget the first one asked for is refused and + # said so, and the heads come back when the resident finishes instead. Every stream + # still finishes its tokens and nothing gets the context error. + global server + server.n_slots = 3 + server.n_ctx = 2048 + server.enable_ctx_shift = True + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" + server.start() + n_predict = 6000 + prompts = [ + "Once upon a time there was a brave knight who", + "The quick brown fox jumps over the lazy dog and", + "In a small village by the sea there lived a fisherman who", + ] + results = parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for p in prompts + ]) + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == n_predict + text = open(server.log_path).read() + assert "no rotation: --preempt-ram 2 MiB" in text + assert "resumed after" in text + assert "Context size has been exceeded" not in text + + +def test_a_recurrent_model_is_served_without_preemption(): + # A recurrent cache holds one state per sequence whatever its length, so the token + # count the planner measures says nothing about it: preemption is off for such a + # model, said so at load, and the forced-park knob parks nothing. + global server + path = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") + if path: + server.model_file = path + else: + server.model_file = None + server.model_hf_repo = "Felladrin/gguf-mamba-130m-hf" + server.model_hf_file = "mamba-130m-hf.Q2_K.gguf" + server.offline = False + server.n_ctx = 1024 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start(timeout_seconds=300) + results = parallel_function_calls([ + (_complete, (64, "Once upon a time")), + (_complete, (64, "The quick brown fox")), + ]) + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 64 + text = open(server.log_path).read() + assert "preemption: off, the recurrent cache holds one state per sequence" in text + assert "preempted" not in text + assert "Context size has been exceeded" not in text From 81eec0bb3bf2783247150038be4db7e18b6107df Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:06:19 +0000 Subject: [PATCH 23/29] server: a started slot's reservation counts from the prefix it keeps, as its used count already did preempt_kv_reserve() measured a started slot's next prompt chunk from the prompt it still mirrors (task tokens minus prompt tokens), so a request shorter than the previous one on that slot reserved one cell for a chunk of hundreds, while preempt_kv_used() and preempt_n_need() already counted from the retained prefix. The three now share preempt_n_retained(): the common prefix for a started slot, nothing when the request does not cache its prompt, the prompt otherwise. The new chunk always fits where the mirrored prompt was, so the undercount could only bite a restore decided in the same pass, and the retry ladder hid it for ordinary requests; the figures agree now either way. --- tools/server/server-context.cpp | 40 ++++++++++++++++++--------------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a2bbc5347544..a53b51b1d59d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2997,20 +2997,29 @@ struct server_context_impl { return preempt_ram_used() + slot.preempt_state_required() <= budget; } + // cells of the slot's that its next step keeps: a slot just given a task still mirrors + // the previous request's prompt until the batch builder keeps the prefix the two share + // and drops the rest (all of it when the request does not cache its prompt), so what it + // holds, and what it is about to ask for, both count from that prefix + int32_t preempt_n_retained(const server_slot & slot) const { + if (slot.state == SLOT_STATE_STARTED && slot.task) { + if (!slot.task->params.cache_prompt) { + return 0; + } + + return (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); + } + + return slot.prompt.n_tokens(); + } + // cells the slot will ask for on its next step once it is back in the pool int32_t preempt_n_need(const server_slot & slot) const { - int32_t res = slot.prompt.n_tokens(); + int32_t res = preempt_n_retained(slot); if (slot.state_before_preempt == SLOT_STATE_GENERATING) { res += 1 + preempt_n_spec(slot); } else { - // a slot just given a task still mirrors the previous request's prompt; the batch - // builder keeps the prefix the two share and drops the rest, so what it holds and - // what it is about to ask for both count from that prefix, not from the old prompt - if (slot.state == SLOT_STATE_STARTED && slot.task) { - res = (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); - } - const int32_t n_left = slot.task ? slot.task->n_tokens() - res : 0; res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); @@ -3055,15 +3064,7 @@ struct server_context_impl { charged.push_back(family); } - // a slot just given a task still mirrors the previous request's prompt until the - // batch builder keeps the prefix the two share and drops the rest; what stays is - // the prefix, so that is what the pool holds for it - if (slot.state == SLOT_STATE_STARTED && slot.task) { - res += (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); - continue; - } - - res += slot.prompt.n_tokens(); + res += preempt_n_retained(slot); } return res; @@ -3086,7 +3087,10 @@ struct server_context_impl { case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { - const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + // from the prefix a started slot keeps, not from the prompt it still + // mirrors: measured by the mirror, a request shorter than the last one + // reserved one cell for a chunk of hundreds + const int32_t n_left = slot.task ? slot.task->n_tokens() - preempt_n_retained(slot) : 0; res_pmt += std::max(1, std::min(n_batch, n_left)); } break; From af560907bf2cdb42582e1dd9c7c57067af12ee21 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:18:31 +0000 Subject: [PATCH 24/29] server: the context shift and the planner run inside the guarded part of the step Both ran ahead of the try block around pre_decode() and batch.render(), so a shift that failed to rebuild a slot's tokens, or a park that failed to allocate, left update_slots() on an uncaught exception and ended the loop instead of telling the slots. They run inside the same guard now, and the existing handler ends the affected slots with the message; a parked slot took no part in what failed and is spared, as before. --- tools/server/server-context.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a53b51b1d59d..c1474940e768 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3536,12 +3536,14 @@ struct server_context_impl { } } - // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured after - // any context shift - pre_decode_shift(); - update_preemption(); - try { + // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured + // after any context shift. Inside the guard with the rest of the step: a shift + // rebuilds a slot's tokens and a park allocates, and either can throw, which the + // slots are told about rather than the loop ending on an uncaught exception + pre_decode_shift(); + update_preemption(); + scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); batch.render(); From 00b27d2dfc8fb1637b35a874c42591ddf467915e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:44:35 +0000 Subject: [PATCH 25/29] server: what a started slot keeps is decided by one rule, the batch builder's The retained-prefix figure and the started-slot normalisation each took the common prefix whole, while the batch builder keeps none of it when the request does not cache its prompt and cuts it short of an aLoRA invocation. Both now ask preempt_n_keep(), which applies the same rule, so the planner counts, reserves and parks a started slot by the state the request will actually keep. --- tools/server/server-context.cpp | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c1474940e768..97ae8884146e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2997,17 +2997,30 @@ struct server_context_impl { return preempt_ram_used() + slot.preempt_state_required() <= budget; } + // cells of the mirrored prompt that a started slot's request keeps, by the rule the batch + // builder applies when it takes the slot: nothing when the request does not cache its + // prompt, otherwise the prefix the two share, cut short of an aLoRA invocation + size_t preempt_n_keep(const server_slot & slot) const { + if (!slot.task->params.cache_prompt) { + return 0; + } + + size_t n_keep = slot.prompt.tokens.get_common_prefix(slot.task->tokens); + + if (slot.alora_invocation_start > 0) { + n_keep = std::min(n_keep, (size_t) (slot.alora_invocation_start - 1)); + } + + return n_keep; + } + // cells of the slot's that its next step keeps: a slot just given a task still mirrors - // the previous request's prompt until the batch builder keeps the prefix the two share - // and drops the rest (all of it when the request does not cache its prompt), so what it - // holds, and what it is about to ask for, both count from that prefix + // the previous request's prompt until the batch builder keeps what preempt_n_keep() + // says and drops the rest, so what it holds, and what it is about to ask for, both + // count from that int32_t preempt_n_retained(const server_slot & slot) const { if (slot.state == SLOT_STATE_STARTED && slot.task) { - if (!slot.task->params.cache_prompt) { - return 0; - } - - return (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); + return (int32_t) preempt_n_keep(slot); } return slot.prompt.n_tokens(); @@ -3118,7 +3131,7 @@ struct server_context_impl { return; } - const size_t n_keep = slot.prompt.tokens.get_common_prefix(slot.task->tokens); + const size_t n_keep = preempt_n_keep(slot); if (n_keep >= slot.prompt.tokens.size()) { return; From ebfe23bb176418741f481f9f7b5b7f82d96c7d3d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:30:10 +0000 Subject: [PATCH 26/29] server: the resume order is read from the environment on every load The flag is process-global and was set only when LLAMA_SERVER_PREEMPT_RESUME was pass, so a reload with the variable unset, or another context loaded in the same process, kept the previous load's order. It is reset to the default and read again on every load. --- tools/server/server-context.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 97ae8884146e..883404311625 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1450,6 +1450,10 @@ struct server_context_impl { } { + // read on every load, so a reload after the variable changed, or another context + // loaded in the same process, gets its own order rather than the previous one's + g_preempt_resume_head_of_line = true; + const char * LLAMA_SERVER_PREEMPT_RESUME = getenv("LLAMA_SERVER_PREEMPT_RESUME"); if (LLAMA_SERVER_PREEMPT_RESUME && strcmp(LLAMA_SERVER_PREEMPT_RESUME, "head") != 0) { if (strcmp(LLAMA_SERVER_PREEMPT_RESUME, "pass") != 0) { From ad89538d23ef400abca5b94b50027788b5e97a0f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:38:12 +0000 Subject: [PATCH 27/29] tests: the rotation-budget test keeps its resident cycling past the rotation delay At 6000 tokens the resident finished in under three seconds on a fast host, before the parked heads had waited the two seconds a rotation needs, so nothing was ever asked for or refused and the test failed on timing. --- tools/server/tests/unit/test_preempt.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 674bf0a9d461..0c8f5dc5f291 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -497,7 +497,10 @@ def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_w server.enable_ctx_shift = True os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" server.start() - n_predict = 6000 + # long enough that the resident is still cycling through shifts two seconds after the + # heads were parked, which is when a rotation is first asked for: at 6000 this model + # finished in under three seconds on a fast host and nothing was ever refused + n_predict = 12000 prompts = [ "Once upon a time there was a brave knight who", "The quick brown fox jumps over the lazy dog and", From 301caa31da27f641aefec3f0f86265e399e36485 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:56:11 +0000 Subject: [PATCH 28/29] server: the resume order belongs to the context; the recurrent flag is assigned on every load The resume order was a file-static flag, so two contexts in one process shared it and the later load decided for both. It is a member now, read from the environment at load. And preempt_recurrent was set and never cleared, so a context reloaded with an attention model after a recurrent one kept preemption off; it is assigned from the model on every load. --- tools/server/server-context.cpp | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 883404311625..ef564b034ac4 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -87,11 +87,6 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is // LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest // parked, and a smaller slot may pass a head that does not fit. // LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. -static bool g_preempt_resume_head_of_line = true; - -static bool preempt_resume_head_of_line() { - return g_preempt_resume_head_of_line; -} constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long @@ -1450,9 +1445,9 @@ struct server_context_impl { } { - // read on every load, so a reload after the variable changed, or another context - // loaded in the same process, gets its own order rather than the previous one's - g_preempt_resume_head_of_line = true; + // read on every load and kept on this context, so a reload after the variable + // changed, or another context loaded in the same process, has an order of its own + preempt_resume_head = true; const char * LLAMA_SERVER_PREEMPT_RESUME = getenv("LLAMA_SERVER_PREEMPT_RESUME"); if (LLAMA_SERVER_PREEMPT_RESUME && strcmp(LLAMA_SERVER_PREEMPT_RESUME, "head") != 0) { @@ -1461,7 +1456,7 @@ struct server_context_impl { LLAMA_SERVER_PREEMPT_RESUME); return false; } - g_preempt_resume_head_of_line = false; + preempt_resume_head = false; SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_RESUME = pass (parked slots come back most-preempted first, and a smaller slot may pass a head that does not fit)\n"); } @@ -1491,8 +1486,11 @@ struct server_context_impl { SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_PLANNER = off (test knob: nothing is parked ahead of the decode, only as a last resort)\n"); } - if (llama_model_is_recurrent(model_tgt)) { - preempt_recurrent = true; + // assigned, not only set: the same context reloaded with an attention model after + // a recurrent one gets its preemption back + preempt_recurrent = llama_model_is_recurrent(model_tgt); + + if (preempt_recurrent) { SRV_WRN("%s", "preemption: off, the recurrent cache holds one state per sequence whatever its length, so there is no cell pool to run out of\n"); } } @@ -2949,6 +2947,11 @@ struct server_context_impl { // and the context error bool preempt_planner_off = false; + // LLAMA_SERVER_PREEMPT_RESUME: head (the default) puts parked slots back in the order they + // were parked and only the first until it fits; pass lets a smaller slot pass a head + // that does not fit. Read at load, per context. + bool preempt_resume_head = true; + // a recurrent cache holds one state per sequence whatever its length: no cell pool, // nothing to run out of, and the token count the planner measures says nothing about // it. Preemption is off for those models; a hybrid keeps its attention cache and stays on. @@ -3263,11 +3266,11 @@ struct server_context_impl { const int32_t n_cells = n_ctx; - // Put back what fits, in the order preempt_resume_head_of_line() describes: by default + // Put back what fits, in the order preempt_resume_head describes: by default // the slot parked longest, and only that one until it fits; under // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked // longest, and a smaller slot may pass a head that does not fit. - const bool head_of_line = preempt_resume_head_of_line(); + const bool head_of_line = preempt_resume_head; for (;;) { std::vector parked; From 3c99fafdf03619b51ad80ee55705ec52fb765a7c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 03:26:35 +0000 Subject: [PATCH 29/29] server : count a started slot by the cells it holds and trim it before a resume A slot just given a task keeps the previous request's prompt in the pool until the batch builder trims it to the shared prefix, which with continuous batching off can be well behind a running generation. The planner counted such a slot by that prefix, so a resume could be found to fit and attempted against cells still occupied, and the parked request then failed its 60 s wait. The pool is now measured by what each slot physically holds. When nothing fits, every started slot is trimmed to the prefix its request keeps before idle slots are cleared, so the cells it will not use are released ahead of the batch builder and a parked slot that fits without them comes back. --- tools/server/server-context.cpp | 44 +++++++++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index ef564b034ac4..b18fa4e2f23a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3084,7 +3084,13 @@ struct server_context_impl { charged.push_back(family); } - res += preempt_n_retained(slot); + // what the pool holds now, the previous request's prompt included for a slot just + // given a task: the batch builder trims that to the prefix the two share, but + // not until the slot is built into a batch, and with continuous batching off that + // can be a long time behind a running generation. Measured by the prefix, a + // restore was found to fit and attempted against cells still occupied. Under + // pressure the planner trims such slots itself, see preempt_normalize_started_all() + res += slot.prompt.n_tokens(); } return res; @@ -3133,6 +3139,29 @@ struct server_context_impl { // budget or stay parked for room it will never use. Keeping only the shared prefix now // is what the batch builder does anyway; the chunk reuse it can add on top is given up // for a slot the planner has to touch, which is rare. + // every started slot, when the pool is short: true when any of them gave cells up + bool preempt_normalize_started_all() { + bool res = false; + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_STARTED || !slot.task) { + continue; + } + + const int32_t before = slot.prompt.n_tokens(); + + preempt_normalize_started(slot); + + if (slot.prompt.n_tokens() < before) { + SLT_INF(slot, "trimmed to the %d cells its request keeps ahead of the batch builder, %d released\n", + slot.prompt.n_tokens(), before - slot.prompt.n_tokens()); + res = true; + } + } + + return res; + } + void preempt_normalize_started(server_slot & slot) { if (slot.state != SLOT_STATE_STARTED || !slot.task) { return; @@ -3341,7 +3370,18 @@ struct server_context_impl { } } - if (best || !try_clear_idle_slots()) { + if (best) { + break; + } + + // a slot just given a task still holds the previous request's prompt until + // the batch builder trims it; trimmed here instead, the cells it will not + // keep are counted out and a parked slot that fits without them comes back + if (preempt_normalize_started_all()) { + continue; + } + + if (!try_clear_idle_slots()) { break; } }