Commit f174553e1 for llama.cpp
commit f174553e1621c6a22a64e7bfcbb58d3a5c49e6fb
Author: Tarek Dakhran <tarek@liquid.ai>
Date: Mon Oct 19 15:58:14 2026 +0200
server : fix swa checkpoint invalidation logic (#30666)
* server : fix swa checkpoint invalidation logic
Assisted-by: Claude
* add testcase that fails with > n_past, but works with > n_prefix
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index a18b7542c..9c3c54c3a 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -3781,6 +3781,9 @@ private:
llama_pos pos_next = slot.prompt.tokens.pos_next(n_past);
+ // number of cached tokens shared with the new prompt, before any checkpoint is restored
+ const int n_prefix = n_past;
+
// ref: https://github.com/ggml-org/llama.cpp/pull/24110
const bool has_new_tokens = (n_past < slot.task->n_tokens());
@@ -3891,10 +3894,10 @@ private:
}
{
- // erase any checkpoints with pos_max > pos_next
+ // erase any checkpoints with pos_max > pos_next, or that cover tokens not shared with the new prompt
for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end();) {
const auto & cur = *it;
- if (cur.pos_max > pos_next) {
+ if (cur.pos_max > pos_next || cur.n_tokens > n_prefix) {
SLT_TRC(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024);
it = slot.prompt.checkpoints.erase(it);
} else {
diff --git a/tools/server/tests/unit/test_completion.py b/tools/server/tests/unit/test_completion.py
index b62abc661..572f5368c 100644
--- a/tools/server/tests/unit/test_completion.py
+++ b/tools/server/tests/unit/test_completion.py
@@ -689,3 +689,81 @@ def test_completion_prompt_cache():
assert "prompt_n" in timings and timings["prompt_n"] + timings["cache_n"] == n_prompt
assert "predicted_n" in timings and timings["predicted_n"] == n_predict
assert "tokens" in res.body and isinstance(res.body["tokens"], list)
+
+
+# a slot that switched to an unrelated prompt must not restore a context checkpoint of the previous one
+def test_completion_checkpoint_after_unrelated_prompt():
+ global server
+ server = ServerPreset.tinygemma3()
+ server.n_slots = 1
+ server.cache_ram = 0
+ server.temperature = 0.0
+ server.start()
+
+ def tokenize(content):
+ res = server.make_request("POST", "/tokenize", data={"content": content})
+ assert res.status_code == 200
+ return res.body["tokens"]
+
+ def top_logprobs(prompt, cache_prompt):
+ res = server.make_request("POST", "/completion", data={
+ "prompt": prompt,
+ "n_predict": 1,
+ "n_probs": 5,
+ "cache_prompt": cache_prompt,
+ })
+ assert res.status_code == 200
+ return [(p["id"], p["logprob"]) for p in res.body["completion_probabilities"][0]["top_logprobs"]]
+
+ tokens = tokenize("Purple elephants quietly compose sonnets about tax law")
+ expected = top_logprobs(tokens[:2], False)
+
+ # 5 tokens: the checkpoint 4 tokens before the end of the prompt holds only its first token
+ top_logprobs(tokenize("The quick brown fox jumps over the lazy dog")[:5], True)
+ top_logprobs(tokens[:1], True)
+ actual = top_logprobs(tokens[:2], True)
+
+ assert [tid for tid, _ in actual] == [tid for tid, _ in expected]
+ for (_, lp), (_, lp_expected) in zip(actual, expected):
+ assert lp == pytest.approx(lp_expected, abs=1e-3)
+
+
+# a restored context checkpoint must be kept for the next request that shares the same prefix
+def test_completion_checkpoint_kept_after_restore():
+ global server
+ server = ServerPreset.tinygemma3()
+ server.n_slots = 1
+ server.n_batch = 16
+ server.n_ubatch = 16
+ server.cache_ram = 0
+ server.temperature = 0.0
+ server.start()
+
+ def tokenize(content):
+ res = server.make_request("POST", "/tokenize", data={"content": content})
+ assert res.status_code == 200
+ return res.body["tokens"]
+
+ def cache_n(prompt):
+ res = server.make_request("POST", "/completion", data={
+ "prompt": prompt,
+ "n_predict": 1,
+ "cache_prompt": True,
+ })
+ assert res.status_code == 200
+ return res.body["timings"]["cache_n"]
+
+ prefix = tokenize("Purple elephants quietly compose sonnets about tax law in the moonlight")
+ suffixes = [tokenize(s) for s in [
+ "The quick brown fox jumps over the lazy dog while the farmer counts his sheep and the sun slowly sets behind the distant hills",
+ "Seven tiny robots repaired the broken lighthouse lens before the storm arrived, then danced on the rocks until the morning tide came in",
+ "Green apples rolled down the hill toward the river, where a curious otter collected them one by one and stacked them on a flat stone",
+ ]]
+
+ # 4 tokens: the checkpoint 4 tokens before the end of the prompt holds exactly the prefix
+ cache_n(prefix + suffixes[0][:4])
+
+ # suffixes longer than n_ubatch + 4: no new checkpoint within the prefix, only the restored one can be reused
+ for suffix in suffixes[1:]:
+ assert len(suffix) > server.n_ubatch + 4
+ assert cache_n(prefix + suffix) == len(prefix) - 1