fix(prompt-cache): completion-origin snapshots claim one more token than their caches hold
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem
The completion-origin donation path at `src/server/batch/scheduler/prompt_cache.rs:1114-1150` builds `tokens = prompt_tokens ++ generated_tokens` and donates with `token_len = tokens.len()`, but the model-owned KV caches at that moment hold only `tokens.len() - 1` tokens of KV. The decode step forwards `*seq.generated_tokens.last()` (`src/server/batch/scheduler/decode_tick.rs:501`, `:744`, `:1273-1293`), so the final sampled token is never fed back through the model.
The gap is conditional on the stop reason. A merged-EOS stop returns before pushing the token (`decode_tick.rs:1448-1456`), so an EOS-terminated run is exact, as is the eos-at-prefill donate, which passes an empty generated slice (`src/server/batch/scheduler/prefill.rs:1418`). The off-by-one holds for `Length`, `StopSequence`, `RepetitionLoop`, `Cancelled`, and the finished-during-prefill donate (`prefill.rs:1443`, `:1461-1470`, `:1515`; `decode_tick.rs:1458`, `:1474-1532`, `:1590`).
On restore the scheduler resumes prefill from `matched_len` (`prompt_cache.rs:214`, then `src/server/batch/scheduler/admission.rs:341-344`, then `prefill.rs:727`), not from the restored cache offset. The last generated token's KV is therefore never computed and every later token lands one slot early.
## Scope
Only a full-entry exact match is affected. Truncating restores are safe because the truncate predicates demand `target_len <= offset` (`src/models/kv_snapshot.rs:214`, `:628-640`), and the store consults the truncate predicate only when no exact entry matched (`src/server/prompt_cache/store.rs:1252-1273`).
This is pre-existing and affects every snapshot-reuse family, not only the three that PR #1752 adds: `origin/main`'s Gemma 4 stored and restored the same `.offset` scalar (`src/models/gemma4.rs:1221`, `:1314-1315`).
## Observed symptom
Recorded in PR #1752's validation section: on `models/gemma-3-4b-it-4bit`, a snapshot covering exactly one decoded token (restore length 333 of a 422-token target) makes the restored run's distribution differ materially from cold, with 13 nats between cold's top two and warm choosing a different token. Neighbouring boundaries at 334, 335, 339 and 344 are clean because those are truncating probes, which the `target_len <= offset` predicate already keeps inside the real KV.
## Proposed solution
The fix belongs in the scheduler, not in any model, and it changes behavior for every snapshot family. Two candidate repairs: donate `token_len = tokens.len() - 1` when the stop reason left the tail token unforwarded (keeping the EOS and eos-at-prefill paths exact), or have the restore path resume prefill from the restored cache offset rather than from `matched_len`. Pick one and apply it uniformly; do not special-case a family.
## Acceptance criteria
- [ ] The donated `token_len` equals the number of tokens whose KV the caches actually hold, for every finish reason (`Stop` via merged EOS, `Stop` via structured stop, `StopSequence`, `Length`, `RepetitionLoop`, `Cancelled`, finished-during-prefill, eos-at-prefill).
- [ ] A restore from a completion-origin snapshot produces the same next-token distribution as a cold run at the same target length, including when the snapshot covers exactly one decoded token.
- [ ] A regression test pins the one-decoded-token boundary, so a future edit cannot reintroduce the skew silently.
- [ ] Per `CLAUDE.md`, real-checkpoint validation across at least Gemma 4, Gemma 3 and one recurrent family before merging; build-only or synthetic-only validation is not enough.
## Verification
```bash
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
```
Manual: reproduce the 333-of-422 boundary on `models/gemma-3-4b-it-4bit` with the prompt cache on in both arms (vary only restored versus cold, per PR #1752's methodology) and confirm the restored run's top-1 and top-5 logprobs match the cold run.
Context: PR #1752 and issue #1335. This is not a fix for #1335 itself.
Contributor guide
Research direction
Start with the completion-origin donation path in src/server/batch/scheduler/prompt_cache.rs and the token-forwarding and stop-reason branches in decode_tick.rs and prefill.rs; compare those with restore admission in admission.rs. Reproduce the one-decoded-token boundary described for PR #1752, add a regression test, and verify cold and restored distributions match before running the workspace test and clippy commands.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai-infra-agents, backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100