diff --git a/common/arg.cpp b/common/arg.cpp
index 86f8610a56d0..5bfa4adcdf0d 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -1709,6 +1709,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
params.cache_ram_mib = value;
}
).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+ add_opt(common_arg(
+ {"--preempt-ram"}, "N",
+ string_format("with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; "
+ "N is the maximum host RAM for parked sequences in MiB (default: %d, -1 - no limit, 0 - disable)", params.preempt_ram_mib),
+ [](common_params & params, int value) {
+ params.preempt_ram_mib = value;
+ }
+ ).set_env("LLAMA_ARG_PREEMPT_RAM").set_examples({LLAMA_EXAMPLE_SERVER}));
add_opt(common_arg(
{"-kvu", "--kv-unified"},
{"-no-kvu", "--no-kv-unified"},
diff --git a/common/common.h b/common/common.h
index de49dac9f63a..c99269f9a967 100644
--- a/common/common.h
+++ b/common/common.h
@@ -614,6 +614,7 @@ struct common_params {
int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot
int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints
int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
+ int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = disable preemption
std::string hostname = "127.0.0.1";
std::string public_path = ""; // NOLINT
diff --git a/tools/server/README.md b/tools/server/README.md
index 93736c3edfa9..7b4a0330340f 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -164,6 +164,7 @@ For the full list of features, please refer to [server's changelog](https://gith
| `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) |
| `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) |
| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) |
+| `--preempt-ram N` | with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; N is the maximum host RAM for parked sequences in MiB (default: 8192, -1 - no limit, 0 - disable)
(env: LLAMA_ARG_PREEMPT_RAM) |
| `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)
(env: LLAMA_ARG_KV_UNIFIED) |
| `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)
(env: LLAMA_ARG_CACHE_IDLE_SLOTS) |
| `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) |
@@ -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 put back. |
+| `llamacpp:requests_preempted` | Gauge | Requests currently parked, waiting for room in the unified KV cache. |
+| `llamacpp:preempt_ram_bytes` | Gauge | Host RAM held by parked sequences. |
### POST `/slots/{id_slot}?action=save`: Save the prompt cache of the specified slot to a file.
diff --git a/tools/server/server-common.h b/tools/server/server-common.h
index f8ea82ef4cf5..f0cf76b8c501 100644
--- a/tools/server/server-common.h
+++ b/tools/server/server-common.h
@@ -467,6 +467,10 @@ struct server_metrics {
uint64_t n_decode = 0;
uint64_t n_busy_slots = 0;
+ // [TAG_PREEMPT] slots parked to make room in the unified KV pool, and put back
+ uint64_t n_preempt = 0;
+ uint64_t n_resume = 0;
+
uint64_t n_draft_tokens = 0; // Total draft tokens generated
uint64_t n_draft_accepted = 0; // Draft tokens actually accepted
uint64_t n_draft_verif_steps = 0; // Total draft token verification steps by the target model
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index a9edbd7be8b4..6723c51397ed 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -59,8 +59,27 @@ enum slot_state {
SLOT_STATE_PROCESSING_PROMPT,
SLOT_STATE_DONE_PROMPT,
SLOT_STATE_GENERATING,
+ SLOT_STATE_PREEMPTED, // [TAG_PREEMPT] cells released, everything needed to resume is in host RAM
};
+// [TAG_PREEMPT] server-side request preemption
+//
+// With --kv-unified the cells are one pool shared by every slot, and each slot believes it
+// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves
+// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size
+// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the
+// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and
+// continue with the rest".
+//
+// Nothing is terminated here. The cells of one slot are taken back and given to it again
+// later: its sequence is copied to host RAM, its cells are released, and when the pool has
+// room the copy goes back and the slot carries on with the same sampler, the same generated
+// text and the same open stream. A streaming client sees a pause, not an error.
+constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation
+constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected
+constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on
+constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked
+
struct server_slot; // forward declaration
struct server_batch {
@@ -293,6 +312,123 @@ struct server_slot {
prompt.clear();
}
+ // [TAG_PREEMPT] state of a slot whose cells were taken back
+ //
+ // Only the KV cells leave. The task, the sampler, the generated text and the position
+ // the stream has reached stay on the slot, so a resume is a memcpy and not a new
+ // request: no retokenisation, no replayed prompt, no seam in the output.
+ slot_state state_before_preempt = SLOT_STATE_IDLE;
+ std::vector preempt_state_tgt;
+ std::vector preempt_state_dft;
+ int32_t n_preempt = 0; // times the CURRENT task has been preempted
+ int32_t n_preempt_fail = 0; // consecutive failed restores
+ int64_t t_preempt_us = 0; // when it was parked
+
+ size_t preempt_state_size() const {
+ return preempt_state_tgt.size() + preempt_state_dft.size();
+ }
+
+ void preempt_state_free() {
+ preempt_state_tgt.clear();
+ preempt_state_tgt.shrink_to_fit();
+ preempt_state_dft.clear();
+ preempt_state_dft.shrink_to_fit();
+ }
+
+ // bytes preempt_save() would need for this slot right now
+ size_t preempt_state_required() const {
+ return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) +
+ (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0);
+ }
+
+ // copy the sequence out of the cache and release its cells
+ bool preempt_save() {
+ const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
+ const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0;
+
+ try {
+ preempt_state_tgt.resize(size_tgt);
+ preempt_state_dft.resize(size_dft);
+ } catch (const std::bad_alloc & e) {
+ SLT_ERR(*this, "failed to allocate %.3f MiB for the preemption state: %s\n",
+ (size_tgt + size_dft) / (1024.0 * 1024.0), e.what());
+ preempt_state_free();
+ return false;
+ }
+
+ if (llama_state_seq_get_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) {
+ SLT_ERR(*this, "%s", "failed to copy the target sequence out of the KV cache\n");
+ preempt_state_free();
+ return false;
+ }
+
+ if (size_dft > 0 &&
+ llama_state_seq_get_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) {
+ SLT_ERR(*this, "%s", "failed to copy the draft sequence out of the KV cache\n");
+ preempt_state_free();
+ return false;
+ }
+
+ // The draft is a prediction, not a result, so it goes with the cells. Preemption
+ // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already
+ // holds exactly the tokens the state above covers -- including the rollback done by
+ // the checkpoint path when a draft was only partially accepted.
+ spec_draft.clear();
+ spec_i_batch.clear();
+ spec_ckpt.clear();
+ spec_is_replay = false;
+
+ i_batch = -1;
+
+ // note: prompt.tokens is deliberately kept. It is the mirror of the state just
+ // copied out, and the resume needs it to know how many cells to ask for.
+ mem.seq_rm(id, -1, -1);
+
+ state_before_preempt = state;
+ state = SLOT_STATE_PREEMPTED;
+ t_preempt_us = ggml_time_us();
+
+ n_preempt++;
+
+ return true;
+ }
+
+ // put the sequence back; the slot then continues from the token it was about to decode
+ bool preempt_restore() {
+ const size_t size_tgt = preempt_state_tgt.size();
+ const size_t size_dft = preempt_state_dft.size();
+
+ if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) {
+ // no room after all: drop the half-written sequence and stay parked
+ mem.seq_rm(id, -1, -1);
+ n_preempt_fail++;
+ return false;
+ }
+
+ if (size_dft > 0 &&
+ llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) {
+ mem.seq_rm(id, -1, -1);
+ n_preempt_fail++;
+ return false;
+ }
+
+ preempt_state_free();
+
+ n_preempt_fail = 0;
+
+ state = state_before_preempt;
+
+ // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks
+ // that the draft context is where it should be, which the restore above ensures.
+ // A slot parked while still processing its prompt makes that transition itself
+ // once the prompt is done.
+ if (state == SLOT_STATE_GENERATING && can_speculate()) {
+ common_speculative_begin(spec, id, prompt.tokens.get_text_tokens());
+ }
+
+ return true;
+ }
+
std::vector lora;
int32_t alora_invocation_start = -1;
@@ -351,6 +487,13 @@ struct server_slot {
n_predict_max = -1;
+ // [TAG_PREEMPT]
+ preempt_state_free();
+ state_before_preempt = SLOT_STATE_IDLE;
+ n_preempt = 0;
+ n_preempt_fail = 0;
+ t_preempt_us = 0;
+
llama_set_sampler(ctx_tgt, id, nullptr);
// clear alora start
@@ -503,6 +646,14 @@ struct server_slot {
t_last_used = ggml_time_us();
+ // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be
+ // released while parked), so the mirror of them must not outlive them: the next
+ // task on this slot would otherwise take a prefix match against an empty cache
+ if (state == SLOT_STATE_PREEMPTED) {
+ preempt_state_free();
+ prompt_clear();
+ }
+
state = SLOT_STATE_IDLE;
// do not keep context of the child slots - the parent's context is enough
@@ -645,6 +796,8 @@ struct server_slot {
{"n_ctx", n_ctx},
{"speculative", can_speculate()},
{"is_processing", is_processing()},
+ {"is_preempted", state == SLOT_STATE_PREEMPTED},
+ {"n_preempt", n_preempt},
};
const auto & ptask = task ? task : task_prev;
@@ -1249,6 +1402,16 @@ struct server_context_impl {
}
}
+ {
+ const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY");
+ preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0;
+
+ if (preempt_test_every > 0) {
+ SRV_WRN("LLAMA_SERVER_PREEMPT_EVERY = %d (test knob: preempting every slot every %d tokens)\n",
+ preempt_test_every, preempt_test_every);
+ }
+ }
+
{
const char * LLAMA_SERVER_SLOTS_N_DIFF = getenv("LLAMA_SERVER_SLOTS_N_DIFF");
slots_n_diff = LLAMA_SERVER_SLOTS_N_DIFF ? atoi(LLAMA_SERVER_SLOTS_N_DIFF) : 0;
@@ -2385,11 +2548,15 @@ struct server_context_impl {
case SERVER_TASK_TYPE_METRICS:
{
int n_processing_slots = 0;
+ int n_preempted_slots = 0;
for (server_slot & slot : slots) {
if (slot.is_processing()) {
n_processing_slots++;
}
+ if (slot.state == SLOT_STATE_PREEMPTED) {
+ n_preempted_slots++;
+ }
}
SRV_DBG("n_processing_slots = %d\n", n_processing_slots);
@@ -2397,6 +2564,8 @@ struct server_context_impl {
res->id = task.id;
res->n_processing_slots = n_processing_slots;
res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size();
+ res->n_preempted_slots = n_preempted_slots;
+ res->preempt_ram_bytes = preempt_ram_used();
res->metrics = metrics;
if (task.metrics_reset_bucket) {
@@ -2674,6 +2843,317 @@ struct server_context_impl {
};
#endif
+ //
+ // [TAG_PREEMPT] server-side request preemption
+ //
+
+
+ // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens,
+ // whether or not the pool is under pressure. It exists to answer the only question that
+ // matters about a resume: with one request on an idle server the batch has the same
+ // shape at every step, so a preempted continuation that is not byte-identical to an
+ // uninterrupted one is the preemption's fault and nothing else's.
+ int32_t preempt_test_every = 0;
+
+ int32_t preempt_n_spec_max() const {
+ return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0;
+ }
+
+ // host RAM the parked sequences hold right now
+ size_t preempt_ram_used() const {
+ size_t res = 0;
+
+ for (const auto & slot : slots) {
+ res += slot.preempt_state_size();
+ }
+
+ return res;
+ }
+
+ // whether parking this slot stays under --preempt-ram
+ bool preempt_fits_budget(const server_slot & slot) const {
+ if (params_base.preempt_ram_mib < 0) {
+ return true;
+ }
+
+ const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024;
+
+ return preempt_ram_used() + slot.preempt_state_required() <= budget;
+ }
+
+ // cells the slot will ask for on its next step once it is back in the pool
+ int32_t preempt_n_need(const server_slot & slot) const {
+ int32_t res = slot.prompt.n_tokens();
+
+ if (slot.state_before_preempt == SLOT_STATE_GENERATING) {
+ res += 1 + preempt_n_spec_max();
+ } else {
+ const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0;
+
+ res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left));
+ }
+
+ return res;
+ }
+
+ // Cells the pool is holding right now. A released slot keeps its prompt in the cache
+ // for the next request to reuse as a prefix, so idle slots count too: the first version
+ // of this counted only the running ones, decided a pool holding 8185 cached cells was
+ // empty, and every resume failed against a cache that was actually full.
+ int32_t preempt_kv_used() const {
+ int32_t res = 0;
+
+ for (const auto & slot : slots) {
+ if (slot.state == SLOT_STATE_PREEMPTED) {
+ continue; // parked: its cells are in host RAM, not in the pool
+ }
+
+ res += slot.prompt.n_tokens();
+ }
+
+ return res;
+ }
+
+ // cells those slots are about to ask for on the next decode
+ int32_t preempt_kv_reserve() const {
+ const int32_t n_spec = preempt_n_spec_max();
+ const int32_t n_batch = llama_n_batch(ctx_tgt);
+
+ int32_t res = 0;
+ int32_t res_pmt = 0;
+
+ for (const auto & slot : slots) {
+ switch (slot.state) {
+ case SLOT_STATE_GENERATING:
+ case SLOT_STATE_DONE_PROMPT:
+ {
+ res += 1 + n_spec;
+ } break;
+ case SLOT_STATE_STARTED:
+ case SLOT_STATE_PROCESSING_PROMPT:
+ {
+ const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0;
+
+ res_pmt += std::max(1, std::min(n_batch, n_left));
+ } break;
+ default:
+ break;
+ }
+ }
+
+ // one batch is all the prompt slots get between them, however many are waiting
+ return res + std::min(res_pmt, n_batch);
+ }
+
+ // Keep the slot that is furthest along -- it is the closest to finishing and to giving
+ // its cells back -- and among the rest prefer one that has not been preempted
+ // PREEMPT_N_STARVED times already, then the smallest.
+ server_slot * preempt_pick_victim() {
+ server_slot * leader = nullptr;
+ int32_t n_running = 0;
+
+ for (auto & slot : slots) {
+ if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) {
+ n_running++;
+
+ if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) {
+ leader = &slot;
+ }
+ }
+ }
+
+ if (n_running < 2) {
+ // a single conversation that does not fit the pool on its own is a real context
+ // overflow and not a scheduling problem - leave it to the existing error path
+ return nullptr;
+ }
+
+ server_slot * victim = nullptr;
+
+ for (auto & slot : slots) {
+ // Before the batch is built every one of these is at a token boundary: a
+ // generating slot between two sampled tokens, a prompt-processing slot between
+ // two chunks of its prompt, a started slot with only a cached prefix (or
+ // nothing) in the pool. A slot holding no cells is still worth parking - it
+ // is about to ask for a whole batch of them.
+ if (slot.state != SLOT_STATE_GENERATING &&
+ slot.state != SLOT_STATE_PROCESSING_PROMPT &&
+ slot.state != SLOT_STATE_STARTED) {
+ continue;
+ }
+
+ if (&slot == leader) {
+ continue;
+ }
+
+ if (slot.task && (slot.task->is_parent() || slot.task->is_child())) {
+ continue; // n_cmpl > 1 slots share one sequence, out of scope here
+ }
+
+ if (!preempt_fits_budget(slot)) {
+ continue;
+ }
+
+ const bool starved = slot.n_preempt >= PREEMPT_N_STARVED;
+ const bool starved_cur = victim && victim->n_preempt >= PREEMPT_N_STARVED;
+
+ if (!victim ||
+ (starved_cur && !starved) ||
+ (starved_cur == starved && slot.prompt.n_tokens() < victim->prompt.n_tokens())) {
+ victim = &slot;
+ }
+ }
+
+ return victim;
+ }
+
+ // called once per update_slots(), before the batch is built: at that point every slot is
+ // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no
+ // draft is in flight, so a slot can be removed from the picture without unpicking a
+ // half-decoded batch
+ void update_preemption() {
+ if (!params_base.kv_unified || slots.size() < 2) {
+ return; // with a cache per slot, no slot can take another one's cells
+ }
+
+ if (params_base.preempt_ram_mib == 0) {
+ return; // --preempt-ram 0: the KV-full retry ladder, as before
+ }
+
+ const int32_t n_cells = n_ctx;
+
+ // Put back what fits: the most-preempted slot first, then the one parked longest.
+ // A slot that does not fit yet must not hold up a smaller one that does: it keeps
+ // its place at the head of the line, and the smaller one is the first to be parked
+ // again if the pool fills, so letting it through costs the head nothing.
+ for (;;) {
+ std::vector parked;
+
+ for (auto & slot : slots) {
+ if (slot.state == SLOT_STATE_PREEMPTED) {
+ parked.push_back(&slot);
+ }
+ }
+
+ if (parked.empty()) {
+ break;
+ }
+
+ std::sort(parked.begin(), parked.end(), [](const server_slot * a, const server_slot * b) {
+ if (a->n_preempt != b->n_preempt) {
+ return a->n_preempt > b->n_preempt;
+ }
+
+ return a->t_preempt_us < b->t_preempt_us;
+ });
+
+ server_slot * best = nullptr;
+
+ // Room for the sequence AND for the next step of everything already running,
+ // so that a resume cannot immediately trigger the preemption of someone else.
+ // A cached prompt on an idle slot is worth less than a conversation waiting to
+ // continue, so give those cells up first - same call the KV-full path makes.
+ 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;
+ }
+
+ const int64_t t_start = ggml_time_us();
+
+ if (!best->preempt_restore()) {
+ // update_slots() runs in a tight loop while tasks are pending, so a counter
+ // alone burns its whole budget in a couple of milliseconds. Give up only on
+ // a slot that has been failing for a while, and keep the log quiet.
+ if (best->n_preempt_fail % 64 == 1) {
+ SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n",
+ best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6);
+ }
+
+ if (best->n_preempt_fail >= PREEMPT_N_FAIL_MAX &&
+ ggml_time_us() - best->t_preempt_us > PREEMPT_FAIL_US) {
+ send_error(*best, "failed to restore the preempted sequence");
+ best->release();
+ }
+
+ break;
+ }
+
+ metrics.n_resume++;
+
+ SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n",
+ (ggml_time_us() - best->t_preempt_us) / 1e6,
+ best->prompt.n_tokens(),
+ (ggml_time_us() - t_start) / 1e3,
+ preempt_kv_used(), n_cells,
+ best->n_preempt);
+ }
+
+ // forced preemption, for the determinism test only
+ if (preempt_test_every > 0) {
+ for (auto & slot : slots) {
+ if (slot.state == SLOT_STATE_GENERATING &&
+ (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every &&
+ preempt_fits_budget(slot) &&
+ slot.preempt_save()) {
+ metrics.n_preempt++;
+
+ SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n",
+ (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0));
+ }
+ }
+ }
+
+ // 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 (parked %.1f MiB of the %d MiB --preempt-ram budget)\n",
+ n_used, n_cells, preempt_ram_used() / (1024.0 * 1024.0), params_base.preempt_ram_mib);
+ break;
+ }
+
+ const int32_t n_tokens = victim->prompt.n_tokens();
+ const int64_t t_start = ggml_time_us();
+
+ if (!victim->preempt_save()) {
+ break; // could not park it; the existing retry ladder is still behind us
+ }
+
+ metrics.n_preempt++;
+
+ SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n",
+ n_tokens,
+ (ggml_time_us() - t_start) / 1e3,
+ victim->preempt_state_size() / (1024.0 * 1024.0),
+ preempt_kv_used(), n_cells, n_used,
+ victim->n_preempt);
+ }
+ }
+
void update_slots() {
#ifdef DEBUG_TIMINGS
static int64_t t_prev = 0;
@@ -2715,6 +3195,9 @@ struct server_context_impl {
}
}
+ // [TAG_PREEMPT] make the pool fit the step that is about to be built
+ update_preemption();
+
try {
scoped_timer t(t_pre_decode, n_pre_decode);
pre_decode();
@@ -3000,7 +3483,9 @@ struct server_context_impl {
return; // batch is full, skip remaining slots
}
- if (!slot.is_processing()) {
+ // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to
+ // batch; it takes no part in this pass until it is restored
+ if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) {
return;
}
@@ -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] a parked slot is processing but took no part in this decode
+ if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) {
metrics.n_busy_slots++;
}
metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens());
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index 0d3beb313cea..9afe3c7f06a8 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -1562,6 +1562,14 @@ std::string server_task_result_metrics::to_metrics() {
"spec_decode_num_drafts_total",
"Speculative: Total speculative decoding verification steps",
(double) metrics.n_draft_verif_steps
+ }, {
+ "n_preempt_total",
+ "Preemption: Total slots parked to make room in the unified KV cache",
+ (double) metrics.n_preempt
+ }, {
+ "n_resume_total",
+ "Preemption: Total parked slots put back",
+ (double) metrics.n_resume
},
};
@@ -1586,6 +1594,14 @@ std::string server_task_result_metrics::to_metrics() {
"n_busy_slots_per_decode",
"Average number of busy slots per llama_decode() call",
(double) metrics.n_busy_slots / std::max((double) metrics.n_decode, 1.0)
+ }, {
+ "requests_preempted",
+ "Preemption: Number of requests currently parked, waiting for room in the unified KV cache",
+ (double) n_preempted_slots
+ }, {
+ "preempt_ram_bytes",
+ "Preemption: Host RAM held by parked sequences",
+ (double) preempt_ram_bytes
},
};
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 9c99143f8e19..00734924bc63 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -494,6 +494,8 @@ struct server_task_result_metrics : server_task_result {
// these are immediate stats, not accumulated (server_metrics is cumulative)
int n_processing_slots = 0;
int n_tasks_deferred = 0;
+ int n_preempted_slots = 0; // [TAG_PREEMPT] processing slots currently parked
+ size_t preempt_ram_bytes = 0; // [TAG_PREEMPT] host RAM their parked sequences hold
server_metrics metrics;
diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py
new file mode 100644
index 000000000000..0da885bcafd5
--- /dev/null
+++ b/tools/server/tests/unit/test_preempt.py
@@ -0,0 +1,269 @@
+import os
+import time
+import tempfile
+import pytest
+from utils import *
+
+# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked
+# (its sequence copied to host RAM, its cells released) instead of every slot being
+# terminated. Both tests need more than one slot and --kv-unified, which is the only
+# configuration where one slot can take another one's cells.
+
+server = ServerPreset.tinyllama2()
+
+
+class LogReader:
+ def __init__(self, path):
+ self.path = path
+ self.pos = 0
+
+ def drain(self):
+ with open(self.path) as f:
+ f.seek(self.pos)
+ content = f.read()
+ self.pos = f.tell()
+ return content
+
+
+@pytest.fixture(autouse=True)
+def create_server():
+ global server
+ server = ServerPreset.tinyllama2()
+ server.n_slots = 2
+ server.kv_unified = True
+ server.server_slots = True
+ server.temperature = 0.0
+ server.seed = 42
+ fd, server.log_path = tempfile.mkstemp(suffix=".log")
+ os.close(fd)
+ yield
+ os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None)
+ os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None)
+
+
+def _complete(n_predict: int, prompt: str = "Hi how are you"):
+ res = server.make_request("POST", "/completion", data={
+ "n_predict": n_predict,
+ "prompt": prompt,
+ "ignore_eos": True,
+ "return_tokens": True,
+ "temperature": 0.0,
+ "seed": 42,
+ })
+ return res
+
+
+def test_forced_preemption_does_not_change_the_output():
+ # Park and restore the only running slot every 8 tokens. With one request the batch
+ # has the same shape at every step whether or not the slot was parked in between, so
+ # any difference in the output is the preemption's fault and nothing else's.
+ global server
+ server.n_ctx = 512
+ server.start()
+ reference = _complete(64)
+ assert reference.status_code == 200
+ assert reference.body["timings"]["predicted_n"] == 64
+ server.stop()
+
+ os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8"
+ server.start()
+ log = LogReader(server.log_path)
+ assert "LLAMA_SERVER_PREEMPT_EVERY = 8" in log.drain()
+
+ preempted = _complete(64)
+ assert preempted.status_code == 200
+ assert preempted.body["timings"]["predicted_n"] == 64
+
+ text = log.drain()
+ assert text.count("preempted on request") >= 6
+ assert text.count("resumed after") >= 6
+
+ assert preempted.body["content"] == reference.body["content"]
+ assert preempted.body["tokens"] == reference.body["tokens"]
+
+
+def test_two_slots_that_overflow_the_pool_together_both_finish():
+ # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well
+ # under 256. Together they do not, 336 against 256. Without preemption the retry
+ # ladder ends with "Context size has been exceeded" on every processing slot; with it
+ # the smaller slot is parked until the leader finishes and its cells are purged, and
+ # then it resumes from the token it was parked on.
+ global server
+ server.n_ctx = 256
+ server.start()
+ log = LogReader(server.log_path)
+
+ n_predict = 160
+ results = parallel_function_calls([
+ (_complete, (n_predict, "Once upon a time there was a brave knight who")),
+ (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")),
+ ])
+
+ text = log.drain()
+ assert "Context size has been exceeded" not in text
+ assert "preempted:" in text
+ assert "resumed after" in text
+
+ for res in results:
+ assert res.status_code == 200
+ assert res.body["timings"]["predicted_n"] == n_predict
+ assert res.body["truncated"] is False
+ assert len(res.body["tokens"]) == n_predict
+
+
+
+_WORDS = (
+ "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor "
+ "incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud "
+ "exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure "
+ "dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. "
+ "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt "
+ "mollit anim id est laborum. "
+) * 4
+
+
+def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]:
+ """A prompt whose token count is in [n_tokens - 12, n_tokens], measured on the server."""
+ words = (salt + " " + _WORDS).split()
+ while words:
+ text = " ".join(words)
+ res = server.make_request("POST", "/tokenize", data={"content": text})
+ assert res.status_code == 200
+ n = len(res.body["tokens"])
+ if n <= n_tokens:
+ assert n >= n_tokens - 12, f"could not land near {n_tokens} tokens, got {n}"
+ return text, n
+ # about four tokens per word on this model's vocabulary
+ words = words[: len(words) - max(1, (n - n_tokens) // 8)]
+ raise AssertionError("empty prompt")
+
+
+def test_two_prompts_that_overflow_the_pool_together_both_finish():
+ # Neither slot ever generates before the pool is full: both are still processing their
+ # prompts. A prompt-processing slot is between two chunks of its prompt, which is as
+ # clean a boundary as the one between two sampled tokens, so it is parked the same way.
+ global server
+ server.n_ctx = 256
+ server.start()
+ log = LogReader(server.log_path)
+
+ prompt_a, n_a = _prompt_of_about(150, "Alpha")
+ prompt_b, n_b = _prompt_of_about(150, "Bravo")
+ n_predict = 16
+ assert n_a + n_predict <= 256 and n_b + n_predict <= 256
+ assert n_a + n_b + 2 * n_predict > 256
+
+ results = parallel_function_calls([
+ (_complete, (n_predict, prompt_a)),
+ (_complete, (n_predict, prompt_b)),
+ ])
+
+ text = log.drain()
+ assert "Context size has been exceeded" not in text
+ assert "preempted:" in text
+ assert "resumed after" in text
+
+ for res in results:
+ assert res.status_code == 200
+ assert res.body["timings"]["predicted_n"] == n_predict
+ assert len(res.body["tokens"]) == n_predict
+
+
+def test_a_generating_slot_and_a_large_prompt_both_finish():
+ # One slot is generating a long answer to a short prompt when a large prompt arrives
+ # beside it. Together they need far more than the pool has. The prompt is admitted
+ # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish.
+ # This model produces a thousand tokens a second, so the second request is sent right
+ # behind the first rather than after a delay: its prompt takes several batches to
+ # process, which is enough for the two to overlap however fast the first one runs.
+ global server
+ server.n_ctx = 256
+ server.start()
+ log = LogReader(server.log_path)
+
+ prompt_b, n_b = _prompt_of_about(150, "Charlie")
+ # b lives long enough for the two to collide: the first run of this used 16 tokens
+ # and b was finished and purged before a had grown into it
+ n_predict_a = 230
+ n_predict_b = 90
+ assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256
+ assert 8 + n_predict_a + n_b + n_predict_b > 256
+
+ def _late(n_predict, prompt):
+ time.sleep(0.02)
+ return _complete(n_predict, prompt)
+
+ results = parallel_function_calls([
+ (_complete, (n_predict_a, "Hi how are you")),
+ (_late, (n_predict_b, prompt_b)),
+ ])
+
+ text = log.drain()
+ assert "Context size has been exceeded" not in text
+ assert "preempted:" in text
+
+ assert results[0].status_code == 200
+ assert results[0].body["timings"]["predicted_n"] == n_predict_a
+ assert results[1].status_code == 200
+ assert results[1].body["timings"]["predicted_n"] == n_predict_b
+
+
+def test_preempt_ram_zero_disables_preemption():
+ # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the
+ # KV-full path ends the requests the way it always did.
+ global server
+ server.n_ctx = 256
+ os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0"
+ server.start()
+ log = LogReader(server.log_path)
+
+ n_predict = 160
+ results = parallel_function_calls([
+ (_complete, (n_predict, "Once upon a time there was a brave knight who")),
+ (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")),
+ ])
+
+ text = log.drain()
+ assert "preempted:" not in text
+ assert "Context size has been exceeded" in text
+ assert any(res.status_code != 200 for res in results)
+
+
+def test_metrics_and_slots_report_the_parked_state():
+ # A client that wants to tell a parked chat from a slow one reads /slots, and an
+ # operator reads /metrics. Both must show the preemption happening, and the counters
+ # must survive the requests finishing.
+ global server
+ server.n_ctx = 256
+ server.server_metrics = True
+ server.start()
+
+ res = server.make_request("GET", "/slots")
+ assert res.status_code == 200
+ for slot in res.body:
+ assert slot["is_preempted"] is False
+ assert slot["n_preempt"] == 0
+
+ n_predict = 160
+ results = parallel_function_calls([
+ (_complete, (n_predict, "Once upon a time there was a brave knight who")),
+ (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")),
+ ])
+ for res in results:
+ assert res.status_code == 200
+
+ res = server.make_request("GET", "/metrics")
+ assert res.status_code == 200
+ metrics = {}
+ for line in res.body.splitlines():
+ if line.startswith("llamacpp:"):
+ name, value = line.split(" ", 1)
+ metrics[name[len("llamacpp:"):]] = float(value)
+ assert metrics["n_preempt_total"] >= 1
+ assert metrics["n_resume_total"] == metrics["n_preempt_total"]
+ assert metrics["requests_preempted"] == 0
+ assert metrics["preempt_ram_bytes"] == 0
+
+ res = server.make_request("GET", "/slots")
+ assert res.status_code == 200
+ assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot"