Skip to content

Add a use_prefix_cache request flag to bypass the KV prefix cache#2179

Open
mlpy0 wants to merge 2 commits into
exo-explore:mainfrom
mlpy0:feature/ephemeral-prefix-cache-bypass
Open

Add a use_prefix_cache request flag to bypass the KV prefix cache#2179
mlpy0 wants to merge 2 commits into
exo-explore:mainfrom
mlpy0:feature/ephemeral-prefix-cache-bypass

Conversation

@mlpy0

@mlpy0 mlpy0 commented Jun 26, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds an optional use_prefix_cache flag to chat-completion requests. It defaults to true (current behavior). Set it to false to 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:

  • periodic warm-up / liveness probes that keep the model and generation path hot, but should not accumulate throwaway cache entries;
  • one-off or automated requests that shouldn't evict other clients' cached conversation prefixes via LRU;
  • ad-hoc calls where caching the prompt brings no benefit.

BenchChatCompletionRequest already has a use_prefix_cache field for exactly this bypass in bench mode. Following review feedback, this promotes that same flag to the base ChatCompletionRequest instead of adding a separate flag, so there's one knob doing one job.

Implementation

  • ChatCompletionRequest.use_prefix_cache defaults to true. BenchChatCompletionRequest keeps use_prefix_cache = false as an override, so bench behavior is unchanged.
  • The chat-completions adapter maps it onto TextGenerationTaskParams.use_prefix_cache, whose default is now true so every existing caller keeps caching.
  • The generators (mlx_generate and ExoBatchGenerator) skip the prefix cache whenever use_prefix_cache is false, decoupled from bench. bench now 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 its use_prefix_cache = false override.

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 rltakashige left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mlpy0

mlpy0 commented Jul 1, 2026

Copy link
Copy Markdown
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.
@mlpy0 mlpy0 changed the title Add an ephemeral request flag to bypass the KV prefix cache Add a use_prefix_cache request flag to bypass the KV prefix cache Jul 1, 2026
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.

2 participants