Skip to content

Force a clean prefill after chained KV prefix-cache extensions - #2261

Open
mlpy0 wants to merge 1 commit into
exo-explore:mainfrom
mlpy0:fix/kv-prefix-cache-chain-depth
Open

Force a clean prefill after chained KV prefix-cache extensions#2261
mlpy0 wants to merge 1 commit into
exo-explore:mainfrom
mlpy0:fix/kv-prefix-cache-chain-depth

Conversation

@mlpy0

@mlpy0 mlpy0 commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Long tool loops extend the same prefix-cache entry on every iteration. Those chained partial prefills accumulate numerical drift in the cached KV: on the same prompt, P(end-of-turn) inside an unterminated tool call goes from logprob -11.5 after a clean one-shot prefill to -0.375 once the entry has been extended twice — enough to flip borderline end-of-turn decisions and end the generation in the middle of a tool call.

KVPrefixCache now tracks how many successive deep extensions built each entry's tail and stops offering it for deep reuse at MAX_CHAIN_DEPTH (default 2, EXO_KV_MAX_CHAIN_DEPTH). The next request re-prefills from scratch and the following save resets the entry. Shallow prefix borrows — the shared system/tools block — are unaffected, since that region is written once by the entry's first prefill and never rewritten.

Two smaller fixes from the same incident:

  • The master draws a sampling seed when the request doesn't set one, so a retry of a failed generation doesn't replay the exact same failure. Bench requests stay unseeded and deterministic. This is the TODO in mlx_generate.
  • The Ollama adapter no longer echoes error_message as assistant content. Ollama's response shape has no error field, so that string is rendered as the model's answer, and it can carry raw generation text. The detail stays in the server logs.

Tests cover the chain-depth accounting (extend / shallow rebuild / shrink / re-save), the reuse and eviction paths at the ratio boundary, and seed defaulting.

Long tool loops extend the same prefix-cache entry on every iteration. Those
chained partial prefills accumulate numerical drift in the cached KV: on the
same prompt, P(end-of-turn) inside an unterminated tool call goes from logprob
-11.5 after a clean one-shot prefill to -0.375 once the entry has been extended
twice — enough to flip borderline end-of-turn decisions and end the generation
in the middle of a tool call.

KVPrefixCache now tracks how many successive deep extensions built each entry's
tail and stops offering it for deep reuse at MAX_CHAIN_DEPTH (default 2,
EXO_KV_MAX_CHAIN_DEPTH). The next request re-prefills from scratch and the
following save resets the entry. Shallow prefix borrows — the shared
system/tools block — are unaffected.

Two smaller fixes from the same incident:

- The master draws a sampling seed when the request doesn't set one, so a retry
  of a failed generation doesn't replay the exact same failure. Bench requests
  stay unseeded and deterministic. This is the TODO in mlx_generate.
- The Ollama adapter no longer echoes error_message as assistant content.
  Ollama's response shape has no error field, so that string is rendered as the
  model's answer, and it can carry raw generation text. The detail stays in the
  server logs.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant