server: extend live-prefix rewind to Flash on Metal, bounded by raw-cache budget

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
+}

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

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