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..846656b31df6 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,313 @@ 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; + } + + // 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; + } + + 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 +3150,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 +3438,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 +3457,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 +3510,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 +3619,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 +3635,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 +3660,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 +3728,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; + + 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); + 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 +3773,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 +3824,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 +3867,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 +3913,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 +3929,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 +3943,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 +4160,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 +4189,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 +4428,8 @@ struct server_context_impl { void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { metrics.n_decode++; for (const auto & slot : slots) { - if (slot.is_processing()) { + // [TAG_PREEMPT_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"