feat: prompt prefix caching with TTL eviction and TurboQuant support

[eloe]

Summary

Adds cross-request KV cache reuse to the server, significantly reducing time-to-first-token for multi-turn conversations and agentic workflows where the system prompt is stable across turns.

How it works

1. PromptCacheState tracks token IDs + KV cache between requests
2. On each request, find_prefix_length() computes the common prefix with the cached state
3. The KV cache is trimmed to the prefix length and reused, skipping re-computation of shared tokens
4. Cache entries are keyed per model (default) or per prompt_cache_key for multi-session routing

Features

• TurboQuant compatible: Uses cache.trim() instead of raw array slicing, works with QuantizedKVCache, RotatingKVCache, and ArraysCache (hybrid models like Qwen3.5)
• TTL eviction: Background asyncio task evicts stale entries after --prompt-cache-ttl seconds (default 300s), freeing GPU memory
• LRU cap: Max 64 cache entries with least-recently-used eviction to prevent unbounded memory growth
• Cache key routing: Optional prompt_cache_key field enables per-session cache isolation (useful for multi-user gateways)
• cached_tokens reporting: usage.input_tokens_details.cached_tokens in responses
• Graceful recovery: Cache mismatches (shape errors, stale state) are caught and fall back to fresh generation

Configuration

Flag	Env var	Default	Description
--prompt-cache-ttl	PROMPT_CACHE_TTL	300	Seconds before idle cache eviction. 0 = no expiry

Tests

9 tests covering cache state lifecycle, persistence, isolation, prefix matching, and cleanup:

python -m pytest mlx_vlm/tests/ -k "PromptCache" -v

Related

This builds on the same direction as #946 by @trevorgordon981 (basic PromptCacheState wiring) and #944 by @damonvjanis (TurboQuant trim fix). This PR includes both of those fixes plus TTL eviction, LRU capping, per-session cache key routing, cached_tokens reporting, graceful stale-cache recovery, and 9 tests.

[gofastercloud]

+1 from production use — this PR is doing real work for us. A huge thank-you to @eloe for pushing it upstream, and one bug report attached below that I think is worth folding in before merge.

Empirical result

We're running the mlx_vlm.server 0.4.4 + --kv-bits 3 --kv-quant-scheme turboquant configuration on an M4 Pro 24 GB Mac Mini, serving mlx-community/gemma-4-26b-a4b-it-4bit as the backend for a local AI agent stack with a ~21 k-token system prompt per turn. Before this PR, every agent turn re-prefilled the full system prompt at ~500 tok/s → ~42 s of prefill before any useful work. Five tool-call round-trips in a single task was 3.5 minutes of pure prefill, routinely blowing the agent's turn timeout.

With this PR applied (plus the TurboQuant shape-check fix described below), measured on two back-to-back identical agent turns through our gateway:

run	prompt size	duration	prefill rate
turn 1 (cold)	25 906 tok	144 708 ms (~2:24)	~500 tok/s
turn 2 (warm)	25 906 tok	5 949 ms (~5.9 s)	cache hit (no prefill)

~24× speedup on end-to-end agent turns. For direct /v1/chat/completions (non-streaming) the numbers are even cleaner: 13.9 s cold → 0.36 s warm on a 6 647-token prompt with the same prompt_cache_key. The prefill progress bars vanish from mlx-vlm.log on warm calls — the server skips straight to decode.

Without this PR, none of that is possible on our hardware. The agent-use case on a single-machine deployment becomes actually viable with prefix caching where before it was drowning in prefill time. This is exactly the kind of thing a server for local agent use needs, and I hope it lands.

Bug report — TurboQuant shape-check incompatibility

generate.py around L780 (post-patch), inside the cache validation that guards against stale KV state:

if reused_prefix_len > 0:
    try:
        for c in kwargs["prompt_cache"]:
            if hasattr(c, "keys") and c.keys is not None:
                expected_seq = reused_prefix_len
                actual_seq = c.keys.shape[2] if len(c.keys.shape) >= 3 else c.offset
                if actual_seq != expected_seq:
                    raise ValueError(
                        f"Cache shape mismatch: expected seq={expected_seq}, got {actual_seq}"
                    )

For unquantized KV caches this is correct: c.keys is (heads, kv_heads, seq_len, head_dim) and shape[2] equals the token count, so the comparison is meaningful.

For TurboQuant-quantized KV caches, however, c.keys has a grouped/chunked layout where shape[2] reflects the group size (typically 1024), not the actual token offset. The check therefore fails on every call for quantized caches — even when the cache is byte-identical to what was just stored — and falls into the invalidate + rebuild branch. The net effect is that prompt caching looks functionally correct on any prose test (the code paths all run) but delivers zero measured speedup on a TurboQuant-configured server, with logs like this on every call:

[prompt_cache] Cache validation failed, rebuilding: Cache shape mismatch: expected seq=6182, got 1024
[prompt_cache] Cache validation failed, rebuilding: Cache shape mismatch: expected seq=6200, got 1024
[prompt_cache] Cache validation failed, rebuilding: Cache shape mismatch: expected seq=6218, got 1024

The got 1024 is the dead giveaway — that's the quant group size, not the offset.

Suggested fix

Trust c.offset first when present (it's the authoritative counter for both unquantized and quantized caches under the existing MLX cache abstractions), and only fall back to shape-based comparison when offset is unavailable:

if reused_prefix_len > 0:
    try:
        for c in kwargs["prompt_cache"]:
            if hasattr(c, "keys") and c.keys is not None:
                expected_seq = reused_prefix_len
                if hasattr(c, "offset") and c.offset is not None:
                    actual_seq = c.offset
                else:
                    actual_seq = c.keys.shape[2] if len(c.keys.shape) >= 3 else 0
                if actual_seq != expected_seq:
                    raise ValueError(
                        f"Cache shape mismatch: expected seq={expected_seq}, got {actual_seq}"
                    )

That's the ~6-line change we're carrying locally. We went from 0 × speedup (with the upstream check) to 24 × speedup (with c.offset first) on an otherwise-identical setup. The non-TurboQuant path is unchanged because c.offset is already set on those caches too and points to the same value that shape[2] would have.

Happy to open a follow-up PR against your branch with the one-line fix if that's useful, or you might prefer to fold it in here since it's a one-character conceptual change and the commit already advertises TurboQuant support. Whatever's easier for you to land.

Context

We carry both this PR's full runtime diff and the TurboQuant shape-fix as local patches, vendored in-tree with a dedicated README and retirement plan. Happy to share the full apply-mlx-patches.sh harness if anyone else is running into the same TurboQuant issue — the patch file will be at docs/upstream-prs/03-mlx-vlm-prompt-cache-turboquant/ once we push it public.

Again — huge thanks for this PR. This is going to unblock a lot of single-machine agent stacks.