mlxcache — CacheBlend-style KV-cache fusion for MLX-LM

Reuse the KV cache of RAG text chunks without paying for a full prefill.

When you stitch several retrieved passages in front of a query, ordinary prefix caching only helps the first chunk — every later chunk's true KV depends on cross-attention to the chunks before it, which per-chunk caching never computed. Naively concatenating independently-cached KV therefore gives wrong attention and degrades quality.

mlxcache ports CacheBlend to MLX: it precomputes each chunk's KV independently, then selectively recomputes only the small fraction of tokens whose KV actually moves once they can see the other chunks (the High KV-Deviation tokens). That recovers full-prefill behaviour at a fraction of the attention cost.

It is an adapter for mlx-lm: it drives the loaded model's own modules, so its forward is numerically identical to model(ids) (verified in the tests). No monkeypatching; decoding hands back to a stock mlx_lm KVCache.

Install

cd mlxcache && pip install -e .

Requires mlx, mlx-lm. Supports the Llama / Mistral / Qwen2 / Qwen3 decoder family (dense attention).

Usage

from mlx_lm import load
from mlxcache import blend_generate

model, tok = load("mlx-community/Qwen3-0.6B-bf16")
enc = lambda s: tok.encode(s, add_special_tokens=False)

# chunks are token-id sequences, concatenated in order:
#   [ system/prefix,  doc_0,  doc_1,  …,  query+suffix ]
chunks = [enc("Answer using the passages.\n\n"),
          enc("Passage A: The Eiffel Tower in Paris was completed in 1889.\n\n"),
          enc("Passage B: Mount Kilimanjaro is the highest mountain in Africa at 5895 m.\n\n"),
          enc("Question: What is the highest mountain in Africa?\nAnswer:")]

text = blend_generate(model, tok, chunks,
                      suffix_len=len(chunks[-1]),  # query is always recomputed
                      recomp_ratio=0.15,           # fraction of doc tokens refreshed
                      max_tokens=32)
print(text)

Lower level — get a decode-ready cache and the first-token logits:

from mlxcache import blend
cache, res = blend(model, chunks, suffix_len=len(chunks[-1]), recomp_ratio=0.15)
# `cache` is a list of mlx_lm KVCache (offset = full prompt length);
# continue with a plain `model(token, cache=cache)` loop, or `mlx_lm.generate`.
print(res.n_selected, "of", res.n_tokens, "tokens recomputed")

To amortize across queries (the real RAG win), collect once and reuse:

from mlxcache import BlendModel, collect, fuse_gradual
bm = BlendModel(model)
chunk_kv, _, N = collect(bm, chunks)          # do this once per document set
res = fuse_gradual(bm, all_ids, chunk_kv, suffix_len, recomp_ratio=0.15)  # per query

How it works

1. collect — prefill each chunk independently (chunk-local attention), keeping every layer's K/V. Each chunk is RoPE'd at its intended global start position, so the stored keys are already in the fused frame (and in the exact layout mlx_lm's KVCache uses, so decode consumes them as-is).
2. fuse — run the concatenated sequence. At one early check layer, score each token's KV deviation ‖V_fresh − V_chunk-local‖² — how much its value moves once it sees the other chunks. The top recomp_ratio document tokens (plus the always-recomputed query/suffix) are refreshed with their real cross-attended K/V; the rest stay chunk-local.
3. gradual filtering (method="gradual", default) — after the check layer, the non-selected tokens are dropped from the residual stream, so deep layers run on ~recomp_ratio of the tokens. A dropped token can't affect a survivor's context, so the result is identical to the dense path — it just skips the wasted compute.

fuse (dense) and fuse_gradual produce the same cache and logits; gradual is the fast one.

Results

Long-context speedup — Qwen3-4B

Per-query prefill speedup vs a full prefill, over ShareGPT long context (chunks pre-collected, 15% recompute), measured at 4B because the fixed per-layer fusion overhead (scatter / gather / mask) only amortizes against enough per-token compute. Reproduce with python benchmarks/long_context.py && python benchmarks/plot_long_context.py.

context	full prefill	gradual fuse	speedup	deep-layer tokens	KL vs full
1k	3.3 s	1.5 s	2.2×	36%	0.009
2k	8.4 s	2.9 s	2.8×	26%	0.120
4k	20.6 s	6.1 s	3.4×	20%	0.104
8k	51.0 s	15.1 s	3.4×	18%	0.010

The speedup grows with context as the per-layer overhead amortizes. collect costs ≈ one full prefill, so the win requires reusing it across many queries over the same documents (the RAG setting). Absolute times are from a memory-constrained machine — the ratio is the portable number.

Quality — Qwen3-0.6B

mean KL(full-prefill ‖ approx) of the next-token distribution (lower = behaves more like a true full prefill):

recompute ratio	mean KL vs full prefill
0% (naive concat)	0.52
15% (CacheBlend)	0.13
30%	0.088
100% (= full prefill)	~1e-10

Serving with mlx_lm.server

Run the stock OpenAI-compatible server with CacheBlend prefill for long prompts:

python -m mlxcache.server --model mlx-community/Qwen3-4B-4bit --port 8080 \
    --chunk-size 256 --recomp-ratio 0.15 --min-fuse-tokens 1024

A prompt of at least --min-fuse-tokens arriving against an empty cache is split into --chunk-size chunks, fused, and loaded into the cache; the server then streams normally. Shorter prompts, prefix-cache hits, and draft/speculative setups fall back to the stock prefill. All other flags pass through to mlx_lm.server. Single-request streaming path only (batching is force-disabled so the hook fires).

Experiments

• experiments/turboquant_fused — composes the fused cache with mlx-turboquant's rotated 4-bit TurboQuantKVCache. Fusion fixes cross-chunk quality, TurboQuant shrinks bytes, and they stack: fused · TQ-4bit+QJL reaches KL 0.077 vs full prefill at 3.2× less KV memory on ShareGPT. Plain 4-bit KV is lossy; QJL's unbiased residual correction is what makes it work. Full 3×3 grid + mechanics in the experiment README.
• experiments/energy — tokens per joule across plain / MLX 4-bit weights / TurboQuant KV / MLXCache / composition, measured with macmon (no-sudo SoC power). Different techniques save energy in different phases: weight-quant lifts decode (7.6 tok/J), fusion lifts end-to-end (3.4 tok/J) by cutting prefill energy 2.7×; TurboQuant KV trades energy for memory. Includes an honest fanless-M4 thermal caveat.

Status / not yet done

• fp16/bf16 core. The fused cache also composes with rotated 4-bit storage — see the TurboQuant experiment above.
• Dense attention only — sliding-window layers are rejected by assert_supported.
• Single prompt (B=1) during fusion; the mlx_lm.server integration therefore uses the single-request streaming path (batching force-disabled).
• No disk persistence for collect yet — it is recomputed per request in the server; a persistent per-document store is the obvious next step.

Tests

pytest tests/

test_forward_equiv pins the manual forward to model(ids); test_fusion covers RoPE framing, fuse == full prefill at ratio 1.0, gradual == dense (exact in fp32), and the quality monotonicity.