Add SnapKV-D post-prefill KV eviction

[pierre427]

What

Adds SnapKV-D, an opt-in post-prefill KV-cache eviction for long-context decode. After a prompt is prefilled, most middle rows contribute little to future attention; SnapKV-D keeps attention sinks + a recent window + the top observation-window-scored middle rows within a budget (SnapKV, arXiv:2404.14469) and evicts the rest, so the per-token KV read shrinks proportionally. Nothing runs unless a caller compacts a cache — default behavior is unchanged.

Why it needs a new cache

Retained rows are a sparse subset of the prompt, so RoPE position and physical row count must diverge. PositionPreservingKVCache tracks the true sequence position in offset (for future rotations) while storing only the retained rows, and records each row's true position so a prefix trim (prompt-cache reuse) stays exact and speculative rollback trims only the generated suffix. The retained keys keep their original rotation and decode queries use their true offset, so attention is correct against the compact cache.

API (all in mlx_lm/models/cache.py)

• snapkv_keep_indices(seq_len, budget, scores, *, sink_tokens=4, recent_tokens=None, min_tokens=128) — the retained positions. A prompt at/under min_tokens or within budget is kept whole.
• PositionPreservingKVCache(_BaseCache) — the compact, position-preserving cache (state/meta_state roundtrippable, trimmable, speculation-aware).
• evict_prompt_cache(prompt_cache, keep_indices, *, true_offset) — replaces full-attention KVCache layers with compact caches; other layer types untouched. Returns a SnapKVEvictionResult with retained/original counts and compact bytes.
• SnapKVAttentionCapture — scores a prefill by wrapping mx.fast.scaled_dot_product_attention, reducing only the observation-window query rows to a per-key vector so it never holds a prompt-sized attention matrix. The model output still uses the original fused kernel. The patch is process-global for the with block.
• compact_prompt_cache(model, prompt, *, budget, ...) — convenience that prefills, scores, and evicts in one call.

Usage:

from mlx_lm.models.cache import compact_prompt_cache
result = compact_prompt_cache(model, prompt_ids, budget=512)
# decode from result.cache at offset == len(prompt)

When it pays, and the honest caveats

• Long-context decode wins. Keeping a small fraction of prompt rows cuts the dominant per-token KV read. Lab measurements (10 Qwen-family variants × 4 scenarios) saw ~2–3× decode at long context and −2 to −4 GB peak, with a steady compact cache of ~84–102 MB.
• Wall-clock is neutral at short generations. The one-time scoring/prepare pass (~1.5 s on large models) only amortizes over long decodes / large contexts; short generations do not benefit.
• Capture-time peak. Scoring adds transient memory during prefill. The windowed scorer here (observation window + chunked reduction) bounds it; a dense scorer would not.
• Budget floor. There is a hard floor (min_tokens, default 128); single-reference facts degrade below a budget ratio. Quality is workload-level, not per-token lossless.

Because of these, it is default-off and explicit per call.

Tests

tests/test_snapkv_cache.py: the selection policy (sinks/recent/top-scored, no-op and validation paths); the position-preserving cache (offset/stored divergence, state/meta_state roundtrip, logical-prefix and speculative-suffix trims, growth, nbytes); eviction (offset preserved, non-KV layers untouched, byte reduction); and end-to-end capture + compaction + decode on a tiny attention model, including that the scoring hook actually fires and a short prompt is a no-op. black / isort --profile black clean.

This is the foundation of a small stack: a --kv-eviction snapkv server flag and a DuoAttention head-partitioned variant build on it.

[pierre427]

Pushed e8f737d: hardened the sink preservation in snapkv_keep_indices. It already staged sinks + a recency window before the score-ranked fill, but capped the sink floor at max(1, budget // 8) — so under a tight budget it kept fewer than the intended N sinks, which is exactly the attention-sink-eviction → repetition-collapse failure #1552 fixed for the deepseek_v32 DSA indexer.

The guarantee_sinks guard clamps the sink floor only by budget/seq_len, so the first N sinks + recency can never be evicted regardless of middle-row scores. Existing callers use budget >= 32 (cap >= 4), so their behaviour is unchanged; the guard only bites when budget // 8 < sink_tokens. Added 2 regression tests (sinks+recency survive budget=16 with every middle row scored 1e9; flag toggles the floor).