lablup / lablup/mlxcel

fix(prompt-cache): completion-origin snapshots claim one more token than their caches hold

Open
#1,754 1 comment 0 reactions 0 assignees View on GitHub
area:core area:inference priority:high status:ready type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem

The completion-origin donation path at `src/server/batch/scheduler/prompt_cache.rs:1114-1150` builds `tokens = prompt_tokens ++ generated_tokens` and donates with `token_len = tokens.len()`, but the model-owned KV caches at that moment hold only `tokens.len() - 1` tokens of KV. The decode step forwards `*seq.generated_tokens.last()` (`src/server/batch/scheduler/decode_tick.rs:501`, `:744`, `:1273-1293`), so the final sampled token is never fed back through the model.

The gap is conditional on the stop reason. A merged-EOS stop returns before pushing the token (`decode_tick.rs:1448-1456`), so an EOS-terminated run is exact, as is the eos-at-prefill donate, which passes an empty generated slice (`src/server/batch/scheduler/prefill.rs:1418`). The off-by-one holds for `Length`, `StopSequence`, `RepetitionLoop`, `Cancelled`, and the finished-during-prefill donate (`prefill.rs:1443`, `:1461-1470`, `:1515`; `decode_tick.rs:1458`, `:1474-1532`, `:1590`).

On restore the scheduler resumes prefill from `matched_len` (`prompt_cache.rs:214`, then `src/server/batch/scheduler/admission.rs:341-344`, then `prefill.rs:727`), not from the restored cache offset. The last generated token's KV is therefore never computed and every later token lands one slot early.

## Scope

Only a full-entry exact match is affected. Truncating restores are safe because the truncate predicates demand `target_len <= offset` (`src/models/kv_snapshot.rs:214`, `:628-640`), and the store consults the truncate predicate only when no exact entry matched (`src/server/prompt_cache/store.rs:1252-1273`).

This is pre-existing and affects every snapshot-reuse family, not only the three that PR #1752 adds: `origin/main`'s Gemma 4 stored and restored the same `.offset` scalar (`src/models/gemma4.rs:1221`, `:1314-1315`).

## Observed symptom

Recorded in PR #1752's validation section: on `models/gemma-3-4b-it-4bit`, a snapshot covering exactly one decoded token (restore length 333 of a 422-token target) makes the restored run's distribution differ materially from cold, with 13 nats between cold's top two and warm choosing a different token. Neighbouring boundaries at 334, 335, 339 and 344 are clean because those are truncating probes, which the `target_len <= offset` predicate already keeps inside the real KV.

## Proposed solution

The fix belongs in the scheduler, not in any model, and it changes behavior for every snapshot family. Two candidate repairs: donate `token_len = tokens.len() - 1` when the stop reason left the tail token unforwarded (keeping the EOS and eos-at-prefill paths exact), or have the restore path resume prefill from the restored cache offset rather than from `matched_len`. Pick one and apply it uniformly; do not special-case a family.

## Acceptance criteria

- [ ] The donated `token_len` equals the number of tokens whose KV the caches actually hold, for every finish reason (`Stop` via merged EOS, `Stop` via structured stop, `StopSequence`, `Length`, `RepetitionLoop`, `Cancelled`, finished-during-prefill, eos-at-prefill).
- [ ] A restore from a completion-origin snapshot produces the same next-token distribution as a cold run at the same target length, including when the snapshot covers exactly one decoded token.
- [ ] A regression test pins the one-decoded-token boundary, so a future edit cannot reintroduce the skew silently.
- [ ] Per `CLAUDE.md`, real-checkpoint validation across at least Gemma 4, Gemma 3 and one recurrent family before merging; build-only or synthetic-only validation is not enough.

## Verification

```bash
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
```

Manual: reproduce the 333-of-422 boundary on `models/gemma-3-4b-it-4bit` with the prompt cache on in both arms (vary only restored versus cold, per PR #1752's methodology) and confirm the restored run's top-1 and top-5 logprobs match the cold run.

Context: PR #1752 and issue #1335. This is not a fix for #1335 itself.

Contributor guide

Open the contributing guide

Research direction

Start with the completion-origin donation path in src/server/batch/scheduler/prompt_cache.rs and the token-forwarding and stop-reason branches in decode_tick.rs and prefill.rs; compare those with restore admission in admission.rs. Reproduce the one-decoded-token boundary described for PR #1752, add a regression test, and verify cold and restored distributions match before running the workspace test and clippy commands.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
ai-infra-agents, backend, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.