diff --git a/tools/server/README.md b/tools/server/README.md index 6188a084bc0..a775207131c 100644 --- a/tools/server/README.md +++ b/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` diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 9377e3c1757..2e24e413a80 100644 --- a/tools/server/server-context.cpp +++ b/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)) { diff --git a/tools/server/tests/unit/test_kv_keep_only_active.py b/tools/server/tests/unit/test_kv_keep_only_active.py index 44c05fab0cb..78478187777 100644 --- a/tools/server/tests/unit/test_kv_keep_only_active.py +++ b/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 diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index 488b63701d1..94e486b8409 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/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):