fix(prompt-cache): a whole-prompt cache hit re-runs the last token on top of a cache that already holds it #1760

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. (PR fix(prompt-cache): stop a whole-prompt hit from duplicating the last token #2136: gemma-3-4b byte-identical; gemma-4-12b token-identical with bitwise-equal logprobs once forward widths match, and the residual default-flag divergence is the (b) effect at a bf16 tie; trinity-nano is not present on the validation host, so Qwen3 and Llama 3.2 cover the KV branch.)
• If (a): a family that cannot truncate falls back to a cold prefill with the decline recorded, and never installs a partial restore.
• Not applicable, (a) held. 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

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.

Resolved by PR #2136: the measured per-layer offset after the clamp was prompt_tokens.len() on every family tested, so (a) held.