diff --git a/src/exo/api/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py index cbd545318b..0242c8d0b4 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, + use_prefix_cache=request.use_prefix_cache, ) diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py index 02e397adaa..9a868002ca 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 + # 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 29228c363f..bc9d9f0e57 100644 --- a/src/exo/shared/types/text_generation.py +++ b/src/exo/shared/types/text_generation.py @@ -115,7 +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 + # 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 5c7394ddec..aae4c9da1a 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -161,9 +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 - ): + 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 @@ -264,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: + 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 2e3d051251..1774a5794b 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 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: + if not task.use_prefix_cache: kv_prefix_cache = None # Use prefix cache if available, otherwise create fresh cache