Conversation
Long-context decode reads the whole KV cache every step, but 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, cutting the per-token KV read proportionally. Opt-in; nothing runs unless a caller compacts a 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 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. - snapkv_keep_indices computes the retained positions (sinks + recent + top scored) for a budget; a prompt at/under min_tokens or within budget is kept whole. - evict_prompt_cache replaces full-attention KVCache layers with compact position-preserving caches, leaving other layer types untouched. - SnapKVAttentionCapture scores a prefill by wrapping mx.fast SDPA and reducing only the observation-window query rows, so it never holds a prompt-sized attention matrix. compact_prompt_cache ties prefill + scoring + eviction into one call. Tests cover the selection policy, the position-preserving cache (offset/stored divergence, state/meta_state roundtrip, logical-prefix and speculative-suffix trims, growth), eviction, and end-to-end capture + compaction + decode on a tiny attention model. black/isort clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Wire SnapKV-D post-prefill eviction into the server as an opt-in policy.
- New flags: --kv-eviction {none,snapkv} (default none), --kv-budget,
--kv-window, --kv-sink-tokens, --kv-min-tokens.
- On a fresh full prefill longer than the floor, the single-request path
prefills the prompt under a SnapKVAttentionCapture, compacts every
full-attention layer to the budget via evict_prompt_cache, and decodes
from the last prompt token. Partial prefix-cache hits and short prompts
are left uncompacted so retained positions stay exact.
- SnapKV scores a prefill with a process-global attention hook that cannot
separate per-request scores in a shared batch, so enabling it routes the
model to the single-request path.
Tested with a stubbed generator over a tiny attention model: off is a
passthrough, a fresh long prompt is compacted to the budget with the true
offset preserved and decodes, and short / partial-hit prompts are skipped.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Second SnapKV-D stage: HeadPartitionedKVCache stores the union of retained prompt positions once and exposes a per-head prefix mask so retrieval heads attend the full long context while streaming heads attend only their own sink+recent rows. offset still tracks the true sequence position for RoPE, and positions/head_positions make prefix trim (prompt-cache reuse) and speculative rollback exact. evict_prompt_cache_by_head applies a caller-supplied head_keep_indices, replacing plain KVCache layers with the partitioned cache and leaving every other layer type untouched. The head-classification policy (deciding which heads are retrieval vs streaming) is intentionally out of scope and stays application-side; this provides only the cache representation and the eviction op. Default-off: no model wires it up automatically. Provenance: DuoAttention (arXiv:2410.10819). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds a DuoAttention (arXiv:2410.10819) head-partitioned KV cache, a second eviction stage on top of SnapKV-D: different KV heads retain different token subsets — "retrieval" heads keep the long context, "streaming" heads keep only sinks + a recent window — so per-head sparsity cuts the KV read further than a single shared keep set.
How (all in
mlx_lm/models/cache.py)HeadPartitionedKVCache(_BaseCache)— stores the union of retained positions with a per-head keep mask, tracks the trueoffsetfor RoPE likePositionPreservingKVCache, and masks attention so each head only attends to its own retained rows. State/meta_state roundtrippable, trimmable, speculation-aware.evict_prompt_cache_by_head(prompt_cache, head_keep_indices, *, true_offset, query_heads=None)— replaces full-attentionKVCachelayers withHeadPartitionedKVCachegiven caller-supplied per-head keep indices; other layers untouched. Returns aHeadPartitionedEvictionResultwith union/per-head counts and compact bytes.Scope
This PR provides the cache class and the eviction op that applies a caller-supplied
head_keep_indices. The head classification (deciding which heads are retrieval vs streaming — typically a one-time offline pass per model) is deliberately out of scope; this is the runtime substrate it plugs into.A note on the storage model
Storage is a shared union of positions plus a per-head mask, so the memory saving comes from streaming heads narrowing the union, not from any single head. If one head retains the full context, the stored union is the full prompt regardless of what other heads drop — the win requires the retained heads themselves to be bounded.
Test
tests/test_snapkv_duoattention.py:HeadPartitionedKVCacheconstruction from per-head keep,size/offset/positions/head_positions, state/meta_state roundtrip,is_trimmable,trim,nbytes, per-headmake_masksemantics, and decode append;evict_prompt_cache_by_headwith asymmetric per-head keep (offset preserved, correct union count, non-KV layers untouched, real byte reduction); and an end-to-end prefill → evict-by-head → decode on a tiny attention model.black/isort --profile blackclean.