Force a clean prefill after chained KV prefix-cache extensions - #2261
Open
mlpy0 wants to merge 1 commit into
Open
Force a clean prefill after chained KV prefix-cache extensions#2261mlpy0 wants to merge 1 commit into
mlpy0 wants to merge 1 commit into
Conversation
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.
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.
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.
KVPrefixCachenow tracks how many successive deep extensions built each entry's tail and stops offering it for deep reuse atMAX_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:
TODOinmlx_generate.error_messageas 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.