lablup / lablup/mlxcel

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

Open
#1,760 0 comments 0 reactions 0 assignees View on GitHub
area:inference priority:medium status:ready type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.