diff --git a/common/arg.cpp b/common/arg.cpp index 5bfa4adcdf0d..25783957dfd4 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1709,10 +1709,20 @@ 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 put it back 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( {"--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), + string_format("maximum host RAM in MiB for parked (preempted) sequences; a slot that fits is parked by copying " + "its sequence to host RAM, one that does not is parked by dropping its cells and recomputing them later " + "(default: %d, -1 - no limit, 0 - never copy; use --no-preempt to turn preemption off)", params.preempt_ram_mib), [](common_params & params, int value) { params.preempt_ram_mib = value; } diff --git a/common/common.h b/common/common.h index c99269f9a967..d2d9a1682c6a 100644 --- a/common/common.h +++ b/common/common.h @@ -614,7 +614,8 @@ 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 + int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = always recompute + bool preempt = true; // with a unified KV cache, park a slot 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 7b4a0330340f..8530e46a8590 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -164,7 +164,8 @@ 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) | +| `--preempt, --no-preempt` | with a unified KV cache and more than one slot, park a running request and put it back later instead of failing every request when the cache fills (default: enabled)
(env: LLAMA_ARG_PREEMPT) | +| `--preempt-ram N` | maximum host RAM in MiB for parked (preempted) sequences; a slot that fits is parked by copying its sequence to host RAM, one that does not is parked by dropping its cells and recomputing them later (default: 8192, -1 - no limit, 0 - never copy; use `--no-preempt` to turn preemption off). Copying back is exact; recomputing costs no host RAM but can change the continuation of a paused answer. Unlike `--cache-ram`, `0` does not disable the feature
(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) | @@ -1139,8 +1140,11 @@ 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_preempt_total` | Counter | Slots parked to make room in the unified KV cache (0 unless `--kv-unified` with more than one slot, and 0 with `--no-preempt`). | +| `llamacpp:n_preempt_swap_total` | Counter | Slots parked by copying the sequence to host RAM. | +| `llamacpp:n_preempt_recompute_total` | Counter | Slots parked by dropping the cells, to be recomputed later. | | `llamacpp:n_resume_total` | Counter | Parked slots put back. | +| `llamacpp:n_recompute_tokens_total` | Counter | Tokens re-processed through the model to put parked slots back (prompt and generated). | | `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. | diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f0cf76b8c501..111b6ae3e637 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,9 +467,14 @@ 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; + // [TAG_PREEMPT] slots parked to make room in the unified KV pool, split by how they were + // parked, put back, and the tokens re-run through the model to put a recomputed one back + // (n_preempt == n_preempt_swap + n_preempt_recompute) + uint64_t n_preempt = 0; + uint64_t n_preempt_swap = 0; + uint64_t n_preempt_recompute = 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 diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8412d8c9ff33..e3f6bdfc70c8 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include #include @@ -60,7 +61,7 @@ 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 + SLOT_STATE_PREEMPTED, // [TAG_PREEMPT] cells released, everything needed to resume is kept }; // [TAG_PREEMPT] server-side request preemption @@ -73,9 +74,36 @@ enum slot_state { // 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. +// later, 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. +// +// There are two ways to give the cells back, and which one a park uses is decided per victim +// by the host RAM budget --preempt-ram: +// +// SWAP under the budget: the sequence is copied out with llama_state_seq_get_data_ext +// and copied back with llama_state_seq_set_data_ext. Exact to the byte, at about +// 36 KiB per token of host RAM on a 4B model. +// RECOMPUTE over it: the cells are simply dropped and the token list (which is in RAM +// anyway) is re-prefilled when room returns. Zero host RAM, and the price is that +// prefill and decode are different kernels, so a recomputed logit is not +// bit-identical and a greedy argmax can flip on a near tie. +// +// This is #184 (SWAP) and #185 (RECOMPUTE) as one mechanism: the budget decides, and the slot +// remembers which way it was parked so the restore takes the matching path. +enum slot_preempt_mode { + PREEMPT_MODE_NONE = 0, + PREEMPT_MODE_SWAP, + PREEMPT_MODE_RECOMPUTE, +}; + +static const char * preempt_mode_name(slot_preempt_mode mode) { + switch (mode) { + case PREEMPT_MODE_SWAP: return "swap"; + case PREEMPT_MODE_RECOMPUTE: return "recompute"; + default: return "none"; + } +} + 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 @@ -331,16 +359,43 @@ struct server_slot { // [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. + // the stream has reached stay on the slot, so a resume is never a new request: no + // retokenisation, no replayed prompt to the client, no seam in the output. + // + // Exactly one of the two representations below is live at a time, chosen by the mode: + // preempt_state_tgt/dft hold the copied sequence (SWAP), preempt_replay holds the token + // list to run through the model again (RECOMPUTE). + slot_preempt_mode preempt_mode = PREEMPT_MODE_NONE; // how the LAST park was done slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; + server_tokens preempt_replay; + bool preempt_resuming = false; // re-prefilling right now (RECOMPUTE only) int32_t n_preempt = 0; // times the CURRENT task has been preempted + int32_t n_recompute = 0; // tokens the CURRENT task has had re-processed 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 + // 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(); + } + + // Re-prefilling tokens this request has ALREADY generated, which is a resume and must not + // look like a new prompt anywhere: not to the sampler, not to the client, not to the + // timings. A slot parked before its prompt was finished has no replay list; it simply + // starts that prompt again and is an ordinary prefill in every respect, so it is not this. + bool preempt_replaying() const { + return preempt_resuming && !preempt_replay.empty(); + } + size_t preempt_state_size() const { return preempt_state_tgt.size() + preempt_state_dft.size(); } @@ -401,6 +456,7 @@ struct server_slot { // copied out, and the resume needs it to know how many cells to ask for. mem.seq_rm(id, -1, -1); + preempt_mode = PREEMPT_MODE_SWAP; state_before_preempt = state; state = SLOT_STATE_PREEMPTED; t_preempt_us = ggml_time_us(); @@ -410,6 +466,53 @@ struct server_slot { return true; } + // The other way to give the cells back: drop them and keep the token list, which is in RAM + // anyway. Used when the sequence does not fit under --preempt-ram. Cannot fail. + 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. + 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(); + + preempt_mode = PREEMPT_MODE_RECOMPUTE; + state_before_preempt = state; + state = SLOT_STATE_PREEMPTED; + t_preempt_us = ggml_time_us(); + + n_preempt++; + } + + // send a RECOMPUTE-parked slot back through prompt processing; nothing to put back, so it + // cannot fail the way preempt_restore() can + void preempt_resume() { + preempt_resuming = true; + state = SLOT_STATE_STARTED; + } + // 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(); @@ -481,6 +584,10 @@ struct server_slot { common_sampler_ptr smpl; + // the "processing has started" signal goes out once per task, however many times the slot + // is parked and sent back through prompt processing + bool sent_begin = false; + llama_token sampled; // in speculative mode, this is the last accepted token // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state @@ -533,8 +640,13 @@ struct server_slot { // [TAG_PREEMPT] preempt_state_free(); + preempt_replay.clear(); + preempt_mode = PREEMPT_MODE_NONE; state_before_preempt = SLOT_STATE_IDLE; + preempt_resuming = false; + sent_begin = false; n_preempt = 0; + n_recompute = 0; n_preempt_fail = 0; n_ctx_shift = 0; t_preempt_us = 0; @@ -696,6 +808,8 @@ struct server_slot { // task on this slot would otherwise take a prefix match against an empty cache if (state == SLOT_STATE_PREEMPTED) { preempt_state_free(); + preempt_replay.clear(); + preempt_resuming = false; prompt_clear(); } @@ -773,7 +887,8 @@ 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; + // [TAG_PREEMPT] a recomputing slot is measured against its replay list, not its prompt + 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); @@ -843,6 +958,11 @@ struct server_slot { {"is_processing", is_processing()}, {"is_preempted", state == SLOT_STATE_PREEMPTED}, {"n_preempt", n_preempt}, + {"n_recompute", n_recompute}, + // how the last park of THIS task was done; null until it has been parked once + {"preempt_mode", preempt_mode == PREEMPT_MODE_NONE + ? json() + : json(preempt_mode_name(preempt_mode))}, }; const auto & ptask = task ? task : task_prev; @@ -2974,7 +3094,15 @@ struct server_context_impl { return res; } - // whether parking this slot stays under --preempt-ram + 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; + } + + // whether parking this slot by SWAP stays under --preempt-ram. -1 is no limit, so every + // park is a swap and this is exactly #184; 0 leaves no room for any copy, so every park + // falls through to recompute. bool preempt_fits_budget(const server_slot & slot) const { if (params_base.preempt_ram_mib < 0) { return true; @@ -2994,35 +3122,82 @@ struct server_context_impl { 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 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 + // a slot whose sequence cannot be rebuilt from its token list alone + bool preempt_can_recompute(const server_slot & slot) const { + // the token list holds placeholders for the media chunks, not the chunks; a replay + // built from it would prefill the wrong thing + return !slot.task->tokens.has_mtmd && !slot.prompt.tokens.has_mtmd; + } + + // How this slot would be parked right now, and PREEMPT_MODE_NONE if it cannot be parked + // at all. This is the whole of the hybrid decision: the copy is preferred because it is + // exact, and the recompute catches everything the budget will not hold. + slot_preempt_mode preempt_mode_for(const server_slot & slot) const { + if (!slot.task) { + return PREEMPT_MODE_NONE; + } + + if (slot.task->is_parent() || slot.task->is_child()) { + return PREEMPT_MODE_NONE; // n_cmpl > 1 slots share one sequence, out of scope here + } + + if (preempt_fits_budget(slot)) { + return PREEMPT_MODE_SWAP; + } + + // a multimodal slot can only be swapped, so over the budget it is not a candidate + return preempt_can_recompute(slot) ? PREEMPT_MODE_RECOMPUTE : PREEMPT_MODE_NONE; + } + + // Cells a parked slot will ask for before it can take its next step. A slot replaying + // tokens it has already generated needs its whole sequence back first; one that was parked + // before it finished its prompt starts that prompt again a batch at a time, exactly like a + // fresh request, and charging it the whole prompt would make a large prompt unwakeable. int32_t preempt_n_need(const server_slot & slot) const { + if (!slot.task) { + return 0; + } + + const int32_t n_spec = preempt_n_spec(slot); + + if (slot.preempt_mode == PREEMPT_MODE_RECOMPUTE && slot.preempt_replaying()) { + return slot.n_input_tokens() + 1 + n_spec; + } + int32_t res = slot.prompt.n_tokens(); if (slot.state_before_preempt == SLOT_STATE_GENERATING) { - res += 1 + preempt_n_spec(slot); + res += 1 + n_spec; + } else if (slot.preempt_replaying()) { + // swapped part-way through a replay: it still owes the rest of that replay before + // it takes a step, which is the charge preempt_kv_reserve() makes for it too + res += std::max(1, slot.n_input_tokens() - res) + 1 + n_spec; } 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) { + if (slot.state == SLOT_STATE_STARTED && !slot.preempt_replaying()) { 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)); + res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), slot.n_input_tokens() - res)); } return res; } + // tokens the pool would hold for this slot once it is fully back, for reporting + int32_t preempt_n_seq(const server_slot & slot) const { + return slot.preempt_mode == PREEMPT_MODE_RECOMPUTE ? slot.n_input_tokens() : slot.prompt.n_tokens(); + } + // 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 @@ -3037,7 +3212,9 @@ struct server_context_impl { for (const auto & slot : slots) { if (slot.state == SLOT_STATE_PREEMPTED) { - continue; // parked: its cells are in host RAM, not in the pool + // parked: a swapped sequence is in host RAM and a recomputed one is nowhere, + // and in both cases the cells are back in the pool + continue; } // a child waiting for its parent's prompt does not share anything yet: until @@ -3090,22 +3267,40 @@ 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; + int32_t n_held = slot.prompt.n_tokens(); + + // a slot just given a task still mirrors the previous request's prompt; + // the batch builder keeps only the prefix the two share, so what it is + // about to ask for counts from that prefix, exactly as preempt_kv_used() + // charges it. Counting from the old prompt instead makes a shorter new + // one look like a single cell and under-reserves the whole pool. + if (slot.state == SLOT_STATE_STARTED && slot.task && !slot.preempt_replaying()) { + n_held = (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); + } + + const int32_t n_left = slot.task ? slot.n_input_tokens() - n_held : 0; - res_pmt += std::max(1, std::min(n_batch, n_left)); + if (slot.preempt_replaying()) { + // 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 + preempt_n_spec(slot); + } else { + 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 + // one batch is all the fresh 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. // [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, @@ -3141,6 +3336,10 @@ struct server_context_impl { slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); } + // 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 preempted PREEMPT_N_STARVED times, and then + // the one preempt_better_victim() ranks first, which by default is the smallest. server_slot * preempt_pick_victim() { server_slot * leader = nullptr; int32_t n_running = 0; @@ -3179,22 +3378,21 @@ struct server_context_impl { continue; } - if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + if (!slot.task || slot.task->is_parent() || slot.task->is_child()) { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + // trim a reused slot to the prefix it keeps before anything sizes it: the mode, + // the budget and the copy must all be decided on the cells it will really hold preempt_normalize_started(slot); - if (!preempt_fits_budget(slot)) { + // the budget no longer rules a slot out, it only decides how it is parked; the + // only slots left out are the ones neither mechanism can take + if (preempt_mode_for(slot) == PREEMPT_MODE_NONE) { continue; } - 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 && preempt_better_victim(slot, *victim))) { + if (!victim || preempt_prefer_victim(slot, *victim)) { victim = &slot; } } @@ -3202,6 +3400,40 @@ struct server_context_impl { return victim; } + // park one slot, whichever way the budget allows; returns the mode used + slot_preempt_mode preempt_park(server_slot & slot) { + slot_preempt_mode mode = preempt_mode_for(slot); + + if (mode == PREEMPT_MODE_SWAP && !slot.preempt_save()) { + // The copy is the only part of a park that can fail, and it leaves the slot + // untouched when it does. Host RAM refusing the sequence is exactly the case the + // recompute exists for, so fall through to it rather than back to the KV-full + // ladder. + SLT_WRN(slot, "%s", "could not copy the sequence to host RAM, dropping its cells instead\n"); + + mode = preempt_can_recompute(slot) ? PREEMPT_MODE_RECOMPUTE : PREEMPT_MODE_NONE; + } + + switch (mode) { + case PREEMPT_MODE_SWAP: + { + metrics.n_preempt++; + metrics.n_preempt_swap++; + } break; + case PREEMPT_MODE_RECOMPUTE: + { + slot.preempt_drop(); + + metrics.n_preempt++; + metrics.n_preempt_recompute++; + } break; + default: + return PREEMPT_MODE_NONE; + } + + return mode; + } + // 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 @@ -3221,23 +3453,38 @@ struct server_context_impl { return a.prompt.n_tokens() < b.prompt.n_tokens(); } + // Is a the slot to take the cells from, ahead of b? A slot already part-way through a + // recompute goes last, since parking it again throws away work it has just redone; then + // one that has not been preempted PREEMPT_N_STARVED times; and only then the size order + // preempt_better_victim() decides. + bool preempt_prefer_victim(const server_slot & a, const server_slot & b) const { + if (a.preempt_resuming != b.preempt_resuming) { + return b.preempt_resuming; + } + + const bool starved_a = a.n_preempt >= PREEMPT_N_STARVED; + const bool starved_b = b.n_preempt >= PREEMPT_N_STARVED; + + if (starved_a != starved_b) { + return starved_b; + } + + return preempt_better_victim(a, b); + } + // 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 + if (!preempt_enabled()) { + return; // --no-preempt, a cache per slot, or a single slot } 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 - } - const int32_t n_cells = n_ctx; // Put back what fits, in the order preempt_resume_head_of_line() describes: by default @@ -3259,6 +3506,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. It would otherwise sit at the head of the line for + // ever without a restore ever being attempted. + { + 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", + preempt_n_seq(*slot), n_cells); + + send_error(*slot, + string_format("request (%d tokens) exceeds the available context size (%d tokens), try increasing it", + preempt_n_seq(*slot), n_cells), + ERROR_TYPE_EXCEED_CONTEXT_SIZE); + slot->release(); + + gave_up = true; + } + } + + if (gave_up) { + continue; + } + } + 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; @@ -3273,31 +3547,6 @@ 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 @@ -3389,6 +3638,22 @@ struct server_context_impl { const int64_t t_start = ggml_time_us(); + if (best->preempt_mode == PREEMPT_MODE_RECOMPUTE) { + // nothing to put back: the slot re-enters prompt processing and re-prefills + // its own token list, one batch per iteration from here on + best->preempt_resume(); + + metrics.n_resume++; + + SLT_WRN(*best, "resuming after %.2f s: %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); + + continue; + } + 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 @@ -3421,13 +3686,17 @@ struct server_context_impl { 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_fits_budget(slot) && - slot.preempt_save()) { - metrics.n_preempt++; + (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every) { + const int32_t n_tokens = slot.prompt.n_tokens(); + const int32_t n_gen = (int32_t) slot.stats.n_gen; + + const slot_preempt_mode mode = preempt_park(slot); - 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)); + if (mode != PREEMPT_MODE_NONE) { + SLT_WRN(slot, "preempted on request after %d generated tokens: parked by %s, %d cells released, %.1f MiB in host RAM\n", + n_gen, preempt_mode_name(mode), n_tokens, + slot.preempt_state_size() / (1024.0 * 1024.0)); + } } } } @@ -3460,16 +3729,18 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); - if (!victim->preempt_save()) { + const slot_preempt_mode mode = preempt_park(*victim); + + if (mode == PREEMPT_MODE_NONE) { break; // could not park it; the existing retry ladder is still behind us } - metrics.n_preempt++; - - SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + SLT_WRN(*victim, "preempted: parked by %s, %d cells released in %.2f ms, %.1f MiB in host RAM, %d tokens to recompute later, kv %d/%d (wanted %d), preemptions %d\n", + preempt_mode_name(mode), n_tokens, (ggml_time_us() - t_start) / 1e3, victim->preempt_state_size() / (1024.0 * 1024.0), + mode == PREEMPT_MODE_RECOMPUTE ? victim->n_input_tokens() : 0, preempt_kv_used(), n_cells, n_used, victim->n_preempt); } @@ -3838,19 +4109,24 @@ 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; + // [TAG_PREEMPT] a recomputing slot reads its replay list, not its task + 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] 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) { @@ -3887,33 +4163,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; @@ -3996,7 +4272,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)); @@ -4012,7 +4288,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; @@ -4037,7 +4313,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; @@ -4105,21 +4381,40 @@ 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] On a recompute the cells are gone, so n_past is 0 and + // there is nothing cached to report. Leave every figure the first prefill + // recorded exactly as it was: the timings a client reads describe the + // request it made, and the recomputed work is reported separately, as + // n_recompute on /slots and n_recompute_tokens_total on /metrics. + if (slot.preempt_replaying()) { + // the tokens are counted as they are actually decoded, in + // metrics_post_decode(): a replay that is interrupted and started + // over must not be charged twice + SLT_WRN(slot, "resumed after %.2f s parked: recomputing %d tokens, %d cached, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_us) / 1e6, + slot.n_input_tokens() - n_past, 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); + 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] ... which it does once per task. A slot that is here for + // the second time was parked and sent back, and its headers are already + // out; a second one of these would land in the middle of the stream. + if (slot.task->params.stream && !slot.sent_begin) { + slot.sent_begin = true; + if (slot.task->params.return_progress) { // send initial 0% progress update if needed send_partial_response(slot, {}, true); @@ -4132,7 +4427,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; } } @@ -4183,7 +4478,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; @@ -4226,7 +4521,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) { @@ -4272,7 +4567,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; } @@ -4288,13 +4583,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); @@ -4302,10 +4597,21 @@ 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] A replay 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. A slot parked before it finished its + // prompt has generated nothing yet and must still be initialised here, + // which is why this asks about the replay and not about the resume. + if (!slot.preempt_replaying()) { + 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 @@ -4356,7 +4662,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 preempt_enabled() && llama_get_memory(ctx_tgt); } bool preempt_last_resort(int32_t off) { @@ -4407,14 +4713,16 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); - if (!victim->preempt_save()) { + const slot_preempt_mode mode = preempt_park(*victim); + + if (mode == PREEMPT_MODE_NONE) { 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", + SLT_WRN(*victim, "preempted as a last resort: parked by %s, %d cells released in %.2f ms, %.1f MiB in host RAM, kv %d/%d (wanted %d), preemptions %d\n", + preempt_mode_name(mode), n_tokens, (ggml_time_us() - t_start) / 1e3, victim->preempt_state_size() / (1024.0 * 1024.0), @@ -4628,7 +4936,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] a replay is not a new prompt; the client is mid-answer + if (slot.task->params.stream && slot.task->params.return_progress && !slot.preempt_replaying()) { send_partial_response(slot, {}, true); } } @@ -4656,6 +4965,17 @@ struct server_context_impl { GGML_ASSERT(slot.task->need_sampling()); + // [TAG_PREEMPT] 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; @@ -4904,7 +5224,12 @@ struct server_context_impl { n_prompt_tokens++; auto & slot = slots[t.id_slot]; - if (slot.stats.is_set()) { + // [TAG_PREEMPT] a replayed token is not one of the request's prompt tokens; it is + // work the server owes itself, and it is counted as such + if (slot.preempt_replaying()) { + slot.n_recompute++; + metrics.n_recompute_tokens++; + } else if (slot.stats.is_set()) { slot.stats.n_prompt_processed++; } } @@ -4922,7 +5247,10 @@ struct server_context_impl { for (int i = off; i < off + n_tokens; ++i) { const auto & t = batch.tokens[i]; auto & slot = slots[t.id_slot]; - if (t.is_prompt && slot.stats.is_set()) { + // [TAG_PREEMPT] and it must not drag the prompt timestamp forward either: that + // would report the whole parked wall clock as prompt time and inflate the + // generation rate by the same amount + if (t.is_prompt && slot.stats.is_set() && !slot.preempt_replaying()) { slot.stats.set_prompt_last(t_now); } } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 9afe3c7f06a8..30f6c70be58b 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1566,10 +1566,22 @@ std::string server_task_result_metrics::to_metrics() { "n_preempt_total", "Preemption: Total slots parked to make room in the unified KV cache", (double) metrics.n_preempt + }, { + "n_preempt_swap_total", + "Preemption: Total slots parked by copying the sequence to host RAM", + (double) metrics.n_preempt_swap + }, { + "n_preempt_recompute_total", + "Preemption: Total slots parked by dropping the cells, to be recomputed later", + (double) metrics.n_preempt_recompute }, { "n_resume_total", "Preemption: Total parked slots put back", (double) metrics.n_resume + }, { + "n_recompute_tokens_total", + "Preemption: Total tokens re-processed through the model to put parked slots back", + (double) metrics.n_recompute_tokens }, }; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index a2fcd750f18c..96170662f854 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -40,6 +40,7 @@ def create_server(): os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) os.environ.pop("LLAMA_SERVER_PREEMPT_PLANNER", None) os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) + os.environ.pop("LLAMA_ARG_PREEMPT", None) def _complete(n_predict: int, prompt: str = "Hi how are you"): @@ -209,12 +210,16 @@ def _late(n_predict, prompt): 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 +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. + # + # This used to be --preempt-ram 0. That value now means "never copy to host RAM, always + # recompute", so it parks rather than disables, and the switch had to move to its own + # flag. test_preempt_hybrid.py covers what --preempt-ram 0 does instead. global server server.n_ctx = 256 - os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + os.environ["LLAMA_ARG_PREEMPT"] = "0" server.start() log = LogReader(server.log_path) diff --git a/tools/server/tests/unit/test_preempt_hybrid.py b/tools/server/tests/unit/test_preempt_hybrid.py new file mode 100644 index 000000000000..a98a84b942ba --- /dev/null +++ b/tools/server/tests/unit/test_preempt_hybrid.py @@ -0,0 +1,296 @@ +import os +import tempfile +import threading +import pytest +from utils import * + +# The hybrid half of preemption: WHICH mechanism parks a slot, and that the choice is made by +# the --preempt-ram budget. test_preempt.py covers the policy itself (who is parked, in what +# order they come back, that nobody is terminated); everything here is about the two ways of +# parking and the counters that tell them apart. +# +# default budget every park fits, so every park is a swap, and the output is exact +# --preempt-ram 0 nothing fits, so every park is a recompute, and the run still finishes +# (on #184 this configuration disabled preemption and both requests died) +# --preempt-ram 1 a budget small enough to hold some parks and not others: both + +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.server_metrics = 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_RAM", None) + os.environ.pop("LLAMA_ARG_PREEMPT", None) + + +def _complete(n_predict: int, prompt: str = "Hi how are you"): + return server.make_request("POST", "/completion", data={ + "n_predict": n_predict, + "prompt": prompt, + "ignore_eos": True, + "return_tokens": True, + "temperature": 0.0, + "seed": 42, + }) + + +class SlotWatcher: + """Poll /slots while requests are in flight and keep every parked state it saw. + + The counters on /metrics say a park happened; only this says what a client polling /slots + would have been told while it was happening.""" + + def __init__(self): + self.seen = [] + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _run(self): + while not self._stop.is_set(): + try: + res = server.make_request("GET", "/slots", timeout=5) + if res.status_code == 200: + for slot in res.body: + if slot["is_preempted"]: + self.seen.append((slot["preempt_mode"], slot["n_preempt"])) + except Exception: + pass + self._stop.wait(0.02) + + def __enter__(self): + self._thread.start() + return self + + def __exit__(self, *exc): + self._stop.set() + self._thread.join(timeout=5) + + @property + def modes(self): + return {mode for mode, _ in self.seen} + + +def _metrics(): + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + out = {} + for line in res.body.splitlines(): + if line.startswith("llamacpp:"): + name, value = line.split(" ", 1) + out[name[len("llamacpp:"):]] = float(value) + return out + + +_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. " +) * 24 + + +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 + words = words[: len(words) - max(1, (n - n_tokens) // 8)] + raise AssertionError("empty prompt") + + +def test_default_budget_parks_by_swap_and_the_output_is_identical(): + # Under the default 8192 MiB every park fits in host RAM, so the hybrid must behave + # exactly as #184: a copy out and a copy back, byte-identical output, and no recompute. + 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) + + preempted = _complete(64) + assert preempted.status_code == 200 + assert preempted.body["content"] == reference.body["content"] + assert preempted.body["tokens"] == reference.body["tokens"] + + text = log.drain() + assert "parked by swap" in text + assert "parked by recompute" not in text + + m = _metrics() + assert m["n_preempt_total"] >= 6 + assert m["n_preempt_swap_total"] == m["n_preempt_total"] + assert m["n_preempt_recompute_total"] == 0 + assert m["n_recompute_tokens_total"] == 0 + assert m["n_resume_total"] == m["n_preempt_total"] + assert m["preempt_ram_bytes"] == 0 + + +def test_preempt_ram_zero_parks_by_recompute_and_both_requests_finish(): + # The case #184 could not serve. With --preempt-ram 0 no sequence may be copied to host + # RAM, so on #184 nothing was parked and the KV-full ladder ended both requests with + # "Context size has been exceeded". Here the same budget means "always recompute", and + # both requests finish with nothing held in host RAM at all. + global server + server.n_ctx = 256 + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + with SlotWatcher() as watcher: + 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 "parked by recompute" in text + assert "parked by swap" not in text + + # whatever /slots showed while a slot was parked, it can only have been a recompute; the + # poll can miss a short park, so it is the mode that is asserted, not that it saw one + assert watcher.modes <= {"recompute"}, watcher.seen + + 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 + + m = _metrics() + assert m["n_preempt_total"] >= 1 + assert m["n_preempt_recompute_total"] == m["n_preempt_total"] + assert m["n_preempt_swap_total"] == 0 + assert m["n_recompute_tokens_total"] >= 1 + assert m["n_resume_total"] == m["n_preempt_total"] + assert m["preempt_ram_bytes"] == 0 + assert m["requests_preempted"] == 0 + + +def test_a_small_budget_parks_both_ways(): + # The point of the hybrid: the budget is a boundary, not a switch. With four slots on a + # small pool several sequences are parked at once, the first ones fit under 1 MiB and are + # copied, the ones after them do not and are recomputed, and every request still finishes. + # + # Four slots and not two: with two slots only one can ever be parked at a time (the leader + # is never a victim), one park of this model is well under 1 MiB, and the budget would + # never be reached. This model holds about 740 bytes of KV per token, so 1 MiB runs out + # at roughly 1400 parked tokens, which four 1500-token sequences cross and four + # 400-token ones do not. + global server + server.n_ctx = 4096 + server.n_slots = 4 + # the preset's 32-token batch would take thousands of iterations to prefill four 600-token + # prompts into a 4096-cell pool; nothing here depends on the batch being small + server.n_batch = 512 + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + server.start() + log = LogReader(server.log_path) + + prompts = [_prompt_of_about(600, salt)[0] for salt in ("Alpha", "Bravo", "Charlie", "Delta")] + # 4 x 1500 cells wanted from a pool of 4096: the early, smaller parks are copied and the + # later, larger ones do not fit the budget and are recomputed + n_predict = 900 + results = parallel_function_calls([(_complete, (n_predict, p)) for p in prompts]) + + text = log.drain() + assert "Context size has been exceeded" not 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 + + m = _metrics() + assert m["n_preempt_swap_total"] >= 1, "nothing was parked by swap under a 1 MiB budget" + assert m["n_preempt_recompute_total"] >= 1, "the budget never pushed a park to recompute" + assert m["n_preempt_swap_total"] + m["n_preempt_recompute_total"] == m["n_preempt_total"] + assert m["n_resume_total"] == m["n_preempt_total"] + assert m["preempt_ram_bytes"] == 0 + + +def test_a_recomputed_request_does_not_double_count_its_prompt(): + # A recompute runs the prompt through the model again. That is the server's problem and not + # the client's: usage.prompt_tokens is the prompt the client sent, and the timings describe + # the request it made, however many times the server had to prefill it. The work is + # reported separately, as n_recompute on /slots and n_recompute_tokens_total on /metrics. + global server + server.n_ctx = 512 + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + + prompt = "Once upon a time there was a brave knight who" + plain = server.make_request("POST", "/tokenize", data={"content": prompt}) + assert plain.status_code == 200 + n_prompt = len(plain.body["tokens"]) + + n_predict = 64 + res = _complete(n_predict, prompt) + assert res.status_code == 200 + + m = _metrics() + assert m["n_preempt_recompute_total"] >= 6, "the request was never recomputed" + assert m["n_preempt_swap_total"] == 0 + # ... and it really did re-run more tokens than the request ever had + assert m["n_recompute_tokens_total"] > n_prompt + n_predict + + # /tokenize does not add BOS, the prompt path does, so allow the one extra token + assert res.body["tokens_evaluated"] - n_prompt in (0, 1) + timings = res.body["timings"] + assert timings["prompt_n"] == res.body["tokens_evaluated"], "the recompute was billed as prompt" + assert timings["predicted_n"] == n_predict + # a prompt timestamp dragged forward by the recomputes would report the whole wall clock as + # prompt time and inflate the generation rate by the same amount + assert timings["prompt_ms"] < timings["predicted_ms"] + + oai = server.make_request("POST", "/v1/completions", data={ + "prompt": prompt, + "max_tokens": n_predict, + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + }) + assert oai.status_code == 200 + usage = oai.body["usage"] + assert usage["prompt_tokens"] - n_prompt in (0, 1) + assert usage["completion_tokens"] == n_predict + assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]