diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 6ca8a73fb..f85215314 100644 --- a/mlx_lm/generate.py +++ b/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) diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index 17354a02c..54f8711be 100644 --- a/mlx_lm/models/cache.py +++ b/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] ): diff --git a/tests/test_prompt_cache.py b/tests/test_prompt_cache.py index 287965e12..69fb37873 100644 --- a/tests/test_prompt_cache.py +++ b/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()