Add prompt checkpoints for hybrid caches

[stale2000]

Why this matters

Modern hybrid models combine attention KV cache with recurrent state (for example, Mamba or Gated DeltaNet layers). A recurrent state at token T summarizes every preceding token, so it cannot be safely trimmed back to an arbitrary branch point the way attention KV can.

Without an earlier exact state, a request that diverges from a cached prompt must prefill from the beginning. This is most visible in multi-turn chat, tool-use, and agent workloads that share a long system prompt or conversation prefix.

Solution

For a cache containing any non-trimmable layer, split each prompt segment at the existing prefill_step_size boundary. The server already saves a cache at each segment boundary and already knows how to select the nearest valid shorter prefix. This change gives it exact recurrent-state checkpoints to select:

1. Save complete hybrid state every prefill chunk.
2. On a diverged prompt, restore the newest checkpoint fully inside the shared prefix.
3. Replay only tokens after that checkpoint.

No recurrent state is trimmed, reset, or combined with KV state from a different history. Fully trimmable KV-only models keep their existing cache behavior.

Benefit

With the default 2,048-token prefill chunk, a branch can reuse the checkpoint immediately before its divergence and reprocess at most one chunk before the changed suffix. For example, a change after 30K tokens in a 32K-token prompt can reuse the 28,672-token checkpoint and prefill roughly 3.3K tokens instead of 32K: about 90% fewer prompt tokens processed for that request.

The exact saving depends on where the branch falls and cache capacity. The important guarantee is correctness: every restored recurrent state exactly matches the token prefix it represents.

Scope

• Applies only to server batching with a cache that contains a non-trimmable layer.
• Uses the existing prefill_step_size as the checkpoint interval; no new option or cache format.
• Keeps the existing LRU cache limits and cache-type eviction policy.

Tests

• python -m py_compile mlx_lm/server.py tests/test_server.py
• git diff --check
• Added coverage for checkpoint segmentation and for selecting a saved non-trimmable prefix after a branch.

The runtime MLX suite requires Apple Silicon.