Skip to content

server: retry disk KV lookup with token text #943

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
+205 −10
Open

server: retry disk KV lookup with token text #943

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
215 changes: 205 additions & 10 deletions ds4_server.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10349,15 +10349,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,
Expand Down Expand Up @@ -11639,9 +11702,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);
Expand Down Expand Up @@ -18494,6 +18557,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);
Expand All @@ -18502,8 +18684,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;
Expand All @@ -18517,20 +18701,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);
}

Expand Down Expand Up @@ -19522,6 +19716,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();
Expand Down