Skip to content

kvstore: save continued checkpoints after crossed interval frontiers #903

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JordiPosthumus wants to merge 1 commit into
base: main
Choose a base branch
from
+87 −13
Open

kvstore: save continued checkpoints after crossed interval frontiers #903

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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`
Expand Down
23 changes: 20 additions & 3 deletions ds4_kvstore.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#include <ctype.h>
#include <dirent.h>
#include <errno.h>
#include <limits.h>
#include <math.h>
#include <stdarg.h>
#include <stdio.h>
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
}

Expand Down
64 changes: 60 additions & 4 deletions ds4_server.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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};
Expand Down Expand Up @@ -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();
Expand All @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand All @@ -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
Expand Down