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

[Flor1an-B]

Fixes #698.

Background

#668 added live-KV rewind for repeated exact-prefix prompts (retry/regenerate workflows), but gated it to GLM only via ds4_engine_is_glm_dsa(). #698 reports that DeepSeek V4 Flash on Metal still does a full cold prefill in the same scenario, even when the incoming prompt is an exact token-prefix of the live session (e.g. a 175K-token prompt re-prefilled from scratch after a client-side stream abort).

Why the gate exists (this isn't a small oversight)

Read metal_graph_encode_token_raw_swa() rather than guessing: Metal's raw SWA cache is a genuine ring buffer, raw_row = pos % raw_cap. 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 on rewind), but naively enabling rewind for Flash risks reusing a raw row that the discarded continuation has already overwritten — silently conditioning the resumed generation on stale/wrong tokens. That's a correctness bug, not just a missed cache hit, so I didn't want to just flip the gate.

Fix

New ds4_session_raw_rewind_budget():

• Returns 0 for GLM (unaffected, still gated the existing way).
• For Flash, returns raw_cap - raw_window — the margin the graph already reserves specifically so chunked prefill doesn't evict window-relevant rows mid-chunk. That existing margin turns out to be exactly the safe reuse budget for this too.

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 been overwritten.

Empirical validation, not just the theoretical argument

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 (budget becomes an easy-to-exceed 44 tokens), with the budget check temporarily removed to confirm the failure mode is real and not just theoretical — generated 400 tokens (well past 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. Same forced raw_cap, same over-budget scenario, with the guard in place: correctly falls through to a full re-prefill (reason=token-mismatch in the log) instead of rewinding, and stays correct.
3. A short within-budget retry under that same forced raw_cap=300 still rewinds and matches, confirming the guard isn't just conservative to the point of never firing.

Happy to re-run any of this against a different quant/ctx if useful — the scratch negative-control build was throwaway and not included here, only the guarded version in this diff.

🤖 Generated with Claude Code