Only update KV prefix cache on a good cache hit

Author: rltakashige

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -40,7 +40,10 @@
     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 @@ def submit(
                 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 @@ def _save_prefix_cache(
         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 @@ def _save_prefix_cache(
                 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,

src/exo/worker/engines/mlx/generator/generate.py

@@ -55,6 +55,7 @@
     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,

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