From d9a77d6602e901ee6f6e392940344cc32ef571f9 Mon Sep 17 00:00:00 2001 From: Jordi Posthumus Date: Fri, 4 Sep 2026 14:58:05 -0300 Subject: [PATCH] server: retry disk KV lookup with token text --- ds4_server.c | 215 ++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 205 insertions(+), 10 deletions(-) diff --git a/ds4_server.c b/ds4_server.c index adbd8f080..5ca46e021 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -10381,15 +10381,78 @@ static int kv_cache_try_load_text(server *s, server_slot *slot, return loaded; } +typedef int (*kv_cache_text_load_fn)(server *, server_slot *, const char *, + ds4_tokens *, char **, uint8_t *, bool); +typedef char *(*kv_cache_tokens_render_fn)(ds4_engine *, const ds4_tokens *, + size_t *); + +static int kv_cache_try_load_tokenized_text_with( + server *s, server_slot *slot, + const char *prompt_text, const ds4_tokens *prompt_tokens, + ds4_tokens *effective_prompt, char **loaded_path_out, + uint8_t *loaded_ext_flags_out, bool responses_protocol, + kv_cache_text_load_fn load_text, + kv_cache_tokens_render_fn render_tokens) { + if (!s || !slot) return 0; + int loaded = load_text(s, slot, prompt_text, + effective_prompt, loaded_path_out, + loaded_ext_flags_out, responses_protocol); + /* Preserve the disabled-cache path: the ordinary loader is already a + * cheap no-op, and there is no reason to decode a context-sized token + * vector when no second lookup can succeed. */ + if (loaded > 0 || !s->kv.enabled || !prompt_tokens) { + return loaded; + } + + /* Checkpoints are keyed by decoded token text. The canonical decoded + * bytes of the request tokens can differ from the originally rendered + * prompt (for example when the tokenizer canonicalizes a special-token + * spelling). Keep the ordinary raw-text lookup fast, then retry a miss + * with the same representation used by the writer. */ + size_t canonical_len = 0; + char *canonical = render_tokens(s->engine, prompt_tokens, &canonical_len); + const size_t prompt_len = prompt_text ? strlen(prompt_text) : 0; + if (canonical_len == prompt_len && + (canonical_len == 0 || !memcmp(canonical, prompt_text, canonical_len))) + { + free(canonical); + return 0; + } + + loaded = load_text(s, slot, canonical, + effective_prompt, loaded_path_out, + loaded_ext_flags_out, responses_protocol); + if (loaded > 0) { + server_log(DS4_LOG_KVCACHE, + "ds4-server: kv cache recovered canonical rendered prompt cached=%d", + loaded); + } + free(canonical); + return loaded; +} + +static int kv_cache_try_load_tokenized_text(server *s, server_slot *slot, + const char *prompt_text, + const ds4_tokens *prompt_tokens, + ds4_tokens *effective_prompt, + char **loaded_path_out, + uint8_t *loaded_ext_flags_out, + bool responses_protocol) { + return kv_cache_try_load_tokenized_text_with( + s, slot, prompt_text, prompt_tokens, effective_prompt, + loaded_path_out, loaded_ext_flags_out, responses_protocol, + kv_cache_try_load_text, render_tokens_text); +} + static int kv_cache_try_load(server *s, server_slot *slot, const request *req, ds4_tokens *effective_prompt, char **loaded_path_out, uint8_t *loaded_ext_flags_out) { - return kv_cache_try_load_text(s, slot, req ? req->prompt_text : NULL, - effective_prompt, - loaded_path_out, - loaded_ext_flags_out, - req && req->api == API_RESPONSES); + return kv_cache_try_load_tokenized_text( + s, slot, req ? req->prompt_text : NULL, + req ? &req->prompt : NULL, effective_prompt, + loaded_path_out, loaded_ext_flags_out, + req && req->api == API_RESPONSES); } static int live_text_prefix_prompt(server *s, server_slot *slot, @@ -11704,9 +11767,9 @@ static void canonicalize_tool_checkpoint(server *s, server_slot *slot, * a very long conversation from token zero. */ char *path = NULL; ds4_tokens effective = {0}; - int loaded = kv_cache_try_load_text(s, slot, - rendered.ptr ? rendered.ptr : "", - &effective, &path, NULL, false); + int loaded = kv_cache_try_load_tokenized_text( + s, slot, rendered.ptr ? rendered.ptr : "", &canonical, + &effective, &path, NULL, false); if (loaded == 0) { pthread_mutex_lock(&s->inference_mu); ds4_session_invalidate(slot->session); @@ -18833,6 +18896,125 @@ static void test_kv_text_stub_file(const char *dir, const char *text, test_kv_text_stub_file_model(dir, text, 0, reason, tokens, payload_bytes); } +typedef struct { + int load_calls; + int render_calls; + int load_result[2]; + const char *canonical; + char seen[2][96]; +} test_token_text_retry_probe; + +static test_token_text_retry_probe *test_token_text_probe; + +static int test_token_text_load(server *s, server_slot *slot, + const char *text, ds4_tokens *effective, + char **path, uint8_t *flags, + bool responses_protocol) { + (void)s; + (void)slot; + (void)responses_protocol; + test_token_text_retry_probe *p = test_token_text_probe; + TEST_ASSERT(p != NULL); + if (!p) return 0; + const int call = p->load_calls++; + TEST_ASSERT(call < 2); + if (call < 2) { + snprintf(p->seen[call], sizeof(p->seen[call]), "%s", text ? text : ""); + const int loaded = p->load_result[call]; + if (loaded > 0) { + if (effective) ds4_tokens_push(effective, 4242); + if (path) *path = xstrdup("canonical-hit.kv"); + if (flags) *flags = 0xa5; + } + return loaded; + } + return 0; +} + +static char *test_token_text_render(ds4_engine *engine, + const ds4_tokens *tokens, + size_t *out_len) { + (void)engine; + (void)tokens; + test_token_text_retry_probe *p = test_token_text_probe; + TEST_ASSERT(p != NULL); + if (!p) return xstrdup(""); + p->render_calls++; + const char *text = p->canonical ? p->canonical : ""; + if (out_len) *out_len = strlen(text); + return xstrdup(text); +} + +static int test_token_text_retry(test_token_text_retry_probe *probe, + bool cache_enabled, + const char *raw_text, + const char *canonical_text) { + memset(probe, 0, sizeof(*probe)); + probe->canonical = canonical_text; + server s = {0}; + s.kv.enabled = cache_enabled; + server_slot slot = {0}; + ds4_tokens prompt = {0}; + test_token_text_probe = probe; + int loaded = kv_cache_try_load_tokenized_text_with( + &s, &slot, raw_text, &prompt, NULL, NULL, NULL, false, + test_token_text_load, test_token_text_render); + test_token_text_probe = NULL; + return loaded; +} + +static void test_token_text_cache_retry_control_flow(void) { + test_token_text_retry_probe p; + + memset(&p, 0, sizeof(p)); + p.load_result[0] = 512; + p.canonical = "canonical"; + server s = {0}; + s.kv.enabled = true; + server_slot slot = {0}; + ds4_tokens prompt = {0}; + test_token_text_probe = &p; + TEST_ASSERT(kv_cache_try_load_tokenized_text_with( + &s, &slot, "raw", &prompt, NULL, NULL, NULL, false, + test_token_text_load, test_token_text_render) == 512); + test_token_text_probe = NULL; + TEST_ASSERT(p.load_calls == 1); + TEST_ASSERT(p.render_calls == 0); + TEST_ASSERT(!strcmp(p.seen[0], "raw")); + + TEST_ASSERT(test_token_text_retry(&p, true, "same", "same") == 0); + TEST_ASSERT(p.load_calls == 1); + TEST_ASSERT(p.render_calls == 1); + TEST_ASSERT(!strcmp(p.seen[0], "same")); + + memset(&p, 0, sizeof(p)); + p.load_result[1] = 768; + p.canonical = "canonical"; + s.kv.enabled = true; + ds4_tokens effective = {0}; + char *path = NULL; + uint8_t flags = 0; + test_token_text_probe = &p; + TEST_ASSERT(kv_cache_try_load_tokenized_text_with( + &s, &slot, "raw", &prompt, &effective, &path, &flags, false, + test_token_text_load, test_token_text_render) == 768); + test_token_text_probe = NULL; + TEST_ASSERT(p.load_calls == 2); + TEST_ASSERT(p.render_calls == 1); + TEST_ASSERT(!strcmp(p.seen[0], "raw")); + TEST_ASSERT(!strcmp(p.seen[1], "canonical")); + TEST_ASSERT(effective.len == 1 && effective.v[0] == 4242); + TEST_ASSERT(path != NULL && !strcmp(path, "canonical-hit.kv")); + TEST_ASSERT(flags == 0xa5); + ds4_tokens_free(&effective); + free(path); + + TEST_ASSERT(test_token_text_retry(&p, false, "raw", "canonical") == 0); + TEST_ASSERT(p.load_calls == 1); + TEST_ASSERT(p.render_calls == 0); + TEST_ASSERT(!strcmp(p.seen[0], "raw")); +} + static void test_kv_cache_lookup_uses_longest_text_prefix(void) { char tmpl[] = "/tmp/ds4-kv-text-prefix-test.XXXXXX"; char *dir = mkdtemp(tmpl); @@ -18841,8 +19023,10 @@ static void test_kv_cache_lookup_uses_longest_text_prefix(void) { const char *short_text = "transcript prefix"; const char *long_text = "transcript prefix with sampled token bytes"; + const char *canonical_text = "prefix <｜Assistant｜> canonical bytes"; test_kv_text_stub_file(dir, short_text, KV_REASON_COLD, 512, 0); test_kv_text_stub_file(dir, long_text, KV_REASON_COLD, 768, 0); + test_kv_text_stub_file(dir, canonical_text, KV_REASON_CONTINUED, 1024, 0); kv_disk_cache kc = {0}; kc.enabled = true; @@ -18856,20 +19040,30 @@ static void test_kv_cache_lookup_uses_longest_text_prefix(void) { TEST_ASSERT(idx >= 0 && kc.entry[idx].tokens == 768); TEST_ASSERT(idx >= 0 && kc.entry[idx].text_bytes == strlen(long_text)); TEST_ASSERT(kv_cache_find_text_prefix(&kc, "transcript prefiX", 2, 32768) < 0); + TEST_ASSERT(kv_cache_find_text_prefix( + &kc, "prefix <|assistant|> canonical bytes suffix", 2, 32768) < 0); + idx = kv_cache_find_text_prefix( + &kc, "prefix <｜Assistant｜> canonical bytes suffix", 2, 32768); + TEST_ASSERT(idx >= 0 && kc.entry[idx].tokens == 1024); kv_cache_close(&kc); - char short_sha[41], long_sha[41]; + char short_sha[41], long_sha[41], canonical_sha[41]; sha1_bytes_hex(short_text, strlen(short_text), short_sha); sha1_bytes_hex(long_text, strlen(long_text), long_sha); - char short_name[44], long_name[44]; + sha1_bytes_hex(canonical_text, strlen(canonical_text), canonical_sha); + char short_name[44], long_name[44], canonical_name[44]; snprintf(short_name, sizeof(short_name), "%.40s.kv", short_sha); snprintf(long_name, sizeof(long_name), "%.40s.kv", long_sha); + snprintf(canonical_name, sizeof(canonical_name), "%.40s.kv", canonical_sha); char *short_path = path_join(dir, short_name); char *long_path = path_join(dir, long_name); + char *canonical_path = path_join(dir, canonical_name); unlink(short_path); unlink(long_path); + unlink(canonical_path); free(short_path); free(long_path); + free(canonical_path); rmdir(dir); } @@ -19867,6 +20061,7 @@ static void ds4_server_unit_tests_run(void) { test_kv_cache_cold_store_suppresses_duplicate_continued_boundary(); test_kv_cache_file_size_must_fit_budget(); test_sha1_bytes_hex_matches_known_vector(); + test_token_text_cache_retry_control_flow(); test_kv_cache_lookup_uses_longest_text_prefix(); test_kv_cache_lookup_rejects_wrong_model(); test_kv_cache_lookup_rejects_stale_payload_abi();