From 5cc9776c71221ec673dfca5298e19e8ee8982683 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 04:26:46 +0000 Subject: [PATCH 1/2] server: preempt a slot by dropping its cells and recomputing it later 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, the retry ladder halves n_batch and the server ends every conversation in flight with "Context size has been exceeded", including the ones nowhere near their own limit. This parks one slot instead. Its cells are dropped with llama_memory_seq_rm and nothing is copied out: the two things needed to rebuild the sequence are already in RAM, its token list (prompt.tokens, prompt plus everything generated) and its sampler. When the pool has room the slot re-enters the ordinary prompt-processing path, re-prefills its whole sequence and carries on sampling with the same sampler, from the same token. A streaming client sees a pause, not an error. No host RAM is held while it waits. The resume for a generating slot replays prompt.tokens plus `sampled`, the token it had sampled but not yet decoded, so the logits land at exactly the position the interrupted step was about to read them from and the DONE_PROMPT transition samples the token that step would have sampled. A slot parked while still processing its prompt starts its own prompt over and needs no replay list. usage.prompt_tokens stays task->n_tokens() and never grows; timings.prompt_n accumulates the recomputed tokens, which is what it counts, and /slots reports the split as n_recompute. The sampler is never reset on a resume, so penalties, grammar and RNG carry across the pause untouched. Policy: idle finished slots are purged first, the slot with the most tokens is never preempted, the smallest of the rest is parked, a slot parked three times is passed over while another candidate exists, resume is most-preempted first then longest parked with fit-first, prompt-phase slots are victims too, n_cmpl > 1 and multimodal slots are not. --preempt / --no-preempt (env LLAMA_ARG_PREEMPT), on by default and active only with --kv-unified and more than one slot. /metrics gains n_preempt_total, n_resume_total, n_recompute_tokens_total and requests_preempted; /slots gains is_preempted, n_preempt and n_recompute. LLAMA_SERVER_PREEMPT_EVERY=N forces a park every N generated tokens for the exactness test. tools/server/tests/unit/test_preempt_recompute.py covers the forced-preemption identity, the usage accounting, two generating slots and two prompt slots that overflow together, a generating slot beside a large prompt, --no-preempt, and the metrics. --- common/arg.cpp | 9 + common/common.h | 1 + tools/server/README.md | 5 + tools/server/server-common.h | 6 + tools/server/server-context.cpp | 521 ++++++++++++++++-- tools/server/server-task.cpp | 16 + tools/server/server-task.h | 1 + .../tests/unit/test_preempt_recompute.py | 293 ++++++++++ 8 files changed, 821 insertions(+), 31 deletions(-) create mode 100644 tools/server/tests/unit/test_preempt_recompute.py diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d0..b71f8b2c7630 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1709,6 +1709,15 @@ 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"}, + {"--no-preempt"}, + string_format("with a unified KV cache and more than one slot, park a running request and recompute it later " + "instead of failing every request when the cache fills (default: %s)", params.preempt ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.preempt = value; + } + ).set_env("LLAMA_ARG_PREEMPT").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..3656aa0d929b 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. + bool preempt = true; // with a unified KV cache, park a slot and recompute it later instead of failing every slot 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..3ef17ee8d9e4 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, --no-preempt` | with a unified KV cache and more than one slot, park a running request and recompute it later instead of failing every request when the cache fills (default: enabled)
(env: LLAMA_ARG_PREEMPT) | | `-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) | @@ -1138,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 sent back through prompt processing. | +| `llamacpp:n_recompute_tokens_total` | Counter | Prompt tokens re-processed to put parked slots back. | +| `llamacpp:requests_preempted` | Gauge | Requests currently parked, waiting for room in the unified KV cache. | ### 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..b44c49863d63 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,6 +467,12 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; + // [TAG_PREEMPT_RC] slots parked to make room in the unified KV pool, put back, and the + // prompt tokens that had to be run through the model a second time to put them back + uint64_t n_preempt = 0; + uint64_t n_resume = 0; + uint64_t n_recompute_tokens = 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 a9edbd7be8b4..0384261529f8 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -59,8 +60,31 @@ enum slot_state { SLOT_STATE_PROCESSING_PROMPT, SLOT_STATE_DONE_PROMPT, SLOT_STATE_GENERATING, + SLOT_STATE_PREEMPTED, // [TAG_PREEMPT_RC] cells dropped, the token list and the sampler are all that is left }; +// [TAG_PREEMPT_RC] server-side request preemption by recompute +// +// 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. One slot's cells are dropped and given back to it later. Nothing is +// copied out: the two things needed to rebuild the sequence are already in RAM, its token list +// (prompt.tokens, which is the prompt plus everything it has generated) and its sampler. When the +// pool has room the slot re-enters the ordinary prompt-processing path, re-prefills its whole +// sequence and carries on sampling with the same sampler, from the same token. A streaming client +// sees a pause, not an error. +// +// This is vLLM's default preemption mode. The trade against copying the sequence to host RAM +// (PR #184) is host memory against recomputed prefill, plus one thing that copying does not have: +// prefill and decode are different kernels, so a recomputed logit is not bit-identical to the one +// the interrupted decode would have produced, and a greedy argmax can flip on a near tie. +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 + struct server_slot; // forward declaration struct server_batch { @@ -293,6 +317,75 @@ struct server_slot { prompt.clear(); } + // [TAG_PREEMPT_RC] state of a slot whose cells were dropped + // + // preempt_replay is the token list the slot has to run through the model again to get back + // where it was. It is empty for a slot that had not finished its prompt: that one simply + // starts its own prompt over, so task->tokens is already the right list. + slot_state state_before_preempt = SLOT_STATE_IDLE; + server_tokens preempt_replay; + bool preempt_resuming = false; // re-prefilling right now + int32_t n_preempt = 0; // times the CURRENT task has been parked + int32_t n_recompute = 0; // tokens the CURRENT task has had re-processed + int64_t t_preempt_us = 0; // when it was parked + + // What the prompt path must process for this slot. Only the prompt path uses these; every + // reader that reports to the client (usage, progress totals, the context-overflow error) + // deliberately keeps reading task->n_tokens(), so a recompute never double counts a prompt. + const server_tokens & input_tokens() const { + return preempt_replay.empty() ? task->tokens : preempt_replay; + } + + int32_t n_input_tokens() const { + return preempt_replay.empty() ? task->n_tokens() : (int32_t) preempt_replay.size(); + } + + // drop the cells and keep the token list + void preempt_drop() { + if (state == SLOT_STATE_GENERATING) { + // The cache holds prompt.tokens; `sampled` is the token this slot sampled last step + // and has not decoded yet -- handle_last_sampled_token() would have added it to the + // batch. Replaying prompt.tokens + sampled puts the logits back at exactly the + // position the interrupted step was about to read them from, so the DONE_PROMPT -> + // GENERATING transition samples the very token that step would have sampled. + llama_tokens replay = prompt.tokens.get_text_tokens(); + + replay.push_back(sampled); + + preempt_replay = server_tokens(replay, false); + } + // else: a slot part-way through its prompt starts that prompt over. If it is a slot that + // was already resuming, preempt_replay is left as it is and it starts the replay + // over instead -- clearing it here would lose everything it had generated. + + // 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 matches + // the cache, including the rollback the checkpoint path does on a partial acceptance. + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + + // mem is a common_memory over ctx_tgt AND ctx_dft, so this takes the draft context's + // sequence with it, and prompt.clear() drops the context checkpoints that pointed at + // cells which no longer exist + prompt_clear(); + + state_before_preempt = state; + state = SLOT_STATE_PREEMPTED; + t_preempt_us = ggml_time_us(); + + n_preempt++; + } + + // send the slot back through prompt processing; it cannot fail, there is nothing to put back + void preempt_resume() { + preempt_resuming = true; + state = SLOT_STATE_STARTED; + } + std::vector lora; int32_t alora_invocation_start = -1; @@ -351,6 +444,14 @@ struct server_slot { n_predict_max = -1; + // [TAG_PREEMPT_RC] + preempt_replay.clear(); + state_before_preempt = SLOT_STATE_IDLE; + preempt_resuming = false; + n_preempt = 0; + n_recompute = 0; + t_preempt_us = 0; + llama_set_sampler(ctx_tgt, id, nullptr); // clear alora start @@ -503,6 +604,15 @@ struct server_slot { t_last_used = ggml_time_us(); + // [TAG_PREEMPT_RC] a parked slot can be released (cancelled or failed) while its cells + // are gone; drop the replay list with it so the next task on this slot cannot take a + // prefix match against a sequence that is not in the cache + if (state == SLOT_STATE_PREEMPTED) { + preempt_replay.clear(); + preempt_resuming = false; + prompt_clear(); + } + state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough @@ -577,7 +687,7 @@ struct server_slot { } const double n_prompt_second = stats.n_prompt_tps(); - const double f_progress = task->n_tokens() > 0 ? (double) prompt.n_tokens() / task->n_tokens() : 0.0; + const double f_progress = n_input_tokens() > 0 ? (double) prompt.n_tokens() / n_input_tokens() : 0.0; SLT_INF(*this, "prompt processing, n_tokens = %6d, progress = %.2f, t = %6.2f s / %.2f tokens per second\n", (int) stats.n_prompt_processed, f_progress, t_prompt_total / 1e3, n_prompt_second); @@ -645,6 +755,9 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, + {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"n_preempt", n_preempt}, + {"n_recompute", n_recompute}, }; const auto & ptask = task ? task : task_prev; @@ -1249,6 +1362,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: parking every slot every %d generated 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; @@ -2385,11 +2508,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); @@ -2397,6 +2524,7 @@ 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->metrics = metrics; if (task.metrics_reset_bucket) { @@ -2674,6 +2802,286 @@ struct server_context_impl { }; #endif + + // + // [TAG_PREEMPT_RC] server-side request preemption by recompute + // + + // LLAMA_SERVER_PREEMPT_EVERY=N parks 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 + // anything that differs between a run with the knob and a run without it is the pause. + int32_t preempt_test_every = 0; + + bool preempt_enabled() const { + // with a cache per slot no slot can take another one's cells, and with one slot there is + // no one to take them from + return params_base.preempt && params_base.kv_unified && slots.size() > 1; + } + + 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. A parked slot holds nothing: + // its prompt was cleared with its cells. + int32_t preempt_kv_used() const { + int32_t res = 0; + + for (const auto & slot : slots) { + res += slot.prompt.n_tokens(); + } + + return res; + } + + // cells a parked slot needs before it is worth waking: its whole sequence has to go back + // through the model, and then it takes its next step + int32_t preempt_n_need(const server_slot & slot) const { + return slot.n_input_tokens() + 1 + preempt_n_spec_max(); + } + + // cells the slots are about to ask for + 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.n_input_tokens() - slot.prompt.n_tokens(); + + if (slot.preempt_resuming) { + // Already committed: this slot was woken because its WHOLE sequence + // fitted, and it gets there one batch per iteration. Charging it a + // single batch like an ordinary prompt would let the next iteration + // wake a second slot into cells the first one has not claimed yet, + // and both would be parked again a few hundred tokens later. + res += std::max(1, n_left) + 1 + n_spec; + } else { + res_pmt += std::max(1, std::min(n_batch, n_left)); + } + } break; + default: + break; + } + } + + // one batch is all the fresh prompt slots get between them, however many are waiting + return res + std::min(res_pmt, n_batch); + } + + // a slot whose sequence cannot be rebuilt from its token list alone is not a candidate + bool preempt_can_park(const server_slot & slot) const { + if (!slot.task) { + return false; + } + + if (slot.task->is_parent() || slot.task->is_child()) { + return false; // n_cmpl > 1 slots share one sequence through seq_cp, out of scope here + } + + if (slot.task->tokens.has_mtmd || slot.prompt.tokens.has_mtmd) { + // the token list holds placeholders for the media chunks, not the chunks; a replay + // built from it would prefill the wrong thing + return false; + } + + return true; + } + + // Keep the slot that is furthest along -- it is the closest to finishing and to giving its + // cells back -- and among the rest prefer, in order: one that is not part-way through a + // recompute already, one that has not been parked PREEMPT_N_STARVED times, and then the + // smallest, which is the least work to redo. + 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; + + // lower is better + auto rank = [](const server_slot * slot) { + return std::make_tuple(slot->preempt_resuming ? 1 : 0, + slot->n_preempt >= PREEMPT_N_STARVED ? 1 : 0, + slot->prompt.n_tokens()); + }; + + for (auto & slot : slots) { + // 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. + if (slot.state != SLOT_STATE_GENERATING && + slot.state != SLOT_STATE_PROCESSING_PROMPT && + slot.state != SLOT_STATE_STARTED) { + continue; + } + + if (&slot == leader) { + continue; + } + + if (!preempt_can_park(slot)) { + continue; + } + + if (!victim || rank(&slot) < rank(victim)) { + 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 (!preempt_enabled()) { + return; + } + + 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. + for (;;) { + std::vector parked; + + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTED) { + parked.push_back(&slot); + } + } + + if (parked.empty()) { + break; + } + + 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; + } + + return a->t_preempt_us < b->t_preempt_us; + }); + + server_slot * best = nullptr; + + // Room for the whole 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 - the same call the KV-full path makes. + for (;;) { + for (auto * slot : parked) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + best = slot; + break; + } + } + + if (best || !try_clear_idle_slots()) { + break; + } + } + + if (!best) { + break; + } + + best->preempt_resume(); + + metrics.n_resume++; + + SLT_WRN(*best, "resuming after %.2f s parked: %d tokens to recompute, kv %d/%d, preemptions %d\n", + (ggml_time_us() - best->t_preempt_us) / 1e6, + best->n_input_tokens(), + preempt_kv_used(), n_cells, + best->n_preempt); + } + + // forced preemption, for the exactness 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 && + preempt_can_park(slot)) { + const int32_t n_tokens = slot.prompt.n_tokens(); + + slot.preempt_drop(); + + metrics.n_preempt++; + + SLT_WRN(slot, "preempted on request after %d generated tokens, %d cells dropped\n", + (int32_t) slot.stats.n_gen, n_tokens); + } + } + } + + // 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(); + + victim->preempt_drop(); + + metrics.n_preempt++; + + SLT_WRN(*victim, "preempted: %d cells dropped, %d tokens to recompute later, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + victim->n_input_tokens(), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + } + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -2715,6 +3123,9 @@ struct server_context_impl { } } + // [TAG_PREEMPT_RC] 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(); @@ -3000,7 +3411,9 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - if (!slot.is_processing()) { + // [TAG_PREEMPT_RC] a parked slot is processing but has nothing in the cache to + // batch; it takes no part in this pass until it is resumed + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { return; } @@ -3017,19 +3430,23 @@ struct server_context_impl { // this slot still has a prompt to be processed if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) { - const auto & input_tokens = slot.task->tokens; + const auto & input_tokens = slot.input_tokens(); // used to determine the number of tokens added to the batch for the current slot const auto n_tokens_prev = batch.size(); // TODO: maybe move branch to outside of this loop in the future if (slot.state == SLOT_STATE_STARTED) { - slot.stats.update_prompt_start(); + // [TAG_PREEMPT_RC] a resumed slot started long ago: update_prompt_start() + // asserts t_start == 0 and would abort the server on the second pass + if (slot.stats.t_start == 0) { + slot.stats.update_prompt_start(); + } slot.state = SLOT_STATE_PROCESSING_PROMPT; SLT_TRC(slot, "new prompt, n_ctx_slot = %d, n_keep = %d, task.n_tokens = %d\n", - slot.n_ctx, slot.task->params.n_keep, slot.task->n_tokens()); + slot.n_ctx, slot.task->params.n_keep, slot.n_input_tokens()); // print prompt tokens (for debugging) /*if (1) { @@ -3066,33 +3483,33 @@ struct server_context_impl { } if (!slot.can_split()) { - if (slot.task->n_tokens() > n_ubatch) { + if (slot.n_input_tokens() > n_ubatch) { send_error(slot, string_format( "input (%d tokens) is too large to process. increase the physical batch " "size (current batch size: %d)", - slot.task->n_tokens(), n_ubatch), + slot.n_input_tokens(), n_ubatch), ERROR_TYPE_SERVER); slot.release(); return; } - if (slot.task->n_tokens() > slot.n_ctx) { + if (slot.n_input_tokens() > slot.n_ctx) { send_error( slot, string_format( "input (%d tokens) is larger than the max context size (%d tokens). skipping", - slot.task->n_tokens(), slot.n_ctx), + slot.n_input_tokens(), slot.n_ctx), ERROR_TYPE_EXCEED_CONTEXT_SIZE); slot.release(); return; } } else { - if (slot.task->n_tokens() >= slot.n_ctx) { + if (slot.n_input_tokens() >= slot.n_ctx) { send_error(slot, string_format("request (%d tokens) exceeds the available context size (%d " "tokens), try increasing it", - slot.task->n_tokens(), slot.n_ctx), + slot.n_input_tokens(), slot.n_ctx), ERROR_TYPE_EXCEED_CONTEXT_SIZE); slot.release(); return; @@ -3175,7 +3592,7 @@ struct server_context_impl { llama_pos pos_next = slot.prompt.tokens.pos_next(n_past); // ref: https://github.com/ggml-org/llama.cpp/pull/24110 - const bool has_new_tokens = (n_past < slot.task->n_tokens()); + const bool has_new_tokens = (n_past < slot.n_input_tokens()); // the largest pos_min required for a checkpoint to be useful const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); @@ -3191,7 +3608,7 @@ struct server_context_impl { // this is useful for debugging prompt caching if (slots_debug) { const int np0 = std::max(n_past - slots_n_diff, 0); - const int np1 = std::min(n_past + slots_n_diff + 2, std::min(slot.prompt.tokens.size(), slot.task->tokens.size())); + const int np1 = std::min(n_past + slots_n_diff + 2, std::min(slot.prompt.tokens.size(), slot.input_tokens().size())); std::stringstream ss0; std::stringstream ss1; @@ -3216,7 +3633,7 @@ struct server_context_impl { } { - const auto token = slot.task->tokens[i]; + const auto token = slot.input_tokens()[i]; const auto piece = token != LLAMA_TOKEN_NULL ? common_token_to_piece(ctx_tgt, token) : "[mtmd]"; ss1 << piece; st1 << std::setw(8) << token; @@ -3284,21 +3701,39 @@ struct server_context_impl { } // [TAG_PROMPT_LOGITS] - if (n_past == slot.task->n_tokens() && n_past > 0) { - SLT_WRN(slot, "need to evaluate at least 1 token for each active slot (n_past = %d, task.n_tokens() = %d)\n", n_past, slot.task->n_tokens()); + if (n_past == slot.n_input_tokens() && n_past > 0) { + SLT_WRN(slot, "need to evaluate at least 1 token for each active slot (n_past = %d, task.n_tokens() = %d)\n", n_past, slot.n_input_tokens()); n_past--; SLT_WRN(slot, "n_past was set to %d\n", n_past); } - slot.stats.n_prompt_cached = n_past; - slot.stats.n_prompt_processed = 0; + // [TAG_PREEMPT_RC] On a resume the cells are gone, so n_past is 0 and + // there is nothing cached to report. Keep the figure the first prefill + // recorded, so usage.cached_tokens stays <= usage.prompt_tokens, and let + // n_prompt_processed accumulate: the recompute really was prompt work, and + // timings.prompt_n is a work counter, not a bill. usage.prompt_tokens + // stays task->n_tokens() and never grows. + if (slot.preempt_resuming) { + const int32_t n_redo = slot.n_input_tokens() - n_past; - metrics.add_prompt_cached(n_past); + slot.n_recompute += n_redo; + metrics.n_recompute_tokens += n_redo; + + SLT_WRN(slot, "resumed after %.2f s parked: recomputing %d tokens, %d cached, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_us) / 1e6, n_redo, n_past, slot.n_preempt); + } else { + slot.stats.n_prompt_cached = n_past; + slot.stats.n_prompt_processed = 0; + + metrics.add_prompt_cached(n_past); + } slot.prompt.tokens.keep_first(n_past); // this is to signal the client that the request has started processing - if (slot.task->params.stream) { + // [TAG_PREEMPT_RC] ... which a resumed slot did a long time ago; the + // headers are out and a second progress chunk would land mid-answer + if (slot.task->params.stream && !slot.preempt_resuming) { if (slot.task->params.return_progress) { // send initial 0% progress update if needed send_partial_response(slot, {}, true); @@ -3311,7 +3746,7 @@ struct server_context_impl { if (!slot.can_split()) { // cannot fit the prompt in the current batch - will try next iter - if (batch.size() + slot.task->n_tokens() > n_batch) { + if (batch.size() + slot.n_input_tokens() > n_batch) { return; } } @@ -3362,7 +3797,7 @@ struct server_context_impl { while (true) { auto cur_token_idx = slot.prompt.n_tokens(); if ( - cur_token_idx >= slot.task->n_tokens() || + cur_token_idx >= slot.n_input_tokens() || input_tokens[cur_token_idx] != LLAMA_TOKEN_NULL // encountered a text token ) { break; @@ -3405,7 +3840,7 @@ struct server_context_impl { const auto last_user_pos = spans.last_user_message_pos(); // add prompt tokens for processing in the current batch - while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) { + while (slot.prompt.n_tokens() < slot.n_input_tokens() && batch.size() < n_batch) { // get next token to process llama_token cur_tok = input_tokens[slot.prompt.n_tokens()]; if (cur_tok == LLAMA_TOKEN_NULL) { @@ -3451,7 +3886,7 @@ struct server_context_impl { bool should_break = false; for (int offset : checkpoint_offsets) { const int n_last = std::min(n_batch, offset); - if (slot.task->n_tokens() == slot.prompt.n_tokens() + n_last) { + if (slot.n_input_tokens() == slot.prompt.n_tokens() + n_last) { should_break = true; break; } @@ -3467,13 +3902,13 @@ struct server_context_impl { const auto n_tokens_start = slot.prompt.n_tokens() - n_tokens_cur; - const bool near_prompt_end = slot.task->n_tokens() < slot.prompt.n_tokens() + n_ubatch; + const bool near_prompt_end = slot.n_input_tokens() < slot.prompt.n_tokens() + n_ubatch; const bool is_user_start = spans.is_user_start(n_tokens_start); const bool is_last_user_message = n_tokens_start == last_user_pos; // entire prompt has been processed - if (slot.prompt.n_tokens() == slot.task->n_tokens()) { + if (slot.prompt.n_tokens() == slot.n_input_tokens()) { slot.state = SLOT_STATE_DONE_PROMPT; GGML_ASSERT(batch.size() > 0); @@ -3481,10 +3916,19 @@ struct server_context_impl { // extract the logits only for the last token batch.set_output(batch.size() - 1, true); - slot.stats.n_gen = 0; - slot.i_batch = batch.size() - 1; + slot.i_batch = batch.size() - 1; + + // [TAG_PREEMPT_RC] A resume must not restart the count or the sampler. The + // sampler object was never touched by the park, so its penalties, its + // grammar and its RNG are exactly where the interrupted step left them. + // init_sampler() would reset it and replay the tokens with + // accept_grammar = false, which is right for a prompt and wrong for the + // tokens this request generated. + if (!slot.preempt_resuming) { + slot.stats.n_gen = 0; - slot.init_sampler(); + slot.init_sampler(); + } } else { // skip ordinary mid-prompt checkpoints, unless the batch starts a user // message or we are near the end of the prompt @@ -3689,7 +4133,8 @@ struct server_context_impl { iterate(slots, [&](server_slot & slot) { // optionally send prompt processing progress if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_DONE_PROMPT) { - if (slot.task->params.stream && slot.task->params.return_progress) { + // [TAG_PREEMPT_RC] a recompute is not a new prompt; the client is mid-answer + if (slot.task->params.stream && slot.task->params.return_progress && !slot.preempt_resuming) { send_partial_response(slot, {}, true); } } @@ -3717,9 +4162,22 @@ struct server_context_impl { GGML_ASSERT(slot.task->need_sampling()); + // [TAG_PREEMPT_RC] the recompute is over: drop the replay list so the slot reads + // its own task again, and stop suppressing the prompt-phase client updates + if (slot.preempt_resuming) { + SLT_WRN(slot, "recompute done, generating again: %d tokens back in the cache, %d recomputed in total\n", + slot.prompt.n_tokens(), slot.n_recompute); + + slot.preempt_resuming = false; + slot.preempt_replay.clear(); + } + // prompt evaluated for next-token prediction slot.state = SLOT_STATE_GENERATING; + // for MTP this only checks that the draft context is where it should be, which the + // re-prefill guarantees: the draft context is a function of the target's prefill, + // and the whole prefill was just run again if (slot.can_speculate()) { common_speculative_begin(spec.get(), slot.id, slot.prompt.tokens.get_text_tokens()); } @@ -3943,7 +4401,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_RC] 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..8ed7f653cd9b 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1562,6 +1562,18 @@ 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 + }, { + "n_recompute_tokens_total", + "Preemption: Total prompt tokens re-processed to put parked slots back", + (double) metrics.n_recompute_tokens }, }; @@ -1586,6 +1598,10 @@ 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 }, }; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..617596e7ee64 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -494,6 +494,7 @@ 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_RC] processing slots currently parked server_metrics metrics; diff --git a/tools/server/tests/unit/test_preempt_recompute.py b/tools/server/tests/unit/test_preempt_recompute.py new file mode 100644 index 000000000000..431301af4667 --- /dev/null +++ b/tools/server/tests/unit/test_preempt_recompute.py @@ -0,0 +1,293 @@ +import os +import time +import tempfile +import pytest +from utils import * + +# Preemption by recompute on a unified KV pool: when the next decode does not fit, one slot's +# cells are dropped and only its token list and its sampler are kept, instead of every slot +# being terminated. When the pool has room the slot goes back through prompt processing, gets +# its whole sequence back into the cache, and carries on from the same token. Every test here +# needs more than one slot and --kv-unified, which is the only configuration in which 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) + os.environ.pop("LLAMA_ARG_PREEMPT", 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(): + # Drop and recompute 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. + # + # This is the assertion RECOMPUTE has to earn and a host-RAM save gets for free: the + # recomputed sequence is prefilled rather than decoded, and the two take different kernels. + 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("recompute done") >= 6 + + assert preempted.body["tokens"] == reference.body["tokens"] + assert preempted.body["content"] == reference.body["content"] + + +def test_usage_does_not_double_count_the_recompute(): + # The recompute is the server's problem, not the client's: prompt_tokens is the prompt the + # client sent, however many times the server had to put it through the model. + global server + server.n_ctx = 512 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + + res = _complete(64) + assert res.status_code == 200 + + n_prompt = res.body["tokens_evaluated"] + plain = server.make_request("POST", "/tokenize", data={"content": "Hi how are you"}) + assert plain.status_code == 200 + # /tokenize does not add BOS, the prompt path does, so allow the one extra token + assert n_prompt - len(plain.body["tokens"]) in (0, 1) + # ... while the work counter does show the recomputes + assert res.body["timings"]["prompt_n"] > n_prompt + + +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's + # cells are dropped until the leader finishes and its cells are purged, and then it is + # recomputed and carries on 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 "recompute done" 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 + + +_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 -- and it + # is the cheap case for recompute, because a slot that has not generated anything simply + # starts its own prompt over and no replay list is needed at all. + 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 + + 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") + 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_no_preempt_disables_preemption(): + # --no-preempt 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"] = "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) + + +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, the recomputed tokens must be + # counted, 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 + assert slot["n_recompute"] == 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["n_recompute_tokens_total"] >= 1 + assert metrics["requests_preempted"] == 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 b2ba5ca46971e621b9aa48e21e274a28f2b4f5ea Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 04:28:03 +0000 Subject: [PATCH 2/2] server: end a parked request that no longer fits the pool instead of leaving it waiting A slot is parked with only its token list kept, so the cells it needs back are the whole sequence. If the pool shrank below that, or the sequence grew into the pool while it was running, no amount of waiting brings it back and the client would hang forever. Give up on it with the same context-overflow error the ordinary path raises. Nothing else in the resume loop can end a parked slot, so this is the only way one leaves without resuming. --- tools/server/server-context.cpp | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 0384261529f8..846656b31df6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2988,6 +2988,33 @@ struct server_context_impl { break; } + // A parked sequence larger than the whole pool can never come back, however long it + // waits: that is a real context overflow and not pressure, and it must be told so + // rather than left hanging. Nothing else in this loop can end a parked slot, so this + // is the only way one leaves without being resumed. + { + bool gave_up = false; + + for (auto * slot : parked) { + if (preempt_n_need(*slot) + PREEMPT_N_MARGIN > n_cells) { + SLT_WRN(*slot, "parked sequence of %d tokens no longer fits the %d cell pool, giving up\n", + slot->n_input_tokens(), n_cells); + + send_error(*slot, + string_format("request (%d tokens) exceeds the available context size (%d tokens), try increasing it", + slot->n_input_tokens(), n_cells), + ERROR_TYPE_EXCEED_CONTEXT_SIZE); + slot->release(); + + gave_up = true; + } + } + + if (gave_up) { + continue; + } + } + 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;