fix(prompt-cache): a whole-prompt cache hit re-runs the last token on top of a cache that already holds it
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
Replaying an identical prompt in the same server process returns a different continuation from the first request, under greedy decoding, on every family tested. The first request is a cold prefill, the second hits the prompt cache with a prefix that covers the whole prompt, and the two replies diverge instead of matching token for token.
## Current Behavior
`src/server/batch/scheduler/admission.rs:478-492` guards the degenerate hit where the adopted prefix covers the entire tokenized prompt (a client replaying an identical prompt) by backing `prefill_start_offset` off to `prompt_tokens.len() - 1`, so the prefill still runs and the sampler sees fresh logits. The restore has already happened at that point: `try_adopt_cached_prefix` (`src/server/batch/scheduler/prompt_cache.rs:118`) restores the full `matched_len` at `prompt_cache.rs:203-217` and returns it at `prompt_cache.rs:244`, and the clamp lowers only the prefill cursor, not the restored cache. The prefill then forwards `prompt_tokens[prefill_start_offset..]` (`src/server/batch/scheduler/prefill.rs:727`), the last prompt token, onto a cache that already contains it.
## Evidence
Server on port 19335, prompt cache on, `temperature 0`, `seed 0`, `logprobs` on, one process, the same single-user-turn prompt sent twice:
- `models/gemma-4-12b-it-4bit`: request 2 logged `snapshot hit: restored 39/39 tokens, stored=130, partial=true` plus `prompt-cache hit covered the entire prompt; re-running the last token through prefill`; the two replies share only their first 198 characters (lengths 514 and 535).
- `models/gemma-3-4b-it-4bit` with a 35-token prompt: `restored 35/35, stored=125, partial=true`, same back-off log line, replies share their first 126 characters (lengths 501 and 601).
- `models/trinity-nano-preview-4bit` (AFMoE): `restored 1281/1281, stored=1401, partial=true`, same back-off log line, replies diverge at the first character.
## Scope
This is pre-existing, not introduced by PR #1752. Gemma 4 has shipped this path since before that branch and reproduces it identically, which is why it is filed separately. #1752 widens the set of families that can reach it, by giving Gemma 3, AFMoE and Llama 4 snapshot support.
## Open question: two candidate explanations
Both are still open, and distinguishing them is the first task.
- (a) The last token is genuinely duplicated in the cache, so the model attends over a prompt whose final token appears twice.
- (b) The divergence is only the forward-width effect described in `docs/benchmarks.md:407` ("Judging a change that moves the numbers"): the warm request forwards a 1-token prefill where the cold request forwards the whole prompt in one chunk, and a near-tied position flips.
Either assert the per-layer cache offset against `prompt_tokens.len() - 1` right after the clamp, or compare against a replay whose prompt is one token shorter than the stored entry so the clamp never fires.
## Suggested direction if (a) holds
Clamp `matched_len` to `tokens.len() - 1` inside `try_adopt_cached_prefix` before restoring, so the restore installs one fewer token and the prefill's re-run lands at the right position. The complication is that this turns an exact-prefix restore into a truncating one, which a family whose rotating ring has wrapped will refuse (`snapshot_truncatable_to` returns false; `src/models/gemma4.rs:5849` and `src/models/gemma4.rs:1226`), so the fallback on refusal has to be a cold prefill rather than a silently wrong restore.
## Acceptance Criteria
- [ ] The hypothesis is settled and the measurement recorded: the per-layer cache offset after the clamp is either `prompt_tokens.len()` (duplicate) or `prompt_tokens.len() - 1` (forward-width effect).
- [ ] If (a): a warm replay of an identical prompt at `temperature 0` and `seed 0` produces the same continuation as the cold first request on all three checkpoints above.
- [ ] If (a): a family that cannot truncate falls back to a cold prefill with the decline recorded, and never installs a partial restore.
- [ ] If (b): `admission.rs:478-492` gains a comment naming the forward-width effect, so the next reader does not re-derive it.
- [ ] A regression test in `tests/prompt_cache_e2e.rs` covers the whole-prompt hit and asserts the restored cache offset, not merely that the request succeeds.
## Verification
```bash
cargo build --release --features metal,accelerate
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
```
Manual: start `mlxcel-server` with the prompt cache enabled, send the same single-turn chat completion twice at `temperature 0` and `seed 0`, and compare the two `content` strings byte for byte. A pass is identical strings with the `prompt-cache hit covered the entire prompt` line still present for the second request.
## Related
#1752 widens the affected family set. Follow-ups from the same review: #1754, #1755, #1756.
Contributor guide
Research direction
Start with the whole-prompt path in src/server/batch/scheduler/admission.rs:478-492, then trace restoration in prompt_cache.rs:118, 203-217, and 244 and prefill in prefill.rs:727. Measure the per-layer cache offset after the clamp, using the shorter-prompt comparison if useful, and run tests/prompt_cache_e2e.rs. Done means the duplicate-token versus forward-width hypothesis is recorded, with the appropriate fix or comment and a regression test asserting the restored offset.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, performance, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100