diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d0..39818efff274 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1717,6 +1717,36 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.kv_unified = value; } ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); + add_opt(common_arg( + {"--preempt-high"}, "PERCENT", + "unified server pool high watermark (default: 94, 0 disables preemption)", + [](common_params & params, int value) { + if (value < 0 || value > 100) { + throw std::invalid_argument("--preempt-high must be in [0, 100]"); + } + params.preempt_high = value; + } + ).set_env("LLAMA_ARG_PREEMPT_HIGH").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--preempt-low"}, "PERCENT", + "unified server pool restore watermark (default: 80, must be below high)", + [](common_params & params, int value) { + if (value <= 0 || value >= 100) { + throw std::invalid_argument("--preempt-low must be in (0, 100)"); + } + params.preempt_low = value; + } + ).set_env("LLAMA_ARG_PREEMPT_LOW").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--preempt-ram"}, "MIB", + "maximum host RAM for parked sequence snapshots (default: 8192, 0 disables parking)", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("--preempt-ram must be nonnegative"); + } + params.preempt_ram_mib = value; + } + ).set_env("LLAMA_ARG_PREEMPT_RAM").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"--cache-idle-slots"}, {"--no-cache-idle-slots"}, diff --git a/common/common.h b/common/common.h index de49dac9f63a..07a17d142683 100644 --- a/common/common.h +++ b/common/common.h @@ -561,6 +561,9 @@ struct common_params { bool ctx_shift = false; // context shift on infinite text generation bool swa_full = false; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) bool kv_unified = false; // enable unified KV cache + int32_t preempt_high = 94; // unified server pool high watermark, percent + int32_t preempt_low = 80; // restore only below this watermark, percent + int32_t preempt_ram_mib = 8192; // maximum host memory for parked sequence snapshots bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix bool verbose_prompt = false; // print prompt tokens before generation diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b7..66940d4fc61a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2617,8 +2617,42 @@ class llama_io_read_host : public llama_io_read_i { ~llama_io_read_host() { // flush the reads - for (const auto & rinfo : rinfos) { - ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); + for (size_t i = 0; i < rinfos.size();) { + auto * tensor = rinfos[i].tensor; + size_t end = i + 1; + while (end < rinfos.size() && rinfos[end].tensor == tensor) { + end++; + } + const size_t tensor_bytes = ggml_nbytes(tensor); + auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + // A fragmented sequence can require thousands of synchronous device + // transfers per layer. For bounded tensors, stage the tensor once and + // preserve every byte belonging to other sequences. Bound scratch RAM + // and leave ordinary contiguous transfers on their original fast path. + if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && + !ggml_backend_buffer_is_host(buffer)) { + std::vector staging; + try { + staging.resize(tensor_bytes); + } catch (const std::bad_alloc &) { + // Fall back to the individual transfers below. + } + if (!staging.empty()) { + ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); + for (size_t j = i; j < end; ++j) { + const auto & rinfo = rinfos[j]; + GGML_ASSERT(rinfo.offset <= tensor_bytes && rinfo.size <= tensor_bytes - rinfo.offset); + memcpy(staging.data() + rinfo.offset, rinfo.ptr, rinfo.size); + } + ggml_backend_tensor_set(tensor, staging.data(), 0, tensor_bytes); + i = end; + continue; + } + } + for (; i < end; ++i) { + const auto & rinfo = rinfos[i]; + ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); + } } } diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index d5548afba179..428a92529811 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,6 +73,14 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); + // A fragmented restore may stage a whole device tensor. Check every + // sequence byte-for-byte, including the neighbours that must be preserved. + std::vector> before(params.n_parallel); + for (int s = 0; s < params.n_parallel; ++s) { + before[s].resize(llama_state_seq_get_size(ctx, s)); + GGML_ASSERT(llama_state_seq_get_data(ctx, before[s].data(), before[s].size(), s) == before[s].size()); + } + // clear seq 1 to create a "hole" in the KV cache (fragmentation) // 0.20.20.20.2.... llama_memory_t mem = llama_get_memory(ctx); @@ -96,6 +104,13 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : restored state into seq 1, %zu bytes\n", __func__, nset); + for (int s = 0; s < params.n_parallel; ++s) { + std::vector after(llama_state_seq_get_size(ctx, s)); + GGML_ASSERT(llama_state_seq_get_data(ctx, after.data(), after.size(), s) == after.size()); + GGML_ASSERT(before[s] == after); + } + fprintf(stderr, "%s : all %d sequence snapshots are byte-identical after restore\n", __func__, params.n_parallel); + // Verify we can decode with the restored state // Generate one token to verify the restored state is usable auto sparams = llama_sampler_chain_default_params(); diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa9..5789f125998f 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2130,3 +2130,73 @@ You can specify default preferences for the web UI using `--ui-config **Note:** The old flags `--webui-config` and `--webui-config-file` are deprecated but still work as aliases. You may find available preferences in [settings-keys.ts](../ui/src/lib/constants/settings-keys.ts). + +### Unified-pool preemption + +With `--kv-unified`, llama-server proactively parks independent text completion +requests in host RAM when the projected next batch exceeds `--preempt-high` +(default **94 percent**). It parks newest arrivals first until the projection is +at or below `--preempt-low` (default **80 percent**), or only one runnable request +remains. The largest resident request is protected. Requests parked three times +consecutively are considered after other eligible victims. Shared-prompt parents +and children (`n_cmpl > 1`) are excluded. Preemption is disabled on embedding +servers, servers with a multimodal projector, and contexts without sequence removal. + +The projection includes each resident sequence, its next sampled token and the +maximum speculative draft depth. Target and draft contexts have separate pools; +use the smaller capacity without charging the draft pool twice to the target. +Pending prompt chunks share the remaining batch budget after generation; that +budget is reserved before batch preparation. Counts remain conservative for shared +prefixes, sliding-window caches and prompt-cache reuse. +Idle prompt caches are reclaimed before live requests are parked. + +Parking saves full sequence state from both contexts with the host sequence-state +APIs and removes the sequences from device memory. The task, sampler and grammar, +stream offsets, speculative checkpoints and per-sequence MTP hidden-state +carryover remain alive. No generated tokens are replayed or discarded. A streamed +response pauses and resumes on the same connection. Cancellation releases parked +RAM. Tool call text can pause mid-generation; execution of tools and subsequent +requests remain the client's responsibility. + +Restoration tries the most-parked request first, then the one waiting longest, +skipping entries that do not fit. **Every restore must project at or below LOW**; +the LOW-to-HIGH band cannot admit restores. A request whose next step cannot fit +under LOW on its own is not parked, to avoid an unresumable waiter. HIGH is not a +per-request context limit: a protected solo request may grow to its normal context +limit, with speculative depth reduced near that limit by the existing decoder. + +`--preempt-ram MIB` bounds the total serialized target and draft snapshots (default +8192 MiB). It does not bound existing sampler state or the separate prompt cache. +When the budget or host allocation is exhausted, the scheduler declines that +victim. If remaining eligible victims cannot free enough space, the existing +context-full error behavior remains. Fragmented device restores may additionally use up to 64 MiB of temporary +staging memory per tensor, outside the parked-snapshot budget. A failed restore clears any partial device +state and retains the host snapshot for retry. Persistent restore failures can +leave a request waiting indefinitely. + +Set `--preempt-high 0` to disable scheduling or `--preempt-ram 0` to decline all +parks. HIGH must be in [0, 100]; LOW in (0, 100) and strictly below a nonzero HIGH. +The corresponding environment variables are `LLAMA_ARG_PREEMPT_HIGH`, +`LLAMA_ARG_PREEMPT_LOW`, and `LLAMA_ARG_PREEMPT_RAM`. + +For regression testing, `LLAMA_SERVER_PREEMPT_EVERY=N` forces one eligible request +to park at each N-token generation boundary. Speculative acceptance can cross the +boundary by more than one token. The hook permits parking the last or largest +request, retains it across scheduler iterations, and observes the RAM and LOW +constraints. Zero disables the hook. Compare with the same seed, temperature, +batching and sampling settings; GPU batching and different physical KV layouts +can change floating-point results even with greedy sampling. + +With `--metrics`, `/metrics` exports `llamacpp:preemptions_total`, +`preempt_restores_total`, `preempt_forced_total`, `preempt_ram_denied_total`, +`preempt_restore_failures_total`, `preempt_unused_cells_total` and +`preempt_copy_seconds_total` counters (all with the `llamacpp:` prefix). +Gauges expose `preempt_parked`, `preempt_ram_bytes`, `preempt_resident_cells`, +`preempt_projected_cells`, `preempt_high_cells`, `preempt_low_cells`, and +`preempt_restore_max_cells`. The unused-cell counter sums free cells immediately +before each park, including forced parks, using the conservative resident count. + +`/slots` adds `is_parked`, `preempt_count`, `preempt_bytes`, `preempt_parked_ms`, +`preempt_restore_projected_cells`, `preempt_high_cells`, and `preempt_low_cells`. +A parked request still has `is_processing: true`; its prompt-token count is its +logical length, not resident device occupancy. Per-request fields reset on release. diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f8ea82ef4cf5..aec1c09bea36 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -433,6 +433,21 @@ struct server_slot_stats { struct server_metrics { int64_t t_start = 0; + uint64_t preempt_total = 0; + uint64_t preempt_restore_total = 0; + uint64_t preempt_forced_total = 0; + uint64_t preempt_ram_denied_total = 0; + uint64_t preempt_restore_fail_total = 0; + uint64_t preempt_unused_cells_total = 0; + uint64_t preempt_copy_us = 0; + uint64_t preempt_parked = 0; + uint64_t preempt_ram_bytes = 0; + uint64_t preempt_resident_cells = 0; + uint64_t preempt_projected_cells = 0; + uint64_t preempt_high_cells = 0; + uint64_t preempt_low_cells = 0; + uint64_t preempt_restore_max_cells = 0; + struct bucket { uint64_t count = 0; // number of tokens uint64_t steps = 0; // number of decode steps, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..3357070c4203 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -59,6 +59,7 @@ enum slot_state { SLOT_STATE_PROCESSING_PROMPT, SLOT_STATE_DONE_PROMPT, SLOT_STATE_GENERATING, + SLOT_STATE_PARKED, // task and sampler are alive; sequence memory is on the host }; struct server_slot; // forward declaration @@ -248,6 +249,20 @@ struct server_slot { // state slot_state state = SLOT_STATE_IDLE; + slot_state preempt_state = SLOT_STATE_IDLE; + std::vector preempt_tgt; + std::vector preempt_dft; + uint64_t preempt_count = 0; + uint64_t preempt_streak = 0; + uint64_t preempt_forced_at = 0; + uint64_t preempt_cells = 0; + uint64_t preempt_restore_projected = 0; + int64_t preempt_since = 0; + + size_t preempt_bytes() const { + return preempt_tgt.size() + preempt_dft.size(); + } + server_prompt prompt; bool prompt_save(server_prompt_cache & prompt_cache) const { @@ -323,6 +338,13 @@ struct server_slot { void reset() { SLT_DBG(*this, "%s", "\n"); + std::vector().swap(preempt_tgt); + std::vector().swap(preempt_dft); + preempt_state = SLOT_STATE_IDLE; + preempt_count = preempt_streak = preempt_forced_at = preempt_cells = 0; + preempt_restore_projected = 0; + preempt_since = 0; + spec_is_replay = false; last_nl_pos = 0; @@ -503,6 +525,10 @@ struct server_slot { t_last_used = ggml_time_us(); + // A cancelled parked task must not leave tokens claiming a resident cache. + if (state == SLOT_STATE_PARKED) { + prompt_clear(); + } state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough @@ -645,6 +671,11 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, + {"is_parked", state == SLOT_STATE_PARKED}, + {"preempt_count", preempt_count}, + {"preempt_bytes", preempt_bytes()}, + {"preempt_parked_ms", state == SLOT_STATE_PARKED ? (ggml_time_us() - preempt_since) / 1000.0 : 0.0}, + {"preempt_restore_projected_cells", preempt_restore_projected}, }; const auto & ptask = task ? task : task_prev; @@ -867,6 +898,7 @@ struct server_context_impl { std::unique_ptr prompt_cache; server_metrics metrics; + uint64_t preempt_every = 0; // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync // note: kept out of server_metrics, which is copied as-is into the task result @@ -963,6 +995,27 @@ struct server_context_impl { const bool is_resume = sleeping; params_base = params; + if (params.preempt_high != 0 && params.preempt_low >= params.preempt_high) { + SRV_ERR("%s", "--preempt-low must be below --preempt-high\n"); + return false; + } + preempt_every = 0; + if (const char * value = getenv("LLAMA_SERVER_PREEMPT_EVERY")) { + try { + const std::string str(value); + size_t end = 0; + if (str.empty() || str.find_first_not_of("0123456789") != std::string::npos) { + throw std::invalid_argument("invalid interval"); + } + preempt_every = std::stoull(str, &end); + if (end != str.size()) { + throw std::invalid_argument("invalid interval"); + } + } catch (const std::exception &) { + SRV_ERR("%s", "LLAMA_SERVER_PREEMPT_EVERY must be a nonnegative integer\n"); + return false; + } + } const auto output_limits = server_output_limits(params_base); params_base.n_outputs_max = output_limits.total; params_base.n_outputs_max_per_seq = output_limits.per_seq; @@ -2384,6 +2437,7 @@ struct server_context_impl { } break; case SERVER_TASK_TYPE_METRICS: { + preempt_update_metrics(); int n_processing_slots = 0; for (server_slot & slot : slots) { @@ -2406,6 +2460,7 @@ struct server_context_impl { } break; case SERVER_TASK_TYPE_SLOT_GET: { + preempt_update_metrics(); json slots_data = json::array(); int n_idle_slots = 0; @@ -2416,6 +2471,8 @@ struct server_context_impl { } slots_data.push_back(slot.to_json(slots_debug == 0)); + slots_data.back()["preempt_high_cells"] = metrics.preempt_high_cells; + slots_data.back()["preempt_low_cells"] = metrics.preempt_low_cells; } SRV_DBG("n_idle_slots = %d\n", n_idle_slots); @@ -2674,6 +2731,277 @@ struct server_context_impl { }; #endif + bool preempt_enabled() const { + return params_base.kv_unified && params_base.preempt_high > 0 && + llama_get_memory(ctx_tgt) && ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO && + (!ctx_dft || ctx_dft_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) && + !mctx && !params_base.embedding; + } + + uint64_t preempt_capacity() const { + // The draft has a separate pool. Charge the same sequence positions against + // the smaller pool; do not double-charge MTP's cells in the target pool. + return ctx_dft ? std::min(llama_n_ctx(ctx_tgt), llama_n_ctx(ctx_dft)) : llama_n_ctx(ctx_tgt); + } + + uint64_t preempt_resident(const server_slot & slot) const { + if (slot.state == SLOT_STATE_PARKED) { + return 0; + } + // A conservative upper bound for shared prefixes and sliding-window memory. + // This reads memory metadata without synchronizing the GPU, including replay KV. + auto last = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id); + if (ctx_dft) { + last = std::max(last, llama_memory_seq_pos_max(llama_get_memory(ctx_dft), slot.id)); + } + return last < 0 ? 0 : uint64_t(last) + 1; + } + + uint64_t preempt_need(const server_slot & slot, bool restoring = false) const { + if (slot.state == SLOT_STATE_PARKED && !restoring) { + return 0; + } + const auto state = restoring ? slot.preempt_state : slot.state; + const uint64_t resident = restoring ? slot.preempt_cells : preempt_resident(slot); + if (state == SLOT_STATE_IDLE || state == SLOT_STATE_WAIT_OTHER) { + return resident; + } + const uint64_t draft = slot.can_speculate() ? common_speculative_n_max(¶ms_base.speculative) : 0; + if (state == SLOT_STATE_STARTED) { + // Cache lookup/load already ran during slot assignment. Reserve a + // prompt chunk before pre_decode initializes prefix reuse. Counting + // the full input here would prevent long prompts from ever resuming. + return resident + std::min(slot.task->n_tokens(), llama_n_batch(ctx_tgt)) + 1 + draft; + } + if (state == SLOT_STATE_PROCESSING_PROMPT) { + const uint64_t remaining = std::max(0, slot.task->n_tokens() - slot.prompt.n_tokens()); + return resident + std::min(remaining, llama_n_batch(ctx_tgt)) + 1 + draft; + } + return resident + 1 + std::min(draft, std::max(0, slot.get_n_draft_max())); + } + + uint64_t preempt_projected(const server_slot * restoring = nullptr) const { + uint64_t result = 0; + uint64_t prompt_batch = 0; + uint64_t generation_batch = 0; + const uint64_t n_batch = llama_n_batch(ctx_tgt); + for (const auto & slot : slots) { + const bool restore = &slot == restoring; + result += preempt_need(slot, restore); + const auto state = restore ? slot.preempt_state : slot.state; + if (state == SLOT_STATE_STARTED) { + prompt_batch += std::min(slot.task->n_tokens(), n_batch); + } else if (state == SLOT_STATE_PROCESSING_PROMPT) { + prompt_batch += std::min(std::max(0, slot.task->n_tokens() - slot.prompt.n_tokens()), n_batch); + } else if (state == SLOT_STATE_GENERATING) { + generation_batch += 1 + (slot.can_speculate() ? common_speculative_n_max(¶ms_base.speculative) : 0); + } + } + // All prefills share the remainder of ONE batch after generation. Charging + // a full batch for every prefilling slot parks clients needlessly early. + return result - prompt_batch + std::min(prompt_batch, n_batch > generation_batch ? n_batch - generation_batch : 0); + } + + void preempt_update_metrics() { + metrics.preempt_parked = metrics.preempt_ram_bytes = metrics.preempt_resident_cells = 0; + if (!preempt_enabled()) { + return; + } + metrics.preempt_high_cells = preempt_capacity() * params_base.preempt_high / 100; + metrics.preempt_low_cells = preempt_capacity() * params_base.preempt_low / 100; + metrics.preempt_projected_cells = preempt_projected(); + for (const auto & slot : slots) { + metrics.preempt_parked += slot.state == SLOT_STATE_PARKED; + metrics.preempt_ram_bytes += slot.preempt_bytes(); + metrics.preempt_resident_cells += preempt_resident(slot); + } + } + + bool preempt_eligible(const server_slot & slot) const { + return slot.is_processing() && slot.state != SLOT_STATE_PARKED && + !slot.task->is_parent() && !slot.task->is_child() && + slot.task->need_sampling() && + (slot.state == SLOT_STATE_GENERATING || slot.state == SLOT_STATE_PROCESSING_PROMPT || + slot.state == SLOT_STATE_STARTED); + } + + bool preempt_park(server_slot & slot, bool forced) { + const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + if (size_tgt == 0 || (ctx_dft && size_dft == 0)) { + return false; + } + uint64_t used = 0; + for (const auto & other : slots) { + used += other.preempt_bytes(); + } + const uint64_t limit = uint64_t(params_base.preempt_ram_mib) * 1024 * 1024; + if (used > limit || size_tgt > limit - used || size_dft > limit - used - size_tgt) { + metrics.preempt_ram_denied_total++; + return false; + } + + const auto start = ggml_time_us(); + // Allocate and copy both contexts before removing either one. A failed copy + // leaves the runnable sequence untouched. Sampler, grammar, MTP pending_h, + // accepted/replay drafts, checkpoints and stream offsets remain in place. + std::vector target, draft; + try { + target.resize(size_tgt); + draft.resize(size_dft); + } catch (const std::bad_alloc &) { + metrics.preempt_ram_denied_total++; + return false; + } + if (llama_state_seq_get_data_ext(ctx_tgt, target.data(), size_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt || + (ctx_dft && llama_state_seq_get_data_ext(ctx_dft, draft.data(), size_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft)) { + return false; + } + slot.preempt_cells = preempt_resident(slot); + uint64_t resident = 0; + for (const auto & other : slots) { + resident += preempt_resident(other); + } + const auto unused = preempt_capacity() - std::min(preempt_capacity(), resident); + slot.preempt_tgt = std::move(target); + slot.preempt_dft = std::move(draft); + slot.mem.seq_rm(slot.id, -1, -1); + if (spec) { + common_speculative_get_draft_params(spec.get(), slot.id).drafting = false; + } + slot.preempt_state = slot.state; + slot.state = SLOT_STATE_PARKED; + slot.preempt_since = ggml_time_us(); + slot.preempt_count++; + slot.preempt_streak++; + metrics.preempt_total++; + metrics.preempt_forced_total += forced; + metrics.preempt_unused_cells_total += unused; + metrics.preempt_copy_us += ggml_time_us() - start; + SLT_INF(slot, "preemption parked: cells=%" PRIu64 ", unused=%" PRIu64 ", bytes=%zu, forced=%d\n", + slot.preempt_cells, unused, slot.preempt_bytes(), forced); + return true; + } + + bool preempt_restore(server_slot & slot, uint64_t projected) { + const auto start = ggml_time_us(); + if (llama_state_seq_set_data_ext(ctx_tgt, slot.preempt_tgt.data(), slot.preempt_tgt.size(), slot.id, + LLAMA_STATE_SEQ_FLAGS_NONE) != slot.preempt_tgt.size() || + (ctx_dft && llama_state_seq_set_data_ext(ctx_dft, slot.preempt_dft.data(), slot.preempt_dft.size(), slot.id, + LLAMA_STATE_SEQ_FLAGS_NONE) != slot.preempt_dft.size())) { + // A partial restore must not occupy unaccounted cells. Keep the host + // snapshot and task so a later scheduler pass can retry it. + slot.mem.seq_rm(slot.id, -1, -1); + metrics.preempt_restore_fail_total++; + return false; + } + slot.state = slot.preempt_state; + slot.preempt_restore_projected = projected; + std::vector().swap(slot.preempt_tgt); + std::vector().swap(slot.preempt_dft); + metrics.preempt_restore_total++; + metrics.preempt_restore_max_cells = std::max(metrics.preempt_restore_max_cells, projected); + metrics.preempt_copy_us += ggml_time_us() - start; + SLT_INF(slot, "preemption restored: projected=%" PRIu64 ", low=%" PRIu64 ", pause_ms=%.3f\n", + projected, preempt_capacity() * params_base.preempt_low / 100, + (ggml_time_us() - slot.preempt_since) / 1000.0); + return true; + } + + void preempt_schedule() { + if (!preempt_enabled()) { + return; + } + const uint64_t high = preempt_capacity() * params_base.preempt_high / 100; + const uint64_t low = preempt_capacity() * params_base.preempt_low / 100; + uint64_t projected = preempt_projected(); + bool has_parked = false; + for (const auto & slot : slots) { + has_parked |= slot.state == SLOT_STATE_PARKED; + } + if (projected > high || has_parked) { + // Idle cache entries have no waiting client and are cheaper to reclaim. + while (try_clear_idle_slots()) {} + projected = preempt_projected(); + } + if (has_parked) { + std::vector waiting; + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_PARKED) { + waiting.push_back(&slot); + } + } + std::sort(waiting.begin(), waiting.end(), [](const server_slot * a, const server_slot * b) { + if (a->preempt_count != b->preempt_count) { + return a->preempt_count > b->preempt_count; + } + return a->preempt_since < b->preempt_since; + }); + for (auto * slot : waiting) { + const uint64_t after = preempt_projected(slot); + if (after <= low && preempt_restore(*slot, after)) { + projected = preempt_projected(); + } + } + } + if (projected > high) { + std::vector candidates; + server_slot * largest = nullptr; + size_t runnable = 0; + for (auto & slot : slots) { + if (slot.is_processing() && slot.state != SLOT_STATE_PARKED) { + runnable++; + if (!largest || preempt_resident(slot) > preempt_resident(*largest)) { + largest = &slot; + } + } + if (preempt_eligible(slot)) { + candidates.push_back(&slot); + } + } + // A task id is assigned on arrival, unlike prompt-start time which + // changes when a request waits for a slot or begins prefill. + std::sort(candidates.begin(), candidates.end(), [](const server_slot * a, const server_slot * b) { + if ((a->preempt_streak >= 3) != (b->preempt_streak >= 3)) { + return a->preempt_streak < 3; + } + return a->task->id > b->task->id; + }); + for (auto * slot : candidates) { + if (projected <= low || runnable <= 1) { + break; + } + // Never park an item that cannot subsequently fit under LOW on + // its own. The largest running request may consume the full pool. + if (slot == largest || preempt_need(*slot) > low) { + continue; + } + if (preempt_park(*slot, false)) { + for (auto & other : slots) { + if (&other != slot && other.state != SLOT_STATE_PARKED) { + other.preempt_streak = 0; + } + } + runnable--; + projected = preempt_projected(); + } + } + } + // Test hook: a real eviction across scheduler iterations, including for + // a solo request. Only this hook can park the last/longest runnable slot. + if (preempt_every) { + for (auto & slot : slots) { + if (preempt_eligible(slot) && slot.state == SLOT_STATE_GENERATING && + slot.stats.n_gen / preempt_every > slot.preempt_forced_at && preempt_need(slot) <= low) { + if (preempt_park(slot, true)) { + slot.preempt_forced_at = slot.stats.n_gen / preempt_every; + } + break; + } + } + } + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -2717,6 +3045,7 @@ struct server_context_impl { try { scoped_timer t(t_pre_decode, n_pre_decode); + preempt_schedule(); pre_decode(); batch.render(); } catch (const std::exception & e) { @@ -3000,7 +3329,7 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - if (!slot.is_processing()) { + if (!slot.is_processing() || slot.state == SLOT_STATE_PARKED) { return; } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..9a0ebdbd8d47 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1522,6 +1522,13 @@ json server_task_result_metrics::to_json() { // metrics definition: https://prometheus.io/docs/practices/naming/#metric-names std::string server_task_result_metrics::to_metrics() { const std::vector counters = { + {"preemptions_total", "Successful sequence parks", (double) metrics.preempt_total}, + {"preempt_restores_total", "Successful sequence restores", (double) metrics.preempt_restore_total}, + {"preempt_forced_total", "Parks caused by the test interval", (double) metrics.preempt_forced_total}, + {"preempt_ram_denied_total", "Parks denied by the host RAM limit or allocation failure", (double) metrics.preempt_ram_denied_total}, + {"preempt_restore_failures_total", "Sequence restores to retry", (double) metrics.preempt_restore_fail_total}, + {"preempt_unused_cells_total", "Sum of unused pool cells immediately before parks", (double) metrics.preempt_unused_cells_total}, + {"preempt_copy_seconds_total", "Time copying and removing sequence snapshots", metrics.preempt_copy_us / 1.e6}, { "prompt_tokens_total", "Number of prompt tokens processed, excluding cached tokens", @@ -1566,6 +1573,13 @@ std::string server_task_result_metrics::to_metrics() { }; const std::vector gauges = { + {"preempt_parked", "Currently parked requests", (double) metrics.preempt_parked}, + {"preempt_ram_bytes", "Host bytes used by parked snapshots", (double) metrics.preempt_ram_bytes}, + {"preempt_resident_cells", "Conservative resident sequence cell count", (double) metrics.preempt_resident_cells}, + {"preempt_projected_cells", "Projected cells including the next prefill or draft", (double) metrics.preempt_projected_cells}, + {"preempt_high_cells", "High watermark in cells", (double) metrics.preempt_high_cells}, + {"preempt_low_cells", "Low watermark in cells", (double) metrics.preempt_low_cells}, + {"preempt_restore_max_cells", "Largest projected cell count at any restore", (double) metrics.preempt_restore_max_cells}, { "prompt_tokens_seconds", "Average prompt throughput in tokens/s", diff --git a/tools/server/tests/unit/test_preempt_watermark.py b/tools/server/tests/unit/test_preempt_watermark.py new file mode 100644 index 000000000000..04b592fccf6b --- /dev/null +++ b/tools/server/tests/unit/test_preempt_watermark.py @@ -0,0 +1,179 @@ +"""Unified-pool parking integration tests, using only stories260K.""" +import os +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest +from utils import ServerPreset, TMP_DIR, download_file + + +@pytest.fixture(scope="session", autouse=True) +def load_server_presets(): + # This module needs only the tiny model, not every server preset. + return download_file("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf", + os.path.join(TMP_DIR, "stories260K.gguf")) + + +def start_server(model, *, ctx=512): + server = ServerPreset.tinyllama2() + server.model_hf_repo = None + server.model_hf_file = None + server.model_file = model + server.n_ctx = ctx + server.n_slots = 2 + server.n_batch = 32 + server.n_ubatch = 32 + server.n_threads = 2 + server.kv_unified = True + server.server_continuous_batching = True + server.server_metrics = True + server.server_slots = True + server.cache_ram = 0 + server.start() + return server + + +def metrics(server): + text = server.make_request("GET", "/metrics").body + return {line.split()[0].removeprefix("llamacpp:"): float(line.split()[1]) + for line in text.splitlines() if line and not line.startswith("#")} + + +def request(prompt="Once upon a time", n=160): + return dict(prompt=prompt, n_predict=n, temperature=0, seed=1234, + ignore_eos=True, cache_prompt=False, return_tokens=True) + + +def test_forced_knob_identity(load_server_presets, monkeypatch): + monkeypatch.setenv("LLAMA_SERVER_PREEMPT_EVERY", "0") + server = start_server(load_server_presets) + baseline = server.make_request("POST", "/completion", request()) + assert baseline.status_code == 200 + server.stop() + monkeypatch.setenv("LLAMA_SERVER_PREEMPT_EVERY", "32") + server = start_server(load_server_presets) + parked = server.make_request("POST", "/completion", request()) + assert parked.status_code == 200 + assert parked.body["content"].encode() == baseline.body["content"].encode() + assert parked.body["tokens"] == baseline.body["tokens"] + m = metrics(server) + assert m["preempt_forced_total"] == 4 + assert m["preempt_restores_total"] == 4 + assert m["preempt_parked"] == m["preempt_ram_bytes"] == 0 + + +@pytest.mark.parametrize("high,low", [(94, 80), (97, 90), (99, 98)]) +def test_pressure_finishes_and_honours_band(load_server_presets, monkeypatch, high, low): + monkeypatch.setenv("LLAMA_ARG_PREEMPT_HIGH", str(high)) + monkeypatch.setenv("LLAMA_ARG_PREEMPT_LOW", str(low)) + server = start_server(load_server_presets) + # Each request fits by itself; the live total cannot fit in 512 cells. + bodies = [request("Once " * 100 + ending, n=300) + for ending in ("a fox went home", "a cat found a tree")] + samples = [] + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(server.make_request, "POST", "/completion", body) for body in bodies] + while not all(f.done() for f in futures): + samples.extend(server.make_request("GET", "/slots").body) + time.sleep(0.005) + results = [f.result() for f in futures] + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 300 + assert not res.body["truncated"] + m = metrics(server) + assert m["preemptions_total"] >= 1 + assert m["preemptions_total"] == m["preempt_restores_total"] + assert m["preempt_restore_max_cells"] <= m["preempt_low_cells"] + assert m["preempt_high_cells"] == 512 * high // 100 + assert m["preempt_low_cells"] == 512 * low // 100 + assert any(s["is_parked"] for s in samples) + assert all(s["is_processing"] for s in samples if s["is_parked"]) + assert all(s["preempt_restore_projected_cells"] <= s["preempt_low_cells"] for s in samples) + assert m["preempt_parked"] == m["preempt_ram_bytes"] == 0 + assert m["preempt_restore_failures_total"] == 0 + + +def test_metrics_exposed(load_server_presets): + server = start_server(load_server_presets) + m = metrics(server) + for name in ("preemptions_total", "preempt_restores_total", "preempt_forced_total", + "preempt_ram_denied_total", "preempt_restore_failures_total", + "preempt_unused_cells_total", "preempt_copy_seconds_total", + "preempt_parked", "preempt_ram_bytes", "preempt_resident_cells", + "preempt_projected_cells", "preempt_restore_max_cells"): + assert name in m + assert m[name] == 0 + slots = server.make_request("GET", "/slots").body + assert len(slots) == 2 + assert all(not s["is_parked"] for s in slots) + + +def test_zero_ram_declines_forced_parks(load_server_presets, monkeypatch): + monkeypatch.setenv("LLAMA_ARG_PREEMPT_RAM", "0") + monkeypatch.setenv("LLAMA_SERVER_PREEMPT_EVERY", "32") + server = start_server(load_server_presets) + res = server.make_request("POST", "/completion", request()) + assert res.status_code == 200 + m = metrics(server) + assert m["preemptions_total"] == 0 + assert m["preempt_ram_denied_total"] > 0 + + +def test_long_prefills_can_wait(load_server_presets): + server = start_server(load_server_presets) + bodies = [request("Once " * 420 + ending, n=32) for ending in ("a cat", "a dog")] + with ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(lambda body: server.make_request("POST", "/completion", body), bodies)) + assert all(r.status_code == 200 for r in results), [r.body for r in results] + assert all(r.body["tokens_predicted"] == 32 for r in results) + m = metrics(server) + assert m["preemptions_total"] >= 1 + assert m["preemptions_total"] == m["preempt_restores_total"] + + +def test_shared_prompt_is_not_parked(load_server_presets, monkeypatch): + monkeypatch.setenv("LLAMA_SERVER_PREEMPT_EVERY", "16") + server = start_server(load_server_presets) + body = request(n=64) + body["n_cmpl"] = 2 + res = server.make_request("POST", "/completion", body) + assert res.status_code == 200 + assert len(res.body) == 2 + assert all(item["tokens_predicted"] == 64 for item in res.body) + assert metrics(server)["preemptions_total"] == 0 + + +def test_cancel_parked_stream_releases_ram(load_server_presets): + import requests + server = start_server(load_server_presets) + url = f"http://{server.server_host}:{server.server_port}" + responses = [] + try: + with ThreadPoolExecutor(max_workers=2) as pool: + def start(i): + body = request("Once " * 100 + ("a dog" if i else "a cat"), n=350) + body.update(stream=True, id_slot=i) + return requests.post(url + "/completion", json=body, stream=True, timeout=20) + responses = list(pool.map(start, range(2))) + deadline = time.monotonic() + 10 + parked = None + while time.monotonic() < deadline: + slots = server.make_request("GET", "/slots").body + parked = next((s for s in slots if s["is_parked"]), None) + if parked: + break + time.sleep(.002) + assert parked is not None + assert parked["preempt_bytes"] > 0 + responses[parked["id"]].close() + while time.monotonic() < deadline: + m = metrics(server) + if m["preempt_ram_bytes"] == 0: + break + time.sleep(.01) + assert m["preempt_ram_bytes"] == 0 + assert m["preempt_parked"] == 0 + finally: + for response in responses: + response.close()