server: reuse exact multimodal live KV safely

This commit is contained in:
Jordi Posthumus committed 2026-09-01 07:46:20 -03:00
1 parent 3da2da96b2
commit 98b91351ad
4 files changed
+120 -15

No files matched your search

+14
View File
@@ -65909,6 +65909,7 @@ static bool ds4_session_vision_prefix_matches(
const ds4_session *s,
const ds4_vision_span *images,
size_t image_count) {
if (!s || (image_count != 0 && !images)) return false;
if (!s->checkpoint_valid) return true;
if (s->checkpoint_image_count > image_count) return false;
for (size_t i = 0; i < s->checkpoint_image_count; i++) {
@@ -65926,6 +65927,19 @@ static bool ds4_session_vision_prefix_matches(
return true;
}
bool ds4_session_vision_state_matches(
const ds4_session *s,
const ds4_vision_span *images,
size_t image_count) {
return s && s->checkpoint_valid &&
s->checkpoint_image_count == image_count &&
ds4_session_vision_prefix_matches(s, images, image_count);
}
bool ds4_session_has_vision_state(const ds4_session *s) {
return s && (s->checkpoint_image_count != 0 || s->sync_image_count != 0);
}
static bool ds4_session_vision_range_overlaps(
const ds4_session *s,
uint32_t token_start,
+8
View File
@@ -428,6 +428,14 @@ int ds4_session_sync_multimodal(ds4_session *s,
size_t image_count,
char *err,
size_t errlen);
/* Return true only when every image that conditioned the live checkpoint has
* the same token span and embedding fingerprint in the supplied prompt. */
bool ds4_session_vision_state_matches(const ds4_session *s,
const ds4_vision_span *images,
size_t image_count);
/* True while a session contains, or is actively syncing, image-conditioned
* state. Such state must not be written to the text-keyed disk KV cache. */
bool ds4_session_has_vision_state(const ds4_session *s);
bool ds4_session_rewrite_requires_rebuild(int live_len, int canonical_len, int common);
ds4_session_rewrite_result ds4_session_rewrite_from_common(
ds4_session *s, const ds4_tokens *prompt, int common,
+81 -15
View File
@@ -10163,6 +10163,15 @@ static bool kv_cache_store_live_prefix_text(server *s, server_slot *slot,
char err[160] = {0};
ds4_kvstore_trailer_hooks hooks = kv_cache_tool_map_hooks(s, NULL);
pthread_mutex_lock(&s->inference_mu);
/* The payload contains image-conditioned KV rows, but the disk key and
* trailer do not contain image fingerprints. Never let generic image
* placeholder tokens become a cache hit for a different image.
* sync_image_count covers progress-callback writes during prefill;
* checkpoint_image_count covers completed sessions. */
if (ds4_session_has_vision_state(slot->session)) {
pthread_mutex_unlock(&s->inference_mu);
return false;
}
pthread_mutex_lock(&s->kv_mu);
bool ok = ds4_kvstore_store_live_prefix_text(&s->kv, s->engine,
slot->session,
@@ -11182,6 +11191,25 @@ static int server_session_sync(server *s, server_slot *slot,
DS4_SESSION_SYNC_INTERRUPTED : 0;
}
static int server_multimodal_resume_frontier(int live, int common,
int prompt_len,
bool image_state_matches) {
return common == live && prompt_len >= live && image_state_matches
? live : 0;
}
static int server_multimodal_resume_pos(ds4_session *session,
const ds4_tokens *prompt,
const ds4_vision_span *images,
size_t image_count) {
if (!session || !prompt) return 0;
const int live = ds4_session_pos(session);
const int common = ds4_session_common_prefix(session, prompt);
return server_multimodal_resume_frontier(
live, common, prompt->len,
ds4_session_vision_state_matches(session, images, image_count));
}
static int server_session_sync_multimodal(server *s, server_slot *slot,
const ds4_tokens *prompt,
const ds4_vision_span *images,
@@ -11199,7 +11227,13 @@ static int server_session_sync_multimodal(server *s, server_slot *slot,
return rc;
}
int done = 0;
/* Start at the exact live frontier when both tokens and image identities
* match. Starting at zero made the first scheduling slice truncate a
* perfectly reusable long checkpoint, forcing a complete refill. */
pthread_mutex_lock(&s->inference_mu);
int done = server_multimodal_resume_pos(slot->session, prompt,
images, image_count);
pthread_mutex_unlock(&s->inference_mu);
bool called = false;
while (!g_stop_requested && !slot_job_cancelled(slot) &&
(!called || done < prompt->len)) {
@@ -11560,6 +11594,19 @@ static void canonicalize_tool_checkpoint(server *s, server_slot *slot,
}
free(live_text);
if (j->req.image_count != 0) {
/* Rebuilding through the text-only disk cache would either lose the
* vision embeddings or restore rows for an unverified image. Keep the
* correctly conditioned sampled frontier instead. */
server_log(DS4_LOG_WARNING,
"ds4-server: multimodal tool checkpoint canonicalization skipped ctx=%s common=%d live=%d canonical=%d reason=preserve-image-state",
ctx, common, live_len, canonical.len);
trace_event(s, trace_id,
"multimodal tool checkpoint canonicalization skipped: common=%d live=%d canonical=%d",
common, live_len, canonical.len);
goto done;
}
if (common < j->req.prompt.len) {
trace_event(s, trace_id,
"tool checkpoint canonicalization skipped: common=%d prompt=%d live=%d canonical=%d",
@@ -11915,14 +11962,13 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
char err[160];
err[0] = '\0';
const bool multimodal = j->req.image_count != 0;
if (multimodal) {
pthread_mutex_lock(&s->inference_mu);
ds4_session_invalidate(slot->session);
pthread_mutex_unlock(&s->inference_mu);
request_live_state_clear(s, slot);
}
pthread_mutex_lock(&s->inference_mu);
const int old_pos = ds4_session_pos(slot->session);
const int common = ds4_session_common_prefix(slot->session, &j->req.prompt);
const bool live_vision_match =
ds4_session_vision_state_matches(slot->session,
j->req.images, j->req.image_count);
pthread_mutex_unlock(&s->inference_mu);
trace_cache_diag cache_diag = {0};
trace_cache_capture(&cache_diag, ds4_session_tokens(slot->session),
&j->req.prompt, old_pos, common);
@@ -11941,8 +11987,9 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
* exact token-prefix match. Exact token/text/disk matching remains the
* fallback when the live state is absent or no longer describes the
* request. */
int cached = responses_live_visible_prefix_prompt(s, slot, &j->req, old_pos,
&effective_prompt);
int cached = live_vision_match ?
responses_live_visible_prefix_prompt(s, slot, &j->req, old_pos,
&effective_prompt) : 0;
const char *cache_source = cached > 0 ? "responses-visible" : "none";
if (cached > 0) {
responses_live_match = "visible-prefix";
@@ -11953,7 +12000,7 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
responses_live_match_ids = j->req.responses_live_call_ids.len;
}
}
if (cached == 0) {
if (cached == 0 && live_vision_match) {
cached = responses_live_continuation_prompt(s, slot, &j->req, old_pos,
&effective_prompt,
&responses_live_match_ids);
@@ -11963,7 +12010,7 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
if (cached > 0) {
responses_live_continuation = true;
prompt_for_sync = &effective_prompt;
} else {
} else if (live_vision_match) {
cached = anthropic_live_continuation_prompt(s, slot, &j->req, old_pos,
&effective_prompt,
&anthropic_live_match_ids);
@@ -11991,7 +12038,7 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
http_error(j->fd, s->enable_cors, 409,
"Anthropic continuation state is not available; retry by replaying the full messages history");
return;
} else if (cached == 0) {
} else if (cached == 0 && live_vision_match) {
const int rewind_to = live_prefix_rewind_target(
ds4_engine_is_glm_dsa(s->engine), old_pos,
j->req.prompt.len, common);
@@ -12010,7 +12057,7 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
cache_source = cached > 0 ? "memory-token" : "none";
}
}
if (cached == 0) {
if (cached == 0 && live_vision_match) {
int thinking_cached =
thinking_live_visible_prefix_prompt(s, slot, &j->req, old_pos,
&effective_prompt);
@@ -12024,7 +12071,7 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
int disk_cached = 0;
char *disk_cache_path = NULL;
uint8_t disk_cache_ext_flags = 0;
if (cached == 0) {
if (cached == 0 && live_vision_match) {
int text_cached = live_text_prefix_prompt(s, slot, &j->req,
&effective_prompt);
if (text_cached > 0) {
@@ -12035,11 +12082,17 @@ static void generate_job_inner(server *s, server_slot *slot, job *j) {
}
if (cached == 0 && old_pos > 0) {
server_log(DS4_LOG_WARNING,
"ds4-server: live kv cache miss%s live=%d prompt=%d common=%d reason=%s",
"ds4-server: live kv cache miss%s live=%d prompt=%d common=%d vision=%s reason=%s",
responses_protocol ? " RESPPROTO" : "",
old_pos, j->req.prompt.len, common,
live_vision_match ? "match" : "mismatch",
trace_cache_miss_reason(&cache_diag));
}
if (multimodal && cached > 0) {
server_log(DS4_LOG_KVCACHE,
"ds4-server: multimodal live kv hit images=%zu cached=%d prompt=%d identity=fingerprint-match",
j->req.image_count, cached, prompt_for_sync->len);
}
if (cached == 0) slot->continued_last_store_tokens = 0;
if (!multimodal && s->kv.enabled && cached == 0 &&
old_pos >= s->kv.opt.min_tokens) {
@@ -13157,6 +13210,11 @@ static int job_slot_score(server *s, server_slot *slot, const job *j,
if (!s || !slot || !j || slot->busy || slot->assigned) return INT_MIN;
if (required_slot >= 0 && slot->id != required_slot) return INT_MIN;
if (required_slot == slot->id) return INT_MAX;
if (ds4_session_pos(slot->session) > 0 &&
!ds4_session_vision_state_matches(slot->session,
j->req.images, j->req.image_count)) {
return -1;
}
int common = ds4_session_common_prefix(slot->session, &j->req.prompt);
return common;
}
@@ -14450,6 +14508,13 @@ static void test_mixed_prefill_quantum_option(void) {
TEST_ASSERT(server_prefill_quantum_for(&s, true) == 128);
}
static void test_multimodal_prefill_resume_frontier(void) {
TEST_ASSERT(server_multimodal_resume_frontier(160, 160, 170, true) == 160);
TEST_ASSERT(server_multimodal_resume_frontier(160, 159, 170, true) == 0);
TEST_ASSERT(server_multimodal_resume_frontier(160, 160, 159, true) == 0);
TEST_ASSERT(server_multimodal_resume_frontier(160, 160, 170, false) == 0);
}
static void test_batched_live_continuation_slot_binding(void) {
server s = {0};
server_slot slots[3] = {0};
@@ -19329,6 +19394,7 @@ static void test_responses_inline_image_content(void) {
static void ds4_server_unit_tests_run(void) {
test_batched_prefill_round_robin();
test_mixed_prefill_quantum_option();
test_multimodal_prefill_resume_frontier();
test_batched_live_continuation_slot_binding();
test_request_defaults_use_min_p_filtering();
test_chat_ignore_eos_contract();
+17
View File
@@ -78,6 +78,13 @@ int main(int argc, char **argv) {
}
if (ds4_session_sync_multimodal(session, &prompt, &span, 1,
error, sizeof(error)) != 0) goto done;
if (!ds4_session_has_vision_state(session) ||
!ds4_session_vision_state_matches(session, &span, 1) ||
ds4_session_vision_state_matches(session, NULL, 0)) {
snprintf(error, sizeof(error),
"live session did not retain exact image identity");
goto done;
}
progress_counter progress = {0};
ds4_session_set_progress(session, count_progress, &progress);
@@ -130,8 +137,18 @@ int main(int argc, char **argv) {
memset(span.embedding.data, 0,
image_embedding_elems * sizeof(*span.embedding.data));
span.embedding.fingerprint[0] ^= 1u;
if (ds4_session_vision_state_matches(session, &span, 1)) {
snprintf(error, sizeof(error),
"changed image fingerprint matched live session");
goto done;
}
if (ds4_session_sync_multimodal(session, &prompt, &span, 1,
error, sizeof(error)) != 0) goto done;
if (!ds4_session_vision_state_matches(session, &span, 1)) {
snprintf(error, sizeof(error),
"rebuilt session did not retain changed image identity");
goto done;
}
if (progress.calls == 0 || progress.current != prompt.len ||
progress.total != prompt.len) {
snprintf(error, sizeof(error),