From 6b680320d8527b20c97e86b1d3e03c045d65f73d Mon Sep 17 00:00:00 2001 From: Jordi Posthumus Date: Fri, 4 Sep 2026 14:57:40 -0300 Subject: [PATCH] kvstore: save continued checkpoints after crossed frontiers --- README.md | 13 ++++++----- ds4_kvstore.c | 23 +++++++++++++++--- ds4_server.c | 64 +++++++++++++++++++++++++++++++++++++++++++++++---- 3 files changed, 87 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 503f00c630..a239a5dba2 100644 --- a/README.md +++ b/README.md @@ -1723,7 +1723,8 @@ builds for this model layout. The cache stores checkpoints at four moments: - `cold`: after a long first prompt reaches a stable prefix, before generation. -- `continued`: when prefill or generation reaches the next absolute aligned frontier. +- `continued`: at the first valid live checkpoint that reaches or passes the + next absolute interval frontier. - `evict`: before an unrelated request replaces the live in-memory session. - `shutdown`: when the server exits cleanly. @@ -1733,11 +1734,11 @@ future request appends text to the same prompt. The defaults are conservative: store prefixes of at least 512 tokens, cold-save prompts up to 30000 tokens, trim 32 tail tokens, and align to 2048-token chunks. The important knobs are: -Continued saves use the same alignment and are written only when the live graph -naturally reaches an absolute frontier. With the defaults this means roughly -every 10k tokens, independent of where the first cold checkpoint landed, so long -generations leave restart points behind without persisting the fragile final few -tokens. +Continued-save intervals are rounded up to the cold-boundary alignment. The +cache writes the first valid live checkpoint at or beyond each absolute +frontier. This matters when a restored cold checkpoint itself is unaligned: +fixed-size prefill chunks may step over every exact frontier, but they still +leave restart points roughly every 10k tokens instead of silently leaving none. - `--kv-cache-min-tokens` - `--kv-cache-cold-max-tokens` diff --git a/ds4_kvstore.c b/ds4_kvstore.c index 5b73de7b65..d495610f4a 100644 --- a/ds4_kvstore.c +++ b/ds4_kvstore.c @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -732,8 +733,12 @@ static int kv_cache_continued_step(const ds4_kvstore *kc) { int step = kc->opt.continued_interval_tokens; const int align = kc->opt.boundary_align_tokens; if (align > 0) { - step = ((step + align - 1) / align) * align; - if (step <= 0) step = align; + const int64_t rounded = + (((int64_t)step + align - 1) / align) * align; + /* live_tokens is an int, so a rounded frontier above INT_MAX can + * never be reached. Avoid signed overflow and disable that schedule. */ + if (rounded > INT_MAX) return 0; + step = (int)rounded; } return step; } @@ -742,8 +747,20 @@ int ds4_kvstore_continued_store_target(const ds4_kvstore *kc, int live_tokens) { const int step = kv_cache_continued_step(kc); if (step <= 0) return 0; if (live_tokens < kc->opt.min_tokens) return 0; - if (live_tokens % step != 0) return 0; if (live_tokens <= kc->continued_last_store_tokens) return 0; + + /* A restored cold/text checkpoint is not guaranteed to land on an exact + * multiple of step. Prefill normally advances in fixed-size chunks from + * that restored position, so requiring live_tokens % step == 0 can make + * the progression miss every absolute frontier forever. Save the first + * real live checkpoint at or beyond the next frontier instead. Session + * payloads already support arbitrary positions (evict/shutdown snapshots + * routinely use them), and returning live_tokens keeps the saved graph + * state and token vector at the same exact frontier. */ + int64_t last = kc->continued_last_store_tokens; + if (last < 0) last = 0; + const int64_t next = (last / step + 1) * (int64_t)step; + if ((int64_t)live_tokens < next) return 0; return live_tokens; } diff --git a/ds4_server.c b/ds4_server.c index 50efbe6bc1..860b66145a 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -10260,6 +10260,12 @@ static void kv_cache_slot_note_store(server_slot *slot, int tokens) { } } +static void kv_cache_slot_note_load(server_slot *slot, int tokens) { + /* A restore or rebuild may move a slot forward, backward, or to zero. + * Continued scheduling must restart from that exact frontier. */ + if (slot && tokens >= 0) slot->continued_last_store_tokens = tokens; +} + static int kv_cache_slot_suppress_continued(server *s, server_slot *slot, int tokens) { if (kv_cache_slot_continued_target(s, slot, tokens) != tokens) return -1; @@ -10342,6 +10348,7 @@ static int kv_cache_try_load_text(server *s, server_slot *slot, pthread_mutex_unlock(&s->kv_mu); pthread_mutex_unlock(&s->inference_mu); if (loaded > 0) { + kv_cache_slot_note_load(slot, loaded); if (loaded_path_out && lr.path) *loaded_path_out = xstrdup(lr.path); if (loaded_ext_flags_out) *loaded_ext_flags_out = lr.ext_flags; } @@ -11646,6 +11653,7 @@ static void canonicalize_tool_checkpoint(server *s, server_slot *slot, pthread_mutex_lock(&s->inference_mu); ds4_session_invalidate(slot->session); pthread_mutex_unlock(&s->inference_mu); + kv_cache_slot_note_load(slot, 0); } char sync_err[160] = {0}; @@ -18373,7 +18381,7 @@ static void test_kv_cache_chat_anchor_ignores_multiturn_tail(void) { ds4_tokens_free(&prompt); } -static void test_kv_cache_continued_uses_aligned_frontiers(void) { +static void test_kv_cache_continued_crosses_interval_frontiers(void) { kv_disk_cache kc = {0}; kc.enabled = true; kc.opt = kv_cache_default_options(); @@ -18395,6 +18403,53 @@ static void test_kv_cache_continued_uses_aligned_frontiers(void) { kc.continued_last_store_tokens = 20480; TEST_ASSERT(kv_cache_continued_store_target(&kc, 29999) == 0); TEST_ASSERT(kv_cache_continued_store_target(&kc, 30000) == 30000); + + /* Regression for a real Pi miss: loading a 9703-token cold anchor and + * advancing in 2048-token GLM prefill chunks never lands on an exact + * 16384-token multiple. The first live point past each absolute frontier + * must still become a continued checkpoint. */ + kc.opt.continued_interval_tokens = 16384; + kc.opt.boundary_align_tokens = 2048; + kc.continued_last_store_tokens = 9703; + TEST_ASSERT(kv_cache_continued_store_target(&kc, 15847) == 0); + TEST_ASSERT(kv_cache_continued_store_target(&kc, 17895) == 17895); + kc.continued_last_store_tokens = 17895; + TEST_ASSERT(kv_cache_continued_store_target(&kc, 32231) == 0); + TEST_ASSERT(kv_cache_continued_store_target(&kc, 34279) == 34279); + kc.continued_last_store_tokens = 34279; + TEST_ASSERT(kv_cache_continued_store_target(&kc, 34279) == 0); + + /* A large prefill step may cross several frontiers. Persist the exact + * live state once; the following interval is then based on that state. */ + kc.continued_last_store_tokens = 9703; + TEST_ASSERT(kv_cache_continued_store_target(&kc, 65536) == 65536); + kc.continued_last_store_tokens = 65536; + TEST_ASSERT(kv_cache_continued_store_target(&kc, 81919) == 0); + TEST_ASSERT(kv_cache_continued_store_target(&kc, 81920) == 81920); + + /* CLI values are int-sized. Rounding near INT_MAX must not overflow. */ + kc.opt.continued_interval_tokens = INT_MAX; + kc.opt.boundary_align_tokens = 2; + kc.continued_last_store_tokens = 0; + TEST_ASSERT(kv_cache_continued_store_target(&kc, INT_MAX) == 0); + kc.opt.boundary_align_tokens = 1; + TEST_ASSERT(kv_cache_continued_store_target(&kc, INT_MAX) == INT_MAX); +} + +static void test_kv_cache_disk_load_resets_slot_frontier(void) { + server_slot slot = {0}; + slot.continued_last_store_tokens = 0; + kv_cache_slot_note_load(&slot, 65536); + TEST_ASSERT(slot.continued_last_store_tokens == 65536); + + /* Restoring an older checkpoint must move the schedule backward too. */ + slot.continued_last_store_tokens = 100000; + kv_cache_slot_note_load(&slot, 65536); + TEST_ASSERT(slot.continued_last_store_tokens == 65536); + + /* A full rebuild starts scheduling again from token zero. */ + kv_cache_slot_note_load(&slot, 0); + TEST_ASSERT(slot.continued_last_store_tokens == 0); } static void test_kv_cache_cold_store_suppresses_duplicate_continued_boundary(void) { @@ -19029,7 +19084,7 @@ static void test_kv_cache_eviction_decayed_hits_tie_break_by_age(void) { rmdir(dir); } -static void test_kv_cache_eviction_keeps_aligned_continued_frontiers(void) { +static void test_kv_cache_eviction_keeps_continued_frontiers(void) { char tmpl[] = "/tmp/ds4-kv-live-prefix-test.XXXXXX"; char *dir = mkdtemp(tmpl); TEST_ASSERT(dir != NULL); @@ -19518,7 +19573,8 @@ static void ds4_server_unit_tests_run(void) { test_kv_cache_store_len_uses_configured_boundary(); test_kv_cache_chat_anchor_uses_last_user_before_assistant(); test_kv_cache_chat_anchor_ignores_multiturn_tail(); - test_kv_cache_continued_uses_aligned_frontiers(); + test_kv_cache_continued_crosses_interval_frontiers(); + test_kv_cache_disk_load_resets_slot_frontier(); test_kv_cache_cold_store_suppresses_duplicate_continued_boundary(); test_kv_cache_file_size_must_fit_budget(); test_sha1_bytes_hex_matches_known_vector(); @@ -19533,7 +19589,7 @@ static void ds4_server_unit_tests_run(void) { test_kv_cache_eviction_keeps_smaller_context_prefix(); test_kv_cache_eviction_score_decays_stale_hits(); test_kv_cache_eviction_decayed_hits_tie_break_by_age(); - test_kv_cache_eviction_keeps_aligned_continued_frontiers(); + test_kv_cache_eviction_keeps_continued_frontiers(); } #ifndef DS4_SERVER_TEST_NO_MAIN