Add checkpoint-based prefix caching for hybrid/recurrent models

[d0lphin]

Summary

• Adds RecurrentCheckpointStore, an LRU-bounded store that saves deep-copied cache snapshots at 512-token intervals during prefill
• Integrates checkpoint save/restore into both CacheWrapper (sequential) and BatchedModelKit (batched) code paths
• Zero overhead for standard transformer models — the store is only created when can_trim_prompt_cache() returns False

Problem

Hybrid models like Qwen3-Coder-Next mix KVCache (trimmable attention) with ArraysCache (non-trimmable recurrent/DeltaNet layers). Since can_trim_prompt_cache() requires all layers to be trimmable, it returns False for the entire model — completely disabling prefix caching.

This means every request with a diverging suffix (the most common real-world pattern: same system prompt, different user message) reprefills the entire context from scratch. On a ~13.5k token context, this produces ~5.9s warm TTFT instead of the expected ~200ms.

The LRUPromptCache in batched mode handles exact matches and prefix extensions correctly, but the "diverging suffix" path requires can_trim_prompt_cache() to return True — which it never does for hybrid models.

Why you can't just trim recurrent state

For KV cache (attention layers): trimming to position N means "keep K/V tensors up to N, discard the rest."

For ArraysCache (recurrent/DeltaNet layers): the hidden state at position N encodes the entire sequence up to N. You can't "un-process" tokens — you need an explicit checkpoint saved at position N.

Solution

Bypass trim_prompt_cache() entirely by implementing checkpoint-based prefix reuse at the mlx-engine level:

1. During prefill, process the prompt in chunks (512 tokens), saving a copy.deepcopy() + mx.eval() snapshot of the cache after each chunk
2. On subsequent requests, find the longest checkpoint matching the prompt prefix, restore it, and only process the remaining tokens
3. LRU eviction bounds memory usage (8 checkpoints for sequential, 16 for batched)

This approach requires no changes to mlx-lm — the checkpoint store sits alongside the existing caching infrastructure.

Inspiration

This is directly inspired by llama.cpp's approach to the same problem:

• PR #13979 — Hybrid recurrent cache foundation (Jun 2025)
• PR #16382 — Context checkpointing for hybrid/recurrent models (Oct 2025)

llama.cpp saves recurrent state snapshots at strategic positions and restores them when a new request shares a prefix, instead of reprocessing from scratch.

Benchmark Results

Measured on Qwen3-Coder-Next-4bit (Apple Silicon), ~950-token shared system prompt:

Scenario	TTFT	vs Cold
Cold (no cache)	1.19s	—
Warm prefix (checkpoint restore at 512, diverging suffix)	0.39s	3.0x faster
Exact match (full checkpoint restore)	0.09s	13.4x faster

Changes

File	Change
mlx_engine/recurrent_checkpoint_store.py	New: LRU-bounded checkpoint store with save(), find_longest_prefix(), clear()
mlx_engine/cache_wrapper.py	Checkpoint detection, save during prefill, restore on prefix mismatch, clear on draft model changes
mlx_engine/model_kit/batched_model_kit.py	Checkpoint detection, _prefill_with_checkpoints(), _fetch_cache_with_checkpoint_fallback()
tests/test_recurrent_checkpoint_store.py	11 unit tests for the checkpoint store
tests/test_cache_wrapper.py	7 new tests for checkpoint integration (4 mock-based, 3 model-dependent)

Test plan

• All 15 mock-based tests pass (test_recurrent_checkpoint_store.py + TestCacheWrapperCheckpointing)
• Manual end-to-end test with Qwen3-Coder-Next-4bit: 3 sequential generations confirming checkpoint save/restore and TTFT improvement
• Manual end-to-end test with batched_demo.py --parallel 2: confirmed checkpoint restore across concurrent requests
• Standard transformer models unaffected (checkpoint store is None, zero overhead)

🤖 Generated with Claude Code

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a8602c8a98

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Preserve at least one token on checkpoint restore

When restoring from a checkpoint, the code returns prompt_tokens[prefix_len:] without applying the num_tokens_to_exclude adjustment that _find_common_prefix enforces. If the checkpoint matches the entire prompt (or leaves fewer than num_tokens_to_exclude tokens), update_cache later sees an empty prompt_tokens, sets num_tokens_to_exclude to 0, and returns an empty array. Downstream generation in the text-only path expects at least one token (the code normally keeps the last token out of the cache), so exact-prefix hits can yield zero-length prompts and break generation. Consider reducing prefix_len so that at least num_tokens_to_exclude tokens remain outside the cache before returning.

Useful? React with 👍 / 👎.

[d0lphin]

Good catch — this is a valid edge case. It triggers when a new prompt exactly matches a checkpoint key in length (e.g. request 1 is 100 tokens → checkpoint saved at key tokens[:99], then request 2 is exactly 99 tokens matching that prefix). The checkpoint restore returns zero remaining tokens and generation breaks.

Fixed in 123296a: both checkpoint restore paths in _get_unprocessed_tokens now truncate the query to find_longest_prefix by num_tokens_to_exclude, mirroring the same guard that _find_common_prefix already applies (line 108–113). Since recurrent state can't be trimmed, too-long checkpoints must be excluded from the search entirely.

max_prefix = len(prompt_tokens) - num_tokens_to_exclude
result = self._checkpoint_store.find_longest_prefix(prompt_tokens[:max_prefix])

Note: the batched path (_fetch_cache_with_checkpoint_fallback) was already safe — it always returns prompt_tokens[-1:].

[d0lphin]

I have read the CLA Document and I hereby sign the CLA

[d0lphin]

recheck

[d0lphin]

Codex comment addressed. Tested in LM Studio as well, prefix cache works, however we are serializing now everything which messes up concurrency completely. Fix incoming.