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";