Skip to content

Add cache resume optimization for identical prompts#569

Closed
blightbow wants to merge 2 commits into
ml-explore:mainfrom
blightbow:cache-resume
Closed

Add cache resume optimization for identical prompts#569
blightbow wants to merge 2 commits into
ml-explore:mainfrom
blightbow:cache-resume

Conversation

@blightbow

Copy link
Copy Markdown

Add support for skipping prompt prefill when the KV cache already
contains the full prompt from a previous identical request. This
provides 2.7-15x speedup for retry scenarios (e.g., chat "swiping").

Implementation:

  • Add can_resume_from_cache() to validate cache state
  • Add resume_from_cache parameter to generate_step()
  • Skip prefill loop when cache.offset >= prompt_length
  • Maintain all integrity constraints (RoPE, masks, slicing)

The optimization is backwards compatible (defaults to False) and
does not manipulate cache.offset directly, avoiding potential
integrity issues with positional embeddings and attention masks.

Add comprehensive unit tests for the cache resume functionality:

  • test_can_resume_from_cache(): Tests the validation function with
    various cache states (empty, sufficient, insufficient tokens)
  • test_resume_from_cache_with_generate(): Tests end-to-end behavior
    with actual generation, verifying identical token outputs

The tests verify that:

  1. can_resume_from_cache() correctly identifies when cache is valid
  2. Cache resume produces identical generation outputs
  3. Cache offset increases correctly during resumed generation
  4. The optimization works with both KVCache and RotatingKVCache

Tested with Qwen3-235B model: 2.048s → 0.749s (2.73x speedup)
on 7-token prompt. Expected 6-15x on longer prompts (100-500 tokens).

Add support for skipping prompt prefill when the KV cache already
contains the full prompt from a previous identical request. This
provides 2.7-15x speedup for retry scenarios (e.g., chat "swiping").

Implementation:
- Add can_resume_from_cache() to validate cache state
- Add resume_from_cache parameter to generate_step()
- Skip prefill loop when cache.offset >= prompt_length
- Maintain all integrity constraints (RoPE, masks, slicing)

The optimization is backwards compatible (defaults to False) and
does not manipulate cache.offset directly, avoiding potential
integrity issues with positional embeddings and attention masks.

Tested with Qwen3-235B model: 2.048s → 0.749s (2.73x speedup)
on 7-token prompt. Expected 6-15x on longer prompts (100-500 tokens).
Add comprehensive unit tests for the cache resume functionality:

- test_can_resume_from_cache(): Tests the validation function with
  various cache states (empty, sufficient, insufficient tokens)
- test_resume_from_cache_with_generate(): Tests end-to-end behavior
  with actual generation, verifying identical token outputs

The tests verify that:
1. can_resume_from_cache() correctly identifies when cache is valid
2. Cache resume produces identical generation outputs
3. Cache offset increases correctly during resumed generation
4. The optimization works with both KVCache and RotatingKVCache

All 18 prompt_cache tests pass successfully.
@awni

awni commented Dec 11, 2025

Copy link
Copy Markdown
Member

Closing as we have more sophisticated cache reuse in mlx_lm.server as of #625. Also this sort of logic should be handled outside the core generation to avoid complexity / unnecessary coupling.

@awni awni closed this Dec 11, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants