From f18be7c4d837e135e525cd8e58d5ece996480e4e Mon Sep 17 00:00:00 2001 From: mlpy0 <96022931+mlpy0@users.noreply.github.com> Date: Fri, 26 Jun 2026 15:55:44 +0300 Subject: [PATCH 1/2] Add an ephemeral request flag to bypass the KV prefix cache 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. --- src/exo/api/adapters/chat_completions.py | 1 + src/exo/api/types/api.py | 5 +++++ src/exo/shared/types/text_generation.py | 3 +++ src/exo/worker/engines/mlx/generator/batch_generate.py | 8 +++++--- src/exo/worker/engines/mlx/generator/generate.py | 6 ++++-- 5 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/exo/api/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py index cbd545318b..f7c1d62d0d 100644 --- a/src/exo/api/adapters/chat_completions.py +++ b/src/exo/api/adapters/chat_completions.py @@ -175,6 +175,7 @@ async def chat_request_to_text_generation( presence_penalty=request.presence_penalty, frequency_penalty=request.frequency_penalty, images=images, + ephemeral=request.ephemeral or False, ) diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py index 02e397adaa..07d0d53f4d 100644 --- a/src/exo/api/types/api.py +++ b/src/exo/api/types/api.py @@ -241,6 +241,11 @@ class ChatCompletionRequest(BaseModel): tools: list[dict[str, Any]] | None = None reasoning_effort: ReasoningEffort | None = None enable_thinking: bool | None = None + # When true, run prefill+decode normally but never read from or write to the + # shared KV prefix cache: a fresh per-request cache is used and discarded. + # Useful for one-off or throwaway requests that should not evict cached + # prefixes or leave a cache entry behind. + ephemeral: bool | None = None min_p: float | None = None repetition_penalty: float | None = None repetition_context_size: int | None = None diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py index 29228c363f..4ebff6fc50 100644 --- a/src/exo/shared/types/text_generation.py +++ b/src/exo/shared/types/text_generation.py @@ -116,6 +116,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True): tools: list[dict[str, Any]] | None = None bench: bool = False use_prefix_cache: bool = False + # Run prefill+decode but never touch the shared KV prefix cache (fresh + # throwaway cache, no read/write). For one-off or throwaway requests. + ephemeral: bool = False top_k: int | None = None stop: str | list[str] | None = None seed: int | None = None diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 5c7394ddec..a49cc39da9 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -161,8 +161,10 @@ def submit( is_exact_hit = False prompt_tokens = all_prompt_tokens - if self.kv_prefix_cache is not None and ( - not is_bench or task_params.use_prefix_cache + if ( + self.kv_prefix_cache is not None + and (not is_bench or task_params.use_prefix_cache) + and not task_params.ephemeral ): cache, remaining_tokens, matched_index, is_exact_hit = ( self.kv_prefix_cache.get_kv_cache( @@ -264,7 +266,7 @@ def submit( c.values = c._trim(trim_size, c.values) c._idx = c.max_size - if not is_bench or task_params.use_prefix_cache: + if (not is_bench or task_params.use_prefix_cache) and not task_params.ephemeral: min_prefix_hit_length = max( 1000, system_prompt_token_count(task_params, self.tokenizer) ) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 2e3d051251..22bee869de 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -572,9 +572,11 @@ def mlx_generate( all_prompt_tokens = vision.prompt_tokens media_regions: list[MediaRegion] = vision.media_regions if vision else [] - # Do not use the prefix cache if we are trying to do benchmarks. + # Skip the prefix cache for benchmarks, and for ephemeral requests that must + # leave no cache trace. Nulling it here makes both the read below and the + # post-generation store a no-op. is_bench = task.bench - if is_bench and not task.use_prefix_cache: + if (is_bench and not task.use_prefix_cache) or task.ephemeral: kv_prefix_cache = None # Use prefix cache if available, otherwise create fresh cache From cb6f78c6ecce6db3a6ff55f885c74132076a51de Mon Sep 17 00:00:00 2001 From: mlpy0 <96022931+mlpy0@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:12:28 +0300 Subject: [PATCH 2/2] Use a use_prefix_cache request flag instead of a separate ephemeral flag 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. --- src/exo/api/adapters/chat_completions.py | 2 +- src/exo/api/types/api.py | 10 +++++----- src/exo/shared/types/text_generation.py | 8 ++++---- src/exo/worker/engines/mlx/generator/batch_generate.py | 8 ++------ src/exo/worker/engines/mlx/generator/generate.py | 8 ++++---- 5 files changed, 16 insertions(+), 20 deletions(-) diff --git a/src/exo/api/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py index f7c1d62d0d..0242c8d0b4 100644 --- a/src/exo/api/adapters/chat_completions.py +++ b/src/exo/api/adapters/chat_completions.py @@ -175,7 +175,7 @@ async def chat_request_to_text_generation( presence_penalty=request.presence_penalty, frequency_penalty=request.frequency_penalty, images=images, - ephemeral=request.ephemeral or False, + use_prefix_cache=request.use_prefix_cache, ) diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py index 07d0d53f4d..9a868002ca 100644 --- a/src/exo/api/types/api.py +++ b/src/exo/api/types/api.py @@ -241,11 +241,11 @@ class ChatCompletionRequest(BaseModel): tools: list[dict[str, Any]] | None = None reasoning_effort: ReasoningEffort | None = None enable_thinking: bool | None = None - # When true, run prefill+decode normally but never read from or write to the - # shared KV prefix cache: a fresh per-request cache is used and discarded. - # Useful for one-off or throwaway requests that should not evict cached - # prefixes or leave a cache entry behind. - ephemeral: bool | None = None + # Whether this request may use the shared KV prefix cache. Defaults to True. + # Set to False to run a full prefill+decode on a fresh per-request cache that + # is discarded afterwards, for one-off or throwaway requests that should not + # evict other clients' cached prefixes or leave a cache entry behind. + use_prefix_cache: bool = True min_p: float | None = None repetition_penalty: float | None = None repetition_context_size: int | None = None diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py index 4ebff6fc50..bc9d9f0e57 100644 --- a/src/exo/shared/types/text_generation.py +++ b/src/exo/shared/types/text_generation.py @@ -115,10 +115,10 @@ class TextGenerationTaskParams(BaseModel, frozen=True): stream: bool = False tools: list[dict[str, Any]] | None = None bench: bool = False - use_prefix_cache: bool = False - # Run prefill+decode but never touch the shared KV prefix cache (fresh - # throwaway cache, no read/write). For one-off or throwaway requests. - ephemeral: bool = False + # Whether to use the shared KV prefix cache. When False, run prefill+decode + # on a fresh per-request cache and store nothing (for one-off/throwaway + # requests). Defaults to True so normal generation caches as usual. + use_prefix_cache: bool = True top_k: int | None = None stop: str | list[str] | None = None seed: int | None = None diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index a49cc39da9..aae4c9da1a 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -161,11 +161,7 @@ def submit( is_exact_hit = False prompt_tokens = all_prompt_tokens - if ( - self.kv_prefix_cache is not None - and (not is_bench or task_params.use_prefix_cache) - and not task_params.ephemeral - ): + if self.kv_prefix_cache is not None and task_params.use_prefix_cache: cache, remaining_tokens, matched_index, is_exact_hit = ( self.kv_prefix_cache.get_kv_cache( self.model, all_prompt_tokens, media_regions=media_regions @@ -266,7 +262,7 @@ def submit( c.values = c._trim(trim_size, c.values) c._idx = c.max_size - if (not is_bench or task_params.use_prefix_cache) and not task_params.ephemeral: + if task_params.use_prefix_cache: min_prefix_hit_length = max( 1000, system_prompt_token_count(task_params, self.tokenizer) ) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 22bee869de..1774a5794b 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -572,11 +572,11 @@ def mlx_generate( all_prompt_tokens = vision.prompt_tokens media_regions: list[MediaRegion] = vision.media_regions if vision else [] - # Skip the prefix cache for benchmarks, and for ephemeral requests that must - # leave no cache trace. Nulling it here makes both the read below and the - # post-generation store a no-op. + # Skip the prefix cache when the request opts out (use_prefix_cache=False), + # e.g. benchmarks or one-off requests that must leave no cache trace. Nulling + # it here makes both the read below and the post-generation store a no-op. is_bench = task.bench - if (is_bench and not task.use_prefix_cache) or task.ephemeral: + if not task.use_prefix_cache: kv_prefix_cache = None # Use prefix cache if available, otherwise create fresh cache