server: park unified requests with proactive KV watermarks

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"},

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

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<uint8_t> 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);
+            }
         }
     }
 

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<std::vector<uint8_t>> 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<uint8_t> 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();

tools/server/README.md

@@ -2130,3 +2130,73 @@ You can specify default preferences for the web UI using `--ui-config <JSON conf
 > **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.

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,