Add a use_prefix_cache request flag to bypass the KV prefix cache#2179
Open
mlpy0 wants to merge 2 commits into
Open
Add a use_prefix_cache request flag to bypass the KV prefix cache#2179mlpy0 wants to merge 2 commits into
mlpy0 wants to merge 2 commits into
Conversation
An ephemeral chat-completion request runs prefill+decode normally but never reads from or writes to the shared KV prefix cache: it uses a fresh per-request cache and stores nothing. This lets a caller issue a real generation without disturbing the prefix cache - keeping the model and generation path warm without accumulating throwaway entries, and without evicting other callers' cached prefixes via LRU. Threaded request body -> adapter -> TextGenerationTaskParams -> both the sequential (mlx_generate) and batch (ExoBatchGenerator) generators, mirroring the existing enable_thinking plumbing and reusing the bench-mode cache-bypass path. Defaults false, so behavior is unchanged for existing callers.
rltakashige
reviewed
Jun 30, 2026
rltakashige
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the suggestion and PR!
I'm not sure how necessary this is, as BenchChatCompletionRequest contains a use_prefix_cache parameter that seems to be what this does. Perhaps it makes sense to move use_prefix_cache into the ChatCompletionRequest instead.
What do you think? Looks like a harmless improvement so I am also inclined to approve this.
Contributor
Author
|
Yeah, agreed. I've pushed that version. use_prefix_cache now lives on ChatCompletionRequest and defaults to true, with BenchChatCompletionRequest overriding it to false so bench stays the same. The generators only looked at it in bench mode before, so I changed them to skip the cache whenever it's false instead. Ephemeral flag's gone. |
Per review: fold the cache bypass into a single use_prefix_cache knob on ChatCompletionRequest (default True) rather than adding a new ephemeral flag. BenchChatCompletionRequest keeps use_prefix_cache=False as an override, so bench behavior is unchanged. The generators now skip the prefix cache whenever use_prefix_cache is False, decoupled from bench, and TextGenerationTaskParams defaults use_prefix_cache to True so existing callers keep caching.
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.
Summary
Adds an optional
use_prefix_cacheflag to chat-completion requests. It defaults totrue(current behavior). Set it tofalseto run a normal prefill+decode on a fresh per-request cache that is discarded, so the request never reads from or writes to the shared KV prefix cache.Motivation
Sometimes a caller wants to run a real generation without disturbing the prefix cache. For example:
BenchChatCompletionRequestalready has ause_prefix_cachefield for exactly this bypass in bench mode. Following review feedback, this promotes that same flag to the baseChatCompletionRequestinstead of adding a separate flag, so there's one knob doing one job.Implementation
ChatCompletionRequest.use_prefix_cachedefaults totrue.BenchChatCompletionRequestkeepsuse_prefix_cache = falseas an override, so bench behavior is unchanged.TextGenerationTaskParams.use_prefix_cache, whose default is nowtrueso every existing caller keeps caching.mlx_generateandExoBatchGenerator) skip the prefix cache wheneveruse_prefix_cacheisfalse, decoupled frombench.benchnow only governs benchmark measurement, not caching.Compatibility
Defaults to
true, so existing callers are unaffected. The bench path still bypasses the cache by default via itsuse_prefix_cache = falseoverride.