server: make prompt cache state swaps atomic

Author: Regrad

tests/test-server-prompt-cache.cpp

@@ -8,9 +8,10 @@
 static server_prompt make_prompt(
         int64_t n_tokens,
         std::initializer_list<int64_t> checkpoints,
-        size_t checkpoint_size = 0) {
+        size_t checkpoint_size = 0,
+        llama_token token = 1) {
     server_prompt prompt {
-        server_tokens(llama_tokens(n_tokens, 1), false),
+        server_tokens(llama_tokens(n_tokens, token), false),
         {},
     };
 
@@ -23,6 +24,30 @@ static server_prompt make_prompt(
     return prompt;
 }
 
+static bool cache_contains(
+        const server_prompt_cache & cache,
+        const server_prompt_cache_state * state) {
+    for (const auto & cur : cache.states) {
+        if (&cur == state) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+static bool cache_contains_token(
+        const server_prompt_cache & cache,
+        llama_token token) {
+    for (const auto & cur : cache.states) {
+        if (!cur.prompt.tokens.empty() && cur.prompt.tokens[0] == token) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 int main() {
     {
         auto prompt = make_prompt(20000, {4096, 12000, 18000});
@@ -102,5 +127,78 @@ int main() {
         assert(state->prompt.checkpoints.size() == prompt.checkpoints.size());
     }
 
+    {
+        server_prompt_cache cache(4, 10000);
+        auto live   = make_prompt(1000, {100});
+        auto cached = make_prompt(900, {});
+        server_tokens target(llama_tokens(950, 1), false);
+
+        auto * cached_state = cache.alloc(cached, 512*1024, 0);
+        auto * selected = cache.find_better(live, target, true, 1);
+
+        assert(selected == cached_state);
+    }
+
+    {
+        constexpr size_t mib = 1024*1024;
+
+        server_prompt_cache cache(1, 10000);
+        auto cached = make_prompt(900, {}, 0, 1);
+        auto live   = make_prompt(1000, {}, 0, 2);
+        server_tokens target(llama_tokens(950, 1), false);
+
+        auto * cached_state = cache.alloc(cached, 512*1024, 0);
+        auto * selected = cache.find_better(live, target, true, 1);
+        assert(selected == cached_state);
+
+        auto * saved_state = cache.alloc(live, 768*1024, 0);
+        assert(cache.finalize(saved_state, &target, true, selected));
+
+        // The selected entry is pinned during the swap. The cache may exceed
+        // its steady-state budget until that entry is restored and removed.
+        assert(saved_state != nullptr);
+        assert(cache_contains(cache, cached_state));
+        assert(cache_contains(cache, saved_state));
+        assert(cache.size() == 1280*1024);
+        assert(cache.size() > mib);
+    }
+
+    {
+        server_prompt_cache cache(2, 10000);
+        auto useful = make_prompt(600, {}, 0, 1);
+        auto useless = make_prompt(600, {}, 0, 2);
+        auto current = make_prompt(600, {}, 0, 3);
+        server_tokens target(llama_tokens(700, 1), false);
+
+        auto * useful_state = cache.alloc(useful, 512*1024, 0);
+        auto * useless_state = cache.alloc(useless, 512*1024, 0);
+        assert(useful_state != nullptr);
+        assert(useless_state != nullptr);
+
+        auto * current_state = cache.alloc(current, 1280*1024, 0);
+        assert(cache.finalize(current_state, &target, false));
+
+        assert(current_state != nullptr);
+        assert(cache_contains(cache, useful_state));
+        assert(!cache_contains_token(cache, 2));
+        assert(cache_contains(cache, current_state));
+    }
+
+    {
+        server_prompt_cache cache(4, 10000);
+        auto live = make_prompt(1000, {322});
+        auto cached = make_prompt(900, {322});
+
+        llama_tokens target_tokens(950, 1);
+        target_tokens[3] = 2;
+        server_tokens target(target_tokens, false);
+
+        cache.alloc(cached, 512*1024, 0);
+
+        // No exact recurrent state exists before the divergence at token 3.
+        assert(live.reusable_prefix_tokens(3, target.size(), true) == 0);
+        assert(cache.find_better(live, target, true, 1) == nullptr);
+    }
+
     return 0;
 }

tools/server/server-context.cpp

@@ -217,9 +217,13 @@ struct server_slot {
 
     server_prompt prompt;
 
-    bool prompt_save(server_prompt_cache & prompt_cache) const {
+    server_prompt_cache_state * prompt_save(
+            server_prompt_cache & prompt_cache,
+            const server_tokens * tokens_new = nullptr,
+            bool state_exact = false,
+            const server_prompt_cache_state * state_protected = nullptr) const {
         if (prompt.tokens.size() == 0) {
-            return false;
+            return nullptr;
         }
 
         const size_t cur_size_tgt =           llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
@@ -232,23 +236,46 @@ struct server_slot {
 
         auto * cur = prompt_cache.alloc(prompt, cur_size_tgt, cur_size_dft);
         if (cur == nullptr) {
-            return false;
+            return nullptr;
+        }
+
+        const size_t n_tgt = llama_state_seq_get_data_ext(
+                ctx_tgt,
+                cur->data.main.data(),
+                cur_size_tgt,
+                id,
+                LLAMA_STATE_SEQ_FLAGS_NONE);
+        if (n_tgt != cur_size_tgt) {
+            SLT_ERR(*this, "failed to save main prompt state: expected %zu bytes, wrote %zu\n", cur_size_tgt, n_tgt);
+            prompt_cache.discard(cur);
+            return nullptr;
         }
 
-        llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
         if (ctx_dft) {
-            llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE);
+            const size_t n_dft = llama_state_seq_get_data_ext(
+                    ctx_dft,
+                    cur->data.drft.data(),
+                    cur_size_dft,
+                    id,
+                    LLAMA_STATE_SEQ_FLAGS_NONE);
+            if (n_dft != cur_size_dft) {
+                SLT_ERR(*this, "failed to save draft prompt state: expected %zu bytes, wrote %zu\n", cur_size_dft, n_dft);
+                prompt_cache.discard(cur);
+                return nullptr;
+            }
         }
 
-        return true;
+        if (!prompt_cache.finalize(cur, tokens_new, state_exact, state_protected)) {
+            return nullptr;
+        }
+
+        return cur;
     }
 
     bool prompt_load(
             server_prompt_cache & prompt_cache,
-            const server_tokens & tokens,
-            bool state_exact,
-            int64_t min_reuse_gain) {
-        bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id, state_exact, min_reuse_gain);
+            server_prompt_cache_state * state) {
+        bool res = prompt_cache.load(prompt, state, ctx_tgt, ctx_dft, id);
         if (!res) {
             SLT_WRN(*this, "%s", "failed to load prompt from cache\n");
         }
@@ -1619,21 +1646,51 @@ struct server_context_impl {
         }
 
         if (ret) {
-            update_cache = update_cache && prompt_cache;
+            server_prompt_cache_state * state_selected = nullptr;
+            bool state_exact = false;
 
-            // cache prompts only for completion tasks
-            update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION;
+            const bool can_update_cache =
+                    prompt_cache &&
+                    task.type == SERVER_TASK_TYPE_COMPLETION;
+
+            if (can_update_cache) {
+                state_exact = ctx_tgt_state_exact();
+                const int64_t min_reuse_gain = std::max<int64_t>(1, params_base.checkpoint_min_step / 2);
+                state_selected = prompt_cache->find_better(
+                        ret->prompt,
+                        task.tokens,
+                        state_exact,
+                        min_reuse_gain);
+
+                const int lcp_live = ret->prompt.tokens.get_common_prefix(task.tokens);
+                const bool extends_live =
+                        !ret->prompt.tokens.empty() &&
+                        lcp_live == ret->prompt.n_tokens();
+
+                // A direct append can continue in the live slot without a host
+                // round-trip. A better saved branch still requires a state swap.
+                update_cache = state_selected != nullptr || (update_cache && !extends_live);
+            } else {
+                update_cache = false;
+            }
 
             if (update_cache) {
                 SRV_TRC("%s", "updating prompt cache\n");
 
                 const int64_t t_start = ggml_time_us();
 
-                ret->prompt_save(*prompt_cache);
-
-                const int64_t min_reuse_gain = std::max<int64_t>(1, params_base.checkpoint_min_step / 2);
-                if (!ret->prompt_load(*prompt_cache, task.tokens, ctx_tgt_state_exact(), min_reuse_gain)) {
-                    ret->prompt_clear();
+                auto * state_saved = ret->prompt_save(
+                        *prompt_cache,
+                        &task.tokens,
+                        state_exact,
+                        state_selected);
+
+                if (!ret->prompt_load(*prompt_cache, state_selected)) {
+                    // The active state was serialized before the switch. Restore
+                    // it if the selected cache entry could not be loaded.
+                    if (state_saved == nullptr || !ret->prompt_load(*prompt_cache, state_saved)) {
+                        ret->prompt_clear();
+                    }
                 }
 
                 prompt_cache->update();