From 943659b32987ca58ec877e71fe14d539e0af9a1b Mon Sep 17 00:00:00 2001 From: LEFBE Date: Wed, 5 Aug 2026 19:51:41 +0200 Subject: [PATCH] server: extend live-prefix rewind to Flash on Metal, bounded by raw-cache budget #668 added live-KV rewind for repeated exact-prefix prompts, but gated it to GLM only (ds4_engine_is_glm_dsa()) -- DeepSeek V4 Flash on Metal still does a full cold prefill on a retry/regenerate, even when the incoming prompt is an exact token-prefix of the live session, per #698. Root cause of the GLM-only gate (confirmed by reading the encode path, not guessed): Metal's raw SWA cache is a ring buffer, `raw_row = pos % raw_cap` (metal_graph_encode_token_raw_swa()). ds4_session_rewind() only truncates checkpoint.len; it does not shift or revalidate that ring. GLM's dense KV cache has no such ring -- ds4_session_glm_cap_dense_cache() keeps it consistent unconditionally -- but naively enabling rewind for Flash risks reusing a raw row that the discarded continuation has since overwritten, silently conditioning the resumed generation on the wrong tokens. Fix: new ds4_session_raw_rewind_budget() returns 0 for GLM (unaffected, gated separately as before) and, for Flash, (raw_cap - raw_window) -- the margin the graph already reserves so chunked prefill doesn't evict window-relevant rows mid-chunk. The rewind is only allowed when the discarded tail (old_pos - prompt_len) stays under that budget, which guarantees no raw row the rewound tail depends on has wrapped. ## Empirical validation (not just the theoretical argument above) M5 Max, Metal, production Flash q2-q4 quant: 1. Positive control: real ctx=8192 session, ~1000-token prompt, resend the exact same prompt after a short generation -> log shows "rewound Flash live prefix from N to M", continuation is byte-identical to the first generation. 2. Negative control: forced raw_cap=300 via DS4_METAL_GRAPH_RAW_CAP (so the budget is a easy-to-exceed 44 tokens), *with the budget check removed* to confirm the failure this guards against is real, not hypothetical -- generated 400 tokens (>> budget), resent the same prompt: continuation corrupted into a degenerate "140,140,140,..." repetition loop instead of the correct count, while the unmodified first generation was correct. 3. Same forced raw_cap=300, *with* the guard: the identical over-budget scenario correctly falls through to a full re-prefill (`reason=token-mismatch`) instead of rewinding, and the output stays correct -- confirming the guard degrades to today's existing (slower but correct) behavior exactly when it should, and a short within-budget retry under the same forced raw_cap still rewinds and matches. Co-Authored-By: Claude Sonnet 5 --- ds4.c | 21 +++++++++++++++++++++ ds4.h | 1 + ds4_server.c | 17 +++++++++++++---- 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/ds4.c b/ds4.c index aeca27a2da..59c617ec87 100644 --- a/ds4.c +++ b/ds4.c @@ -65546,3 +65546,24 @@ int ds4_session_ctx(ds4_session *s) { int ds4_session_prefill_cap(ds4_session *s) { return s ? (int)s->prefill_cap : 0; } + +/* Metal's raw SWA cache is a ring buffer of raw_cap rows indexed by + * pos % raw_cap (see metal_graph_encode_token_raw_swa()). ds4_session_rewind() + * only truncates checkpoint.len; it does not shift or revalidate that ring. + * A live-prefix rewind is only safe to reuse those rows verbatim while the + * discarded tokens have not yet wrapped a row the rewound tail still needs -- + * i.e. while (old_pos - new_pos) stays under (raw_cap - raw_window), the + * margin the graph already reserves so chunked prefill doesn't evict + * window-relevant rows mid-chunk. Returns 0 ("no budget, do not rewind") for + * GLM, whose dense KV cache has no such ring and is gated separately via + * ds4_engine_is_glm_dsa(), and for CPU-only builds with no raw cache. */ +uint32_t ds4_session_raw_rewind_budget(const ds4_session *s) { +#ifndef DS4_NO_GPU + if (!s || DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) return 0; + if (s->graph.raw_cap <= s->graph.raw_window) return 0; + return s->graph.raw_cap - s->graph.raw_window; +#else + (void)s; + return 0; +#endif +} diff --git a/ds4.h b/ds4.h index a8a0177c03..bc763d7aa3 100644 --- a/ds4.h +++ b/ds4.h @@ -415,6 +415,7 @@ void ds4_session_rewind(ds4_session *s, int pos); int ds4_session_pos(ds4_session *s); int ds4_session_ctx(ds4_session *s); int ds4_session_prefill_cap(ds4_session *s); +uint32_t ds4_session_raw_rewind_budget(const ds4_session *s); int ds4_engine_routed_quant_bits(ds4_engine *e); bool ds4_engine_has_output_head(ds4_engine *e); bool ds4_engine_has_mtp(ds4_engine *e); diff --git a/ds4_server.c b/ds4_server.c index 459288e19c..2564b8f1b5 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -11170,9 +11170,18 @@ static void generate_job(server *s, server_slot *slot, job *j) { "Anthropic continuation state is not available; retry by replaying the full messages history"); return; } else if (cached == 0) { + const bool is_glm = ds4_engine_is_glm_dsa(s->engine); + /* GLM's dense KV cache can always rewind: ds4_session_glm_cap_dense_cache() + * keeps it consistent. Flash's raw SWA cache is a ring buffer instead + * (see ds4_session_raw_rewind_budget()), so only rewind it while the + * discarded tail is still guaranteed not to have wrapped a row the + * rewound tail will need. */ + const uint32_t raw_budget = is_glm ? 0 : ds4_session_raw_rewind_budget(slot->session); + const bool can_rewind = is_glm || + (raw_budget > 0 && j->req.prompt.len < old_pos && + (uint32_t)(old_pos - j->req.prompt.len) < raw_budget); const int rewind_to = live_prefix_rewind_target( - ds4_engine_is_glm_dsa(s->engine), old_pos, - j->req.prompt.len, common); + can_rewind, old_pos, j->req.prompt.len, common); if (rewind_to >= 0) { pthread_mutex_lock(&s->inference_mu); ds4_session_rewind(slot->session, rewind_to); @@ -11181,8 +11190,8 @@ static void generate_job(server *s, server_slot *slot, job *j) { cache_source = "memory-rewind"; cache_diag.rewind_to = rewind_to; server_log(DS4_LOG_KVCACHE, - "ds4-server: rewound GLM live prefix from %d to %d; final prompt token will be reevaluated", - old_pos, rewind_to); + "ds4-server: rewound %s live prefix from %d to %d; final prompt token will be reevaluated", + is_glm ? "GLM" : "Flash", old_pos, rewind_to); } else { cached = common == old_pos && j->req.prompt.len >= old_pos ? common : 0; cache_source = cached > 0 ? "memory-token" : "none";