Add cache resume optimization for identical prompts

mlx_lm/generate.py

@@ -309,6 +309,7 @@ def generate_step(
     quantized_kv_start: int = 0,
     prompt_progress_callback: Optional[Callable[[int], int]] = None,
     input_embeddings: Optional[mx.array] = None,
+    resume_from_cache: bool = False,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
     A generator producing token ids based on the given prompt from the model.
@@ -337,6 +338,10 @@ def generate_step(
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
+        resume_from_cache (bool): If True and prompt_cache contains cache.offset >=
+          len(prompt), skip the prompt prefill phase and resume generation directly.
+          This optimization is safe when the cache already contains the full prompt
+          from a previous identical request. Default: ``False``.
 
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
@@ -416,6 +421,17 @@ def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None):
             len(input_embeddings) if input_embeddings is not None else len(prompt)
         )
         prompt_processed_tokens = 0
+
+        # Check if we can skip prefill by resuming from cache
+        if resume_from_cache and cache.can_resume_from_cache(prompt_cache, total_prompt_tokens):
+            cache_offset = prompt_cache[0].offset if prompt_cache else 0
+            # Cache already contains the full prompt, skip prefill entirely
+            prompt_processed_tokens = total_prompt_tokens
+            # Keep only the last token for the initial generation step
+            prompt = prompt[-1:]
+            if input_embeddings is not None:
+                input_embeddings = input_embeddings[-1:]
+
         prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
         while total_prompt_tokens - prompt_processed_tokens > 1:
             n_to_process = min(prefill_step_size, prompt.size - 1)

mlx_lm/models/cache.py

@@ -109,6 +109,30 @@ def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
     return [c.trim(num_tokens) for c in cache][0]
 
 
+def can_resume_from_cache(cache: List[Any], prompt_length: int) -> bool:
+    """
+    Check if the cache already contains at least prompt_length tokens.
+
+    This function determines if prompt processing can be skipped because
+    the cache already contains a valid prefilled state for the prompt.
+
+    Args:
+        cache (List[Any]): The model's cache.
+        prompt_length (int): The number of prompt tokens.
+
+    Returns:
+        bool: True if cache.offset >= prompt_length, meaning the cache
+            already contains the full prompt and prefill can be skipped.
+    """
+    if not cache or len(cache) == 0:
+        return False
+
+    # Check the offset of the first cache layer (all layers should have same offset)
+    cache_offset = getattr(cache[0], 'offset', 0)
+
+    return cache_offset >= prompt_length
+
+
 def create_attention_mask(
     N: int, offset: int, return_array: bool, window_size: Optional[int]
 ):

tests/test_prompt_cache.py

@@ -18,6 +18,7 @@
     MambaCache,
     QuantizedKVCache,
     RotatingKVCache,
+    can_resume_from_cache,
     load_prompt_cache,
     make_prompt_cache,
     save_prompt_cache,
@@ -568,6 +569,72 @@ def test_rotating_cache_updates(self):
         self.assertEqual(k.shape[2], 10)
         self.assertEqual(v.shape[2], 10)
 
+    def test_can_resume_from_cache(self):
+        # Test with empty cache
+        self.assertFalse(can_resume_from_cache([], 10))
+        self.assertFalse(can_resume_from_cache(None, 10))
+
+        # Test with cache that has sufficient tokens
+        cache = [KVCache() for _ in range(2)]
+        for c in cache:
+            x = mx.random.uniform(shape=(1, 8, 10, 4))
+            c.update_and_fetch(x, x)
+
+        self.assertTrue(can_resume_from_cache(cache, 10))  # exact match
+        self.assertTrue(can_resume_from_cache(cache, 5))   # cache has more
+        self.assertFalse(can_resume_from_cache(cache, 15)) # cache has less
+
+        # Test with RotatingKVCache
+        cache = [RotatingKVCache(max_size=8) for _ in range(2)]
+        for c in cache:
+            x = mx.random.uniform(shape=(1, 8, 12, 4))
+            c.update_and_fetch(x, x)
+
+        self.assertTrue(can_resume_from_cache(cache, 10))
+        self.assertFalse(can_resume_from_cache(cache, 15))
+
+    def test_resume_from_cache_with_generate(self):
+        """Test that cache resume optimization produces identical token outputs."""
+        model, tokenizer = load(HF_MODEL_PATH)
+        prompt = tokenizer.encode("this is a test prompt", return_tensors="mlx")[0]
+
+        # First generation with cache but no resume
+        prompt_cache = make_prompt_cache(model)
+        results1 = list(
+            generate_step(prompt, model, prompt_cache=prompt_cache, max_tokens=5)
+        )
+        toks1, _ = zip(*results1)
+
+        # Verify cache has sufficient tokens
+        self.assertTrue(can_resume_from_cache(prompt_cache, len(prompt)))
+
+        # Trim the generated tokens to simulate identical request (swipe behavior)
+        trim_prompt_cache(prompt_cache, 5)
+
+        # Second generation with cache resume enabled
+        # This should skip prefill and start generating immediately
+        results2 = list(
+            generate_step(
+                prompt,
+                model,
+                prompt_cache=prompt_cache,
+                max_tokens=5,
+                resume_from_cache=True,
+            )
+        )
+        toks2, _ = zip(*results2)
+
+        # Verify generated tokens are identical
+        # This is the key test: cache resume should produce same outputs
+        self.assertEqual(toks1, toks2)
+
+        # Verify cache offset increased (resume processes last prompt token + new tokens)
+        # After first generation: offset = len(prompt) + 5
+        # After trim: offset = len(prompt)
+        # After resume: offset = len(prompt) + 1 (last token) + 5 (new tokens)
+        # The implementation keeps the last token for generation start
+        self.assertGreater(prompt_cache[0].offset, len(prompt))
+
 
 if __name__ == "__main__":
     unittest.main()