server: discard prompt state for cache_prompt=false tasks

tools/server/README.md

@@ -594,7 +594,7 @@ These words will not be included in the completion, so make sure to add them to
 
 `id_slot`: Assign the completion task to an specific slot. If is -1 the task will be assigned to a Idle slot.  Default: `-1`
 
-`cache_prompt`: Re-use KV cache from a previous request if possible. This way the common prefix does not have to be re-processed, only the suffix that differs between the requests. Because (depending on the backend) the logits are **not** guaranteed to be bit-for-bit identical for different batch sizes (prompt processing vs. token generation) enabling this option can cause nondeterministic results. Default: `true`
+`cache_prompt`: Re-use KV cache from a previous request if possible. This way the common prefix does not have to be re-processed, only the suffix that differs between the requests. Because (depending on the backend) the logits are **not** guaranteed to be bit-for-bit identical for different batch sizes (prompt processing vs. token generation) enabling this option can cause nondeterministic results. If `false`, the state of this request is not kept for later requests: no context checkpoints are created for it, it is not saved to the prompt cache (`--cache-ram`), and its slot is cleared when it ends. Default: `true`
 
 `return_tokens`: Return the raw generated token ids in the `tokens` field. Otherwise `tokens` remains empty. Default: `false`
 

tools/server/server-context.cpp

@@ -665,7 +665,8 @@ struct server_slot {
 
             // do not keep context of the child slots - the parent's context is enough
             // nor of the stateless slots - no task ever reuses it
-            if (task->is_child() || task->is_stateless() || audio_stream || task->type == SERVER_TASK_TYPE_REALTIME) {
+            // nor of the tasks that disable the prompt cache
+            if (task->is_child() || task->is_stateless() || audio_stream || task->type == SERVER_TASK_TYPE_REALTIME || !task->params.cache_prompt) {
                 prompt_clear();
             }
 
@@ -1878,6 +1879,9 @@ struct server_context_impl {
         }
 
         if (ret) {
+            // a no-cache task overwrites the whole slot, save it first
+            update_cache = update_cache || !task.params.cache_prompt;
+
             update_cache = update_cache && prompt_cache;
 
             // cache prompts only for the tasks that reuse them
@@ -1890,7 +1894,8 @@ struct server_context_impl {
 
                 ret->prompt_save(*prompt_cache);
 
-                if (!ret->prompt_load(*prompt_cache, task.tokens)) {
+                // a no-cache task cannot use a loaded prompt
+                if (task.params.cache_prompt && !ret->prompt_load(*prompt_cache, task.tokens)) {
                     ret->prompt_clear();
                 }
 
@@ -4403,6 +4408,9 @@ struct server_context_impl {
                             ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS ||
                             n_swa > 0);
 
+                    // skip checkpoints for future prompt reuse when caching is disabled
+                    do_checkpoint = do_checkpoint && slot.task->params.cache_prompt;
+
                     // TODO: do the same for all models, then remove process_mtmd_chunk()
                     if (use_mixed_batch() && !slot.can_split() && input_tokens.has_mtmd) {
                         if (!add_prompt_mixed(slot)) {

tools/server/tests/unit/test_kv_keep_only_active.py

@@ -39,6 +39,15 @@ def create_server():
     "He met many creatures along the way including dragons and fairies "
     "and wizards who helped him on his noble quest to save the kingdom."
 )
+LONG_PROMPT_EXT = LONG_PROMPT + " The knight finally reached the castle gates."
+
+
+def completion(prompt, cache_prompt, **extra):
+    return server.make_request("POST", "/completion", data={
+        "prompt": prompt,
+        "cache_prompt": cache_prompt,
+        **extra,
+    })
 
 
 # idle slot cleared on launch should restore from cache-ram
@@ -89,6 +98,105 @@ def test_clear_and_restore():
     assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
 
 
+# the state of a cache_prompt = false task is not saved to cache-ram
+def test_nocache_prompt_not_saved():
+    global server
+    server.start()
+    log = LogReader(server.log_path)
+    assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log.drain()
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": LONG_PROMPT,
+        "id_slot": 0,
+        "cache_prompt": False,
+    })
+    assert res.status_code == 200
+    log.drain()
+
+    # Launching slot 1 would save idle slot 0, but its state came from a cache_prompt = false task
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "The quick brown fox",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
+
+
+# the state of a cache_prompt = false task is not saved when another task takes its slot
+def test_nocache_prompt_not_saved_on_slot_reuse():
+    global server
+    server.n_slots = 1
+    server.start()
+
+    assert completion(LONG_PROMPT, False).status_code == 200
+    assert completion("The quick brown fox", True).status_code == 200
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    assert res.body["timings"]["cache_n"] <= 1  # only BOS
+
+
+# a cache_prompt = false task overwrites its slot, so the slot is saved first
+def test_nocache_task_saves_replaced_slot():
+    global server
+    server.n_slots = 1
+    server.start()
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    n_prompt = res.body["timings"]["prompt_n"]
+
+    assert completion(LONG_PROMPT_EXT, False).status_code == 200
+    assert completion("The quick brown fox", True).status_code == 200
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    assert res.body["timings"]["cache_n"] > n_prompt // 2
+
+
+# a cache_prompt = false task does not take a matching prompt out of cache-ram
+def test_nocache_task_keeps_cached_prompt():
+    global server
+    server.n_slots = 1
+    server.start()
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    n_prompt = res.body["timings"]["prompt_n"]
+
+    assert completion("The quick brown fox", True).status_code == 200
+    log = LogReader(server.log_path)
+    log.drain()
+    assert completion(LONG_PROMPT_EXT, False).status_code == 200
+    assert "looking for better prompt" not in log.drain()
+    assert completion("Hello there, how are you today?", True).status_code == 200
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    assert res.body["timings"]["cache_n"] > n_prompt // 2
+
+
+# a cache_prompt = false task that fails before it runs does not lose the slot state
+def test_failed_nocache_task_keeps_slot_state():
+    global server
+    server.n_slots = 1
+    server.start()
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    n_prompt = res.body["timings"]["prompt_n"]
+
+    res = completion(LONG_PROMPT + " and then" * 400, False)
+    assert res.status_code != 200
+    assert "exceeds the available context size" in res.body["error"]["message"]
+    assert completion("The quick brown fox", True).status_code == 200
+
+    res = completion(LONG_PROMPT, True)
+    assert res.status_code == 200
+    assert res.body["timings"]["cache_n"] > n_prompt // 2
+
+
 def test_disabled_with_flag():
     global server
     server.no_cache_idle_slots = True

tools/server/tests/unit/test_slot_save.py

@@ -1,3 +1,6 @@
+import os
+import re
+import tempfile
 import pytest
 from utils import *
 import base64
@@ -676,6 +679,31 @@ def test_slot_restore_preserves_context_checkpoints(swa_server, filename):
     assert res.body["timings"]["prompt_n"] == n_live
 
 
+# a cache_prompt = false task creates no context checkpoints, a cache_prompt = true task still does
+def test_nocache_prompt_creates_no_checkpoints(swa_server):
+    server = swa_server
+    server.debug = True
+    fd, server.log_path = tempfile.mkstemp(suffix='.log')
+    os.close(fd)
+    server.start()
+
+    prompt = "The quick brown fox jumps over the lazy dog. " * 20
+    for id_slot, cache_prompt in [(1, False), (0, True)]:
+        res = server.make_request("POST", "/completion", data={
+            "prompt": prompt,
+            "id_slot": id_slot,
+            "cache_prompt": cache_prompt,
+        })
+        assert res.status_code == 200
+
+    # the log is written in order, so the slot 1 lines are complete once the slot 0 lines are there
+    with open(server.log_path) as f:
+        log = f.read()
+    n_created = lambda id_slot: len(re.findall(rf"id +{id_slot} \| task \d+ \| created context checkpoint", log))
+    assert n_created(0) > 0
+    assert n_created(1) == 0
+
+
 # checkpoint appendix: magic(4) version(4) count(4), then per checkpoint
 # n_tokens(8) pos_min(4) pos_max(4) and three blobs (target, draft, speculative), each size(8) + data
 def parse_ckpt_appendix(data):