Reuse DSpark snapshots for immediate session rewinds

[emilianbold]

Speculative decoding can become slow when the server decides to roll back to a token boundary based on the model's protocol (e.g. near tool-call boundaries). For DeepSeek, truncating the token history alone cannot restore the compressed/recurrent state, so the fallback requires rebuilding the retained prefix.

Rollback for speculative rejection is already handled inside the inference engine using snapshots. However, that rollback opportunity was not retained for a subsequent server-requested rewind.

This tiny PR retains a weak handle to the latest usable snapshot (only one) until it is consumed or an operation invalidates its dependencies. It reuses existing snapshot buffers, adding only a small metadata record. This allows a server-triggered rollback to restore the frontier and replay a short tail. On a miss, it preserves the existing truncate-and-invalidate behavior.

The initial fast path covers fully accepted greedy-verifier blocks on resident, text-only, non-TP Metal DSpark sessions.

Experiments demonstrated degenerate cases in heavy tool-use sessions where, with snapshot reuse deliberately disabled, prefix rebuilding consumed over half of the observed session wall time. Individual rebuilds became more expensive as context grew.

This does not mean speculative decoding should be the default. We gain speed in tool-call generation, but I do not see a clear overall advantage over plain decoding on my machine. (Selectively speculating during tool-call generation might do better, but that is a separate idea not implemented here.)

Reviewer Numbers

Measured on an M5 Max, 40 GPU cores, 128 GB RAM, using Flash 0731 IQ2XXS-w2Q2K and matching DSpark support.

Evidence	Result
Additional retained metadata	660 bytes/session on the tested ABI
Existing Flash frontier payload reused	12,206,080 bytes = 11.64 MiB
Additional snapshot tensors or saved logits	None
No-rewind control, 128 generated tokens	Main 2258.37 ms, patched 2258.79 ms, approximately +0.018%
No-rewind output consistency	Identical token hashes and block counts across 10 measured runs per build

For a 4,095-token starting prefix, the conditional rewind benchmark measured 26ms for 1 token, 132 ms for 5 tokens.

The real-use logs provide the stronger motivation:

• Forced-miss session: 21 tool-boundary rebuilds consumed 1,320.5 seconds, or 68% of the 32m15s observed window. Individual rebuilds ranged from 13.4 seconds at 8.5k tokens to roughly 100 seconds near 54k tokens.
• Reuse-enabled session: all 25 requested rewinds used snapshots, costing 2.544 seconds total. The 23 tool-boundary repairs had a 111 ms median.

Note the real-use logs were from different sessions, so do not present their total-duration difference as a measured end-to-end speedup.

[aaa2015]

Thanks @emilianbold for this great PR! The benchmark numbers (111 ms median vs up to 100s full prefill rebuilds) match the severe Degenerate Rebuild latency we observed on long-context tool-calling sessions on Apple Silicon. This provides the exact engine-level complement needed for PR #1000's server-level stop-token rollback.

I went through the implementation and ran local testing on Apple Silicon. Here are a couple of architectural observations and suggestions that could make this even more powerful:

---

1. Zero-Replay Frontier Rollback when pos == start (Handling Total Draft Rejection)

In ds4_session_rewind() (ds4.c:75506):

if (s->checkpoint_valid && !s->engine->tp.active &&
    s->dspark_rollback_end == s->checkpoint.len &&
    pos > s->dspark_rollback_start) {
    ...
    /* Replay at least one token: the snapshot contains no saved logits.
     * Do not propose drafts or expose a partially restored checkpoint. */
    for (int i = start; state_ok && i < pos; i++) {
        state_ok = metal_graph_eval_token_raw_swa(...);
        if (state_ok) s->checkpoint.len++;
    }
}

Notice that it currently requires pos > s->dspark_rollback_start:

• The Gap: If a speculative block has 0 draft tokens accepted (or if the caller needs to rewind strictly back to the pre-speculative baseline pos == start), pos > start evaluates to false. As a result, state_ok stays false, and line 75524 falls through to s->checkpoint_valid = false;, triggering a full rebuild!
• Opportunity: If we preserve s->logits into a cached session buffer (s->dspark_rollback_logits, allocated when need_spec_verifier is true) during spec_frontier_snapshot(), we can expand the condition to pos >= s->dspark_rollback_start:

  if (pos == start) {
      if (state_ok && s->dspark_rollback_logits && s->logits) {
          memcpy(s->logits, s->dspark_rollback_logits,
                 (size_t)DS4_N_VOCAB * sizeof(s->logits[0]));
      } else {
          state_ok = false;
      }
  } else {
      /* Replay tokens up to pos: evaluating token pos-1 leaves s->logits set for pos. */
      for (int i = start; state_ok && i < pos; i++) {
          ...
      }
  }

• Result: Rewinding to pos == start takes 0 token replays (~0.00 ms), and total draft rejections no longer invalidate the entire session checkpoint. We verified this locally with ./tests/test_session_state_gpu, and all GPU assertions pass cleanly.

---

2. Server-Side Routing in ds4_server.c

To allow live traffic in ds4-server to hit this path, there is a minor gatekeeper check in ds4_server.c:12371:

const int rewind_to = live_prefix_rewind_target(
    ds4_engine_is_glm_dsa(s->engine), old_pos,
    j->req.prompt.len, common);

backend_can_rewind currently only passes ds4_engine_is_glm_dsa. When antirez reviews/merges this PR, exposing an engine capability helper (e.g. bool ds4_engine_can_rewind(ds4_engine *e) that returns true for GLM or Metal DSpark) will allow ds4-server to route DeepSeek rewind requests through this new fast path instead of falling back to token mismatch.

(Additionally, relaxing live_prefix_rewind_target to allow rewinding to common when common < prompt_len turns multi-turn tool-call branching from an eviction miss into a fast delta prefill).

---

3. Future Scope: Multimodal / Vision Sessions

Excluding !ds4_session_has_vision_state(s) makes complete sense for keeping this initial PR safe and minimal. As a future follow-up, since image embeddings in DeepSeek are static in the prompt prefix and speculative drafting only occurs at the tail text positions, we can likely safely allow DSpark snapshot reuse as long as pos >= last_vision_token_end.

---

Overall, fantastic work on this PR! Keeping this snapshot handle in ds4.c fundamentally solves the 100s prefill rebuild problem on Metal.

[emilianbold]

@aaa2015 thanks the the feedback. Some of that can be follow-ups but none seem mandatory. I think the current PR has the best size to be reasonably reviewed and merged.