Repository navigation
fix(prompt-cache): stop a whole-prompt hit from duplicating the last token - #2136
Merged
Merged
Conversation
…token Replaying an identical prompt hits the prompt cache over every prompt token. The adopt restored all of them, and admission then backed only the prefill cursor off to `len - 1`, so prefill forwarded the last prompt token onto a cache that already held it. Measured on CUDA with per-layer offsets after the clamp: the cache held `len` tokens (Gemma 3 115/115, Gemma 4 119/119, Qwen3 128/128, Llama 3.2 160/160), which is hypothesis (a) of the issue, the same broadcast abort #2096 reported on M5. `try_adopt_cached_prefix` now caps every hit at `len - 1` (`whole_prompt_reuse_cap`). On the snapshot branch the cap turns the restore into a truncating one; a model that refuses that truncation (recurrent state, wrapped sliding ring) falls back to cold prefill and records a `layout_constraints` reject. On the KV branch the dense truncate and paged block floors install the capped length. The admission back-off stays as a warn-level backstop. `usage.cached_tokens` now reports the restored length instead of over-counting by one. Tests: three tiny-Gemma 3 scheduler tests read the restored per-layer offset from the model's own snapshot, including the refusal fallback, and an ignored e2e test in tests/prompt_cache_e2e.rs asserts `cached_tokens == prompt_tokens - 1` and byte-identical cold and warm replies (it fails on the pre-fix binary with 20 vs 19). Closes #1760
This was referenced Oct 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
A replay of an identical prompt hits the prompt cache over every prompt token. The adopt restored all
lentokens and admission then moved only the prefill cursor back tolen - 1, so prefill forwarded the last prompt token onto a cache that already held it. The adopt now stops one token short on both the snapshot and the KV branch, so the re-run lands at its own position.Hypothesis settled: (a), a real duplicate
Per-layer cache offsets read right after the admission clamp on the pre-fix binary (GB10, CUDA,
mlxcel-server, prompt cache on):The KV path reaches a whole-prompt hit only when the prompt length is a multiple of the paged block (32 by default), which is why #2096 saw it only at those lengths.
What changed
try_adopt_cached_prefixcaps every hit atlen - 1(whole_prompt_reuse_cap). Snapshot branch: the cap turns the restore into a truncating one after askingsnapshot_truncatable_to; a refusal (recurrent state, wrapped sliding ring) falls back to cold prefill and records alayout_constraintsreject with the entry length. KV branch: the cap runs before the dense truncate and paged block floors, with aprefix_too_shortdecline if it drops below the minimum.usage.cached_tokensnow reports the restored length instead of one more.Validation (GB10, CUDA)
Same process, cold request then warm replays,
temperature 0,seed 0, 200 tokens:The remaining Gemma 4 and Qwen3 divergences are the forward-width effect, hypothesis (b): the warm request forwards 1 (Gemma 4) or 32 (Qwen3, paged block floor) tokens where cold forwards a wider chunk. Control with equal widths (
--prefill-chunk-size 1in both arms): after the fix both are token-identical with bitwise-equal logprobs (max |dlogprob| 0.0); before the fix both still diverge (Gemma 4 differs at the first token, -0.5 vs -0.625).scheduler_whole_prompt_hit_tests(3 tests, tiny Gemma 3) read the restored offset from the model's own snapshot: identical replay, whole-entry hit through admission, and the wrapped-ring refusal falling back to cold. The scheduler prompt-cache suites pass (37 tests).identical_prompt_replay_restores_all_but_the_last_tokenintests/prompt_cache_e2e.rs(ignored; qwen3-0.6b-4bit,--apc-block-size 1) assertscached_tokens == prompt_tokens - 1and byte-identical replies. It passes here and fails on the pre-fix binary (20 vs 19).cargo clippy --release --features cuda --lib --tests -- -D warningsandcargo fmt --checkclean.Not verified on this host
trinity-nano-preview-4bit(AFMoE) is not present; qwen3-4b and llama-3.2-1b stand in to cover the KV branch, which the issue's checkpoints did not reach.lenwhile the mask is built forlen - 1).Closes #1760