From 5c7018ed7b419197ffda429999a2334cab5df939 Mon Sep 17 00:00:00 2001
From: Christian Lohwasser <christian.lohwasser@gmx.de>
Date: Thu, 20 Aug 2026 19:20:00 +0200
Subject: [PATCH] server : share checkpoint state and harden prompt cache OOM
 handling

Share target and draft checkpoint backing storage to avoid large deep copies when saving prompts to the in-memory cache.

Also handle `std::bad_alloc` during complete cache entry construction, including `states.push_back()`.

Assisted-by: Qwen3.8-27B
---
 common/common.cpp            |  8 ++++----
 common/common.h              | 39 ++++++++++++++++++++++++++++++++++--
 tools/server/server-task.cpp | 27 ++++++++++++++-----------
 3 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 25ca838dff0a..9b4c2f340832 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -2265,9 +2265,9 @@ void common_prompt_checkpoint::update_tgt(
 
     const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, flags);
 
-    data_tgt.resize(ckpt_size);
+    uint8_t * data = data_tgt.reset(ckpt_size);
 
-    const size_t n = llama_state_seq_get_data_ext(ctx, data_tgt.data(), ckpt_size, seq_id, flags);
+    const size_t n = llama_state_seq_get_data_ext(ctx, data, ckpt_size, seq_id, flags);
     if (n != ckpt_size) {
         GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", ckpt_size, n);
     }
@@ -2283,9 +2283,9 @@ void common_prompt_checkpoint::update_dft(
 
     const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, flags);
 
-    data_dft.resize(ckpt_size);
+    uint8_t * data = data_dft.reset(ckpt_size);
 
-    const size_t n = llama_state_seq_get_data_ext(ctx, data_dft.data(), ckpt_size, seq_id, flags);
+    const size_t n = llama_state_seq_get_data_ext(ctx, data, ckpt_size, seq_id, flags);
     if (n != ckpt_size) {
         GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", ckpt_size, n);
     }
diff --git a/common/common.h b/common/common.h
index de49dac9f63a..651c8e81dc14 100644
--- a/common/common.h
+++ b/common/common.h
@@ -16,6 +16,7 @@
 #include <map>
 #include <algorithm>
 #include <fstream>
+#include <memory>
 
 #if defined(_WIN32) && !defined(_WIN32_WINNT)
 #define _WIN32_WINNT 0x0A00
@@ -1134,6 +1135,37 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *);
 // prompt utils
 //
 
+class common_shared_state_buffer {
+public:
+    size_t size() const {
+        return storage ? storage->size() : 0;
+    }
+
+    bool empty() const {
+        return !storage || storage->empty();
+    }
+
+    const uint8_t * data() const {
+        return storage ? storage->data() : nullptr;
+    }
+
+    // Allocate fresh storage for a complete state snapshot. Copies of this
+    // object keep referencing the previous immutable snapshot; they are never
+    // modified in place. Only the newly allocated, uniquely owned buffer is
+    // exposed as mutable while the caller captures the replacement snapshot.
+    uint8_t * reset(size_t n) {
+        storage = std::make_shared<std::vector<uint8_t>>(n);
+        return storage->data();
+    }
+
+    void clear() {
+        storage.reset();
+    }
+
+private:
+    std::shared_ptr<std::vector<uint8_t>> storage;
+};
+
 struct common_prompt_checkpoint {
     int64_t n_tokens;
 
@@ -1143,8 +1175,11 @@ struct common_prompt_checkpoint {
     llama_pos pos_min;
     llama_pos pos_max;
 
-    std::vector<uint8_t> data_tgt;
-    std::vector<uint8_t> data_dft;
+    // Snapshot payloads are immutable after capture and can be shared safely
+    // between a live slot and host-memory prompt-cache entries. Copying a
+    // checkpoint therefore copies only the ref-counted handles, not the bytes.
+    common_shared_state_buffer data_tgt;
+    common_shared_state_buffer data_dft;
 
     // (optional) speculative-decoding implementation state stashed with the checkpoint
     // (e.g. eagle3's deferred-boundary g_embd row)
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index 258cdcf8fb86..cd03be2b7f33 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -1762,10 +1762,24 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
     std::vector<uint8_t> state_data_tgt;
     std::vector<uint8_t> state_data_dft;
 
-    // check if we can allocate enough memory for the new state
+    // Keep every allocation involved in creating the cache entry under the
+    // same bad_alloc guard. In particular, prompt.tokens.clone(), copying the
+    // checkpoint list (cheap shared snapshot handles) and states.push_back()
+    // can allocate too. A failure must skip caching, not terminate the server.
     try {
         state_data_tgt.resize(state_size_tgt);
         state_data_dft.resize(state_size_dft);
+
+        states.push_back({
+            /*.prompt =*/ {
+                /*.tokens      =*/ prompt.tokens.clone(),
+                /*.checkpoints =*/ prompt.checkpoints,
+            },
+            /*.data   =*/ {
+                /*.main =*/ std::move(state_data_tgt),
+                /*.drft =*/ std::move(state_data_dft),
+            },
+        });
     } catch (const std::bad_alloc & e) {
         SRV_ERR("failed to allocate memory for prompt cache state: %s\n", e.what());
 
@@ -1778,17 +1792,6 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
         return nullptr;
     }
 
-    states.push_back({
-        /*.prompt =*/ {
-            /*.tokens      =*/ prompt.tokens.clone(),
-            /*.checkpoints =*/ prompt.checkpoints,
-        },
-        /*.data   =*/ {
-            /*.main =*/ std::move(state_data_tgt),
-            /*.drft =*/ std::move(state_data_dft),
-        },
-    });
-
     return &states.back();
 }