Skip to content

fix(prompt-cache): stop a whole-prompt hit from duplicating the last token - #2136

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-1760-whole-prompt-hit
Oct 6, 2026
Merged

inureyes merged 2 commits into
mainfrom
fix/issue-1760-whole-prompt-hit

Conversation

@inureyes

@inureyes inureyes commented Oct 6, 2026

Copy link
Copy Markdown
Member

Summary

A replay of an identical prompt hits the prompt cache over every prompt token. The adopt restored all len tokens and admission then moved only the prefill cursor back to len - 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):

Checkpoint Path prompt_len offset after clamp
gemma-3-4b-it-4bit snapshot, truncating 115 115
gemma-4-12b-it-4bit snapshot, truncating 119 119
qwen3-4b-4bit paged KV, prompt padded to 128 128 128
llama-3.2-1b-4bit paged KV, prompt padded to 160 160 160

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_prefix caps every hit at len - 1 (whole_prompt_reuse_cap). Snapshot branch: the cap turns the restore into a truncating one after asking snapshot_truncatable_to; a refusal (recurrent state, wrapped sliding ring) falls back to cold prefill and records a layout_constraints reject with the entry length. KV branch: the cap runs before the dense truncate and paged block floors, with a prefix_too_short decline if it drops below the minimum.
  • The multimodal whole-entry gate keys on the store's match, not the cap, so a whole-entry VLM replay still adopts; its one-token suffix is what the old back-off already forwarded through the token path.
  • The admission back-off stays as a warn-level backstop. usage.cached_tokens now reports the restored length instead of one more.

Validation (GB10, CUDA)

Same process, cold request then warm replays, temperature 0, seed 0, 200 tokens:

Checkpoint Before After
gemma-3-4b-it-4bit diverged at char 226 3/3 byte-identical
llama-3.2-1b-4bit diverged at char 230 3/3 byte-identical
gemma-4-12b-it-4bit diverged at token 71 diverges at token 27, a bf16 tie (warm top-2 gap 0.0)
qwen3-4b-4bit diverged at token 76 diverges at token 82, a bf16 tie (warm top-2 gap 0.0)
qwen3.5-0.8b-4bit (recurrent) identical (boundary hit, not whole-prompt) identical

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 1 in 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).

  • Unit: 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).
  • e2e: identical_prompt_replay_restores_all_but_the_last_token in tests/prompt_cache_e2e.rs (ignored; qwen3-0.6b-4bit, --apc-block-size 1) asserts cached_tokens == prompt_tokens - 1 and byte-identical replies. It passes here and fails on the pre-fix binary (20 vs 19).
  • cargo clippy --release --features cuda --lib --tests -- -D warnings and cargo fmt --check clean.

Not verified on this host

Closes #1760

…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
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:inference Generation, sampling, decoding (incl. speculative, DRY) status:review Under review labels Oct 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) priority:medium Medium priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant