server: avoid full prefix rewind when speculative decode stops

[aaa2015]

Summary

Commit 5b3cc8b introduced server_generation_rewind() to resample across syntax boundaries and discard unconsumed speculative tokens. However, the rewind was called whenever kept < ntok, without guarding against the case where generation had already stopped (stop_decode == true).

On architectures where rewind cannot be performed in place (such as DeepSeek V4 recurrent/compressed DSA cache), ds4_session_rewind() invalidates the checkpoint, which causes server_generation_rewind() to detect rebuild = true and synchronously re-prefill the entire prompt context from token zero before breaking out of the decode loop.

This resulted in a 10–28 second freeze at the tail of streaming responses right after text output finished, right before emission of [DONE].

Observed Behavior & Profiling

On macOS (Apple Silicon / Metal), running DeepSeek-V4-Flash-Vision-Exp with --dspark:

worker_main (ds4_server.c:13446)
  -> generate_job (ds4_server.c:13311)
    -> generate_job_inner (ds4_server.c:12818)
      -> server_generation_rewind (ds4_server.c:11358)
        -> server_session_sync_multimodal (ds4_server.c:11263)
          -> ds4_session_sync (ds4.c:66770)
            -> metal_graph_prefill_chunked_range (ds4.c:36999)

• When decoding ended due to EOS or tool end (saw_tool_end), stop_decode was set to true.
• The speculative draft block still had unconsumed tokens (kept < ntok).
• Because line 12814 lacked !stop_decode, server_generation_rewind() triggered a full prompt re-prefill (5,000–10,500 tokens, taking 14–28 seconds).
• Immediately after that re-prefill finished, if (stop_decode) break; ran, exiting the decode loop. The newly prefilled session was unused, and only then did the server flush [DONE].

Fix

1. Gate server_generation_rewind() with !stop_decode so it only runs during active mid-stream resampling boundaries.
2. When stop_decode is true and kept < ntok, rewind the session position to block_start + kept without rebuilding/re-prefilling, then break immediately.
3. Mark stop_decode = true when reaching max_tokens or the context limit.

Discarding the unconsumed draft is not optional

Skipping the rewind entirely at the stop boundary is not equivalent. A block
such as [CLOSE, EXTRA] where CLOSE completes a tool call leaves EXTRA
committed to the session although the client never saw it (kept = 1,
ntok = 2, stop_decode = true). Finalization then binds the tool-call
continuation to that longer session, so the next matching tool-result request
replays [prior history, CLOSE, EXTRA, tool result] instead of
[prior history, CLOSE, tool result] — silent state poisoning across turns.

The stop branch therefore rewinds unconditionally, on every backend. DeepSeek
has no rollback frontier and drops its checkpoint here, which costs one rebuild
on the next request; that is the correct trade against a wrong prefix. On
engines that can roll back in place (GLM, and DeepSeek once the DSpark snapshot
reuse lands) the same call is cheap.

Validation

make ds4_test && ./ds4_test --server
→ server: OK
→ ds4 tests: ok

make test-session-state
→ session state tests: ok
→ TP command tests: ok

make tests/test_session_state_gpu && ./tests/test_session_state_gpu
→ session state tests: ok

Machine: Apple M4, macOS 26.4.1, 16 GB. Backend: Metal.

• Verified on live DeepSeek V4 Flash Vision with --dspark:
  • Standard EOS response: tail delay dropped from 24.4s to 0.000s (total wall time 0.89s).
  • Tool calls: tool completion emitted immediately (tail delay 0.14s).
  • Multi-turn conversation continuation retains KV cache hits without regression.
• The model-backed suites (--logprob-vectors, --long-context, ...) were not
  run because this machine has no GGUF checked out.

[emilianbold]

This condition on GLM seems odd and AI found a counterexample:

A DeepSeek speculative block returns [CLOSE, EXTRA], where CLOSE completes a tool call and EXTRA is a subsequent token. Both have already been evaluated and committed to the session.

The server consumes CLOSE, detects the tool end, and stops with:
kept = 1
ntok = 2
stop_decode = true

In the final PR revision, stop_decode skips server_generation_rewind(), and the direct terminal rewind is GLM-only. DeepSeek therefore retains EXTRA, even though it was never consumed into the response.

For a successful Responses tool turn, finalization binds the tool-call ID to that longer session. A subsequent matching tool-result request copies the entire live history:

Expected: [prior history, CLOSE,        tool result]
Actual:   [prior history, CLOSE, EXTRA, tool result]

The response can finish quickly and the next request can report a cache hit, while still conditioning generation on the discarded EXTRA.

[aaa2015]

Good catch, thank you! You are completely right.

Restricting the terminal rewind to GLM in 0cb5c28 was an attempt to avoid invalidating the DeepSeek recurrent checkpoint (which happens because DeepSeek's compressed DSA cache cannot be truncated in-place). However, as your counterexample demonstrates, leaving EXTRA in slot->session causes silent state poisoning across tool-result turns in the Responses API.

I have pushed commit 472c63d to remove the GLM-only guard so that all backends properly rewind to block_start + kept on stop_decode.

This ensures:

1. Discarded speculative draft tokens (EXTRA) are strictly excluded from the session before binding tool-call continuation handles.
2. The current streaming request exits cleanly with 0 tail latency (without freezing for 14-28s doing a synchronous full-prompt rebuild before emitting [DONE]).

[emilianbold]

Note I can't merge PRs, but I will look over this later today or tomorrow.

I thought a bit about speculative decoding and this token boundary and found another approach to speed things up (see PR #1003). This may solve your problem too, but I haven't yet checked.

[aaa2015]

Thank you @emilianbold! PR #1003 is brilliant and a huge leap forward.

I looked closely at PR #1003 (Reuse DSpark snapshots for immediate session rewinds). Retaining the weak handle to the DSpark verification snapshot turns what used to be a full prefix rebuild (10s to 100s on long contexts) into an instantaneous ~100ms restore and replay. That fundamentally solves the recurrent DSA rollback bottleneck on DeepSeek!

The two PRs actually complement each other nicely across two different layers:

• PR server: avoid full prefix rewind when speculative decode stops #1000 operates at the server lifecycle layer: it ensures that once decoding has terminated (stop_decode == true, e.g. on EOS, client stop sequence, or tool-end), the server exits cleanly without attempting mid-stream resampling or delaying the emission of [DONE].
• PR Reuse DSpark snapshots for immediate session rewinds #1003 operates at the engine session rewind layer: when ds4_session_rewind() is invoked (whether at the terminal boundary from server: avoid full prefix rewind when speculative decode stops #1000 or during mid-stream syntax resampling), it restores the snapshot directly and preserves checkpoint_valid = true rather than invalidating the checkpoint.

Together, they make the speculative tool-calling lifecycle both mathematically sound (no phantom EXTRA draft tokens leaked across tool-result turns) and lightning fast (0-delay stream termination + ~100ms checkpoint preservation).

Really appreciate your sharp review on #1000 and the fantastic follow-up in #1003!

[aaa2015]

Heads-up before you take a look: the branch has been rebased onto current main (6289c51), so the head is now 39cb595 rather than 472c63d.

The mechanism is the same one described above — gate the mid-stream resample at the terminal boundary, and rewind the session in place instead — with one addition: the terminal rewind is now skipped when finish == "error", since a failed turn has no valid boundary to rewind to.

For your counterexample the net effect is unchanged: [CLOSE, EXTRA] with kept = 1, ntok = 2 rewinds to block_start + 1 on every backend before finalization binds the tool-call continuation, so the next matching tool-result request replays [prior history, CLOSE, tool result].

To be explicit about the cost, since I glossed over it before: on DeepSeek without #1003, ds4_session_rewind() still invalidates the checkpoint, so the next request pays one rebuild. That is deliberate — a wrong prefix is worse than a rebuild — and #1003 is exactly what removes the penalty by letting the rewind restore the snapshot instead. On GLM the terminal rewind is already cheap.

[emilianbold]

I just had my AI do a quick pass and found another counterexample:

The unconditional rewind fixes the EXTRA leak, but 39cb595 has another counterexample:

DeepSeek returns [CLOSE, EXTRA].
Server keeps CLOSE and rewinds to block_start + 1.
The history is correct, but checkpoint_valid becomes false.
Finalization remembers the tool-call ID and retained position.
Client sends the matching tool result without replaying full history.
The IDs and position match, but build_live_prompt_suffix() calls ds4_session_rebase_vision_state(), which rejects invalid checkpoints even with zero images. The continuation helper returns zero, and the server responds HTTP 409 before reaching sync.
So this case does not merely pay a rebuild on the next request: it rejects the continuation. It affects both Responses and Anthropic tool-output-only requests. I reproduced both using the actual session bookkeeping and server response path with model-free fixtures.

[aaa2015]

Sharp catch, @emilianbold! Thank you for the detailed trace and the model-free reproduction. You are completely right.

I traced the failure mode you described:

1. When DeepSeek rewinds without snapshot restoration, ds4_session_rewind() marks s->checkpoint_valid = false.
2. When the client sends an Anthropic or Responses continuation containing only the new tool output (anthropic_requires_live_tool_state = true), anthropic_live_continuation_prompt() calls build_live_prompt_suffix().
3. In build_live_prompt_suffix() (line 10641), ds4_session_rebase_vision_state() checks if (!s || !s->checkpoint_valid ...). Because checkpoint_valid is false, it rejects the session even when image_count == 0.
4. As a result, cached evaluates to 0, and line 12398 immediately rejects the continuation with HTTP 409 because the client did not supply the historical messages to rebuild from.

This also explains why commit 5b3cc8b originally added server_session_sync_multimodal() inside server_generation_rewind(): Salvatore needed checkpoint_valid to be restored to true before returning so that subsequent continuation handles wouldn't 409. But doing a full synchronous prefill on every stop boundary meant paying that 14–28s penalty at the tail of every tool turn (and even regular EOS completions).

This highlights why your PR #1003 (Reuse DSpark snapshots for immediate session rewinds) is the critical missing piece of the puzzle:

• PR server: avoid full prefix rewind when speculative decode stops #1000 fixes the server protocol boundary (truncating unconsumed draft tokens so EXTRA never leaks, and preventing mid-stream resampling on stop).
• PR Reuse DSpark snapshots for immediate session rewinds #1003 enables the engine to restore the DSpark verification snapshot during ds4_session_rewind(), which keeps s->checkpoint_valid = true without paying the full prefill rebuild.

When tested with PR #1003 applied, ds4_session_rewind(block_start + 1) restores the snapshot and replays the single CLOSE token in ~2ms:

• s->checkpoint_valid remains true.
• build_live_prompt_suffix() succeeds.
• Tail latency drops from 14–28s to ~0s.
• Anthropic and Responses tool continuations succeed seamlessly without 409.

For environments where #1003 is not yet present (e.g. non-Metal backends or snapshot misses), we have two possible strategies:

1. Either gate the full prompt rebuild in server: avoid full prefix rewind when speculative decode stops #1000 so that it only runs when finish == "tool_calls" and the engine cannot rewind cleanly (saving the 14–28s freeze for all regular text / EOS completions, while keeping tool continuations alive until Reuse DSpark snapshots for immediate session rewinds #1003 lands).
2. Or land Reuse DSpark snapshots for immediate session rewinds #1003 alongside server: avoid full prefix rewind when speculative decode stops #1000 so DeepSeek Metal DSpark natively retains valid checkpoints on rewind.

What are your thoughts on landing #1003 first / in tandem, or would you prefer a fallback guard in #1000 for engines without snapshot rollback?