server: preempt a slot instead of ending every conversation when the KV pool fills

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

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

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

@@ -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)<br/>(env: LLAMA_ARG_CTX_CHECKPOINTS) |
 | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)<br/>(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)<br/>(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)<br/>(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)<br/>(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)<br/>(env: LLAMA_ARG_CACHE_IDLE_SLOTS) |
 | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)<br/>(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.
 

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