Add a use_prefix_cache request flag to bypass the KV prefix cache

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,
     )
 
 

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

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

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)
             )

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