← back to Exo
Only update KV prefix cache on a good cache hit (#1817)
c6815bfdcef8a2a62f37e819449cf7da298d4b64 · 2026-03-30 15:04:38 +0100 · rltakashige
## Motivation
Addresses #1816
## Changes
Update on min prefix cache > min_prefix_hit_length **and** hit ratio >
_MIN_PREFIX_HIT_RATIO_TO_UPDATE
min_prefix_hit_length = max(1000, system prompt length) -> system
prompts must match exactly.
## Test Plan
### Manual Testing
Test on OpenCode and Claude Code
Files touched
M src/exo/worker/engines/mlx/generator/batch_generate.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/engines/mlx/utils_mlx.py
Diff
commit c6815bfdcef8a2a62f37e819449cf7da298d4b64
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Mar 30 15:04:38 2026 +0100
Only update KV prefix cache on a good cache hit (#1817)
## Motivation
Addresses #1816
## Changes
Update on min prefix cache > min_prefix_hit_length **and** hit ratio >
_MIN_PREFIX_HIT_RATIO_TO_UPDATE
min_prefix_hit_length = max(1000, system prompt length) -> system
prompts must match exactly.
## Test Plan
### Manual Testing
Test on OpenCode and Claude Code
---
.../worker/engines/mlx/generator/batch_generate.py | 13 ++++++++++--
src/exo/worker/engines/mlx/generator/generate.py | 6 ++++--
src/exo/worker/engines/mlx/utils_mlx.py | 23 ++++++++++++++++++++++
3 files changed, 38 insertions(+), 4 deletions(-)
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index a6d9646b..143c4d40 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -40,7 +40,10 @@ from exo.worker.engines.mlx.generator.generate import (
patch_embed_tokens,
prefill,
)
-from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens
+from exo.worker.engines.mlx.utils_mlx import (
+ fix_unmatched_think_end_tokens,
+ system_prompt_token_count,
+)
from exo.worker.engines.mlx.vision import (
MediaRegion,
VisionProcessor,
@@ -211,12 +214,16 @@ class ExoBatchGenerator:
c._idx = c.max_size
if not is_bench:
+ min_prefix_hit_length = max(
+ 1000, system_prompt_token_count(task_params, self.tokenizer)
+ )
self._save_prefix_cache(
all_prompt_tokens,
list(cache),
cache_snapshots,
prefix_hit_length,
matched_index,
+ min_prefix_hit_length,
media_regions,
)
@@ -426,6 +433,7 @@ class ExoBatchGenerator:
cache_snapshots: list[CacheSnapshot] | None,
prefix_hit_length: int,
matched_index: int | None,
+ min_prefix_hit_length: int = 1000,
media_regions: list[MediaRegion] | None = None,
) -> None:
if self.kv_prefix_cache is None:
@@ -438,7 +446,8 @@ class ExoBatchGenerator:
else 0.0
)
if matched_index is not None and (
- prefix_hit_length > 1000 or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+ prefix_hit_length >= min_prefix_hit_length
+ and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
):
self.kv_prefix_cache.update_kv_cache(
matched_index,
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 043a822f..bb89dc7e 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -55,6 +55,7 @@ from exo.worker.engines.mlx.utils_mlx import (
apply_chat_template,
fix_unmatched_think_end_tokens,
mx_barrier,
+ system_prompt_token_count,
)
from exo.worker.engines.mlx.vision import (
MediaRegion,
@@ -498,6 +499,7 @@ def mlx_generate(
# Encode prompt once at the top and fix unmatched think tags
all_prompt_tokens = encode_prompt(tokenizer, prompt)
all_prompt_tokens = fix_unmatched_think_end_tokens(all_prompt_tokens, tokenizer)
+ min_prefix_hit_length = max(1000, system_prompt_token_count(task, tokenizer))
vision: VisionResult | None = None
if vision_processor is not None:
@@ -714,8 +716,8 @@ def mlx_generate(
else 0.0
)
if matched_index is not None and (
- prefix_hit_length > 1000
- or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+ prefix_hit_length >= min_prefix_hit_length
+ and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
):
kv_prefix_cache.update_kv_cache(
matched_index,
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index e84a1934..790dcd8d 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -633,6 +633,29 @@ def apply_chat_template(
return prompt
+def system_prompt_token_count(
+ task_params: TextGenerationTaskParams,
+ tokenizer: TokenizerWrapper,
+) -> int:
+ """Approximate token count of the system prompt portion of the input."""
+ parts: list[str] = []
+ if task_params.chat_template_messages is not None:
+ for msg in task_params.chat_template_messages:
+ if msg.get("role") in ("system", "developer"):
+ content = msg.get("content", "") # type: ignore
+ if isinstance(content, str):
+ parts.append(content)
+ else:
+ if task_params.instructions:
+ parts.append(task_params.instructions)
+ for msg in task_params.input:
+ if msg.role in ("system", "developer"):
+ parts.append(msg.content)
+ if len(parts) == 0:
+ return 0
+ return len(tokenizer.encode(" ".join(parts), add_special_tokens=False))
+
+
def detect_thinking_prompt_suffix(prompt: str, tokenizer: TokenizerWrapper) -> bool:
"""
Detect if prompt ends with a thinking opening tag that should be
← 39c39e81 Integrations helpers (#1810)
·
back to Exo
·
Fix Nemotron cache leak upstream (#1819) d9ed9430 →