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>
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
Exposes SnapKV-D (#1518) as an opt-in server policy:
--kv-eviction snapkvcompacts each full-attention layer to a KV budget after a fresh prefill, cutting long-context decode KV reads.How
--kv-eviction {none,snapkv}(defaultnone),--kv-budget(default 512),--kv-window(24),--kv-sink-tokens(4),--kv-min-tokens(128).--kv-min-tokens, the single-request path prefills the prompt under aSnapKVAttentionCapture, compacts every full-attention layer to the budget viaevict_prompt_cache, and decodes from the last prompt token. Partial prefix-cache hits and short prompts are left uncompacted so the scored/retained positions stay exact.Scope
The eviction runs on the single-request path by design, not as a temporary limitation. SnapKV-D's benefit is reading physically fewer KV rows per decode step, and it compacts each request to a different retained-row set. mlx-lm's batched decode reads one shared
(B, H, T, d)KV tensor with a uniform row count across the batch, so there is no way to give one request 512 rows and another 2000 in the same batch — masking evicted rows would not reduce the (memory-bound) read. So SnapKV is inherently a per-request-decode optimization and routes to the single-request path. Default-off, so no behavior changes unless the flag is set.Test
tests/test_server_kv_eviction.pydrivesResponseGenerator._maybe_snapkv_compactover a tiny attention model with a stubbed args namespace: eviction off is a passthrough; a fresh long prompt is compacted to the budget with the true offset preserved and then decodes; short prompts (≤min_tokens) and partial prefix-cache hits are skipped.black/isort --profile blackclean.