TL;DR
KV cache eviction permanently throws away tokens. KV cache merging corrupts the values of surviving tokens. ResKV (arXiv:2607.29591) does neither — it splits a fixed budget into a main cache that keeps exact tokens and a residual cache that reconstructs the attention contribution of evicted tokens. Both caches participate in the same softmax, so neither numerator nor denominator is distorted. Result: improvements across all 32 LongBench configurations and 63 of 64 RULER configurations vs. prior methods at identical budgets.
The Problem
Long-context LLM inference is memory-bound. A single 32K-token forward pass with Llama-3-8B can exhaust gigabytes of KV cache. The fix everyone reaches for is KV cache compression: keep only the "important" tokens, drop the rest.
The two mainstream approaches both have a fundamental flaw baked in.
Eviction methods (H2O, SnapKV, PyramidKV): drop low-importance tokens permanently. Clean and fast, but those dropped tokens had nonzero attention weights. Once dropped, their contribution to the softmax numerator and denominator is gone. Under tight budgets, performance degrades sharply on tasks requiring distributed evidence across the context window.
Merging methods: instead of dropping, fold evicted tokens into surviving ones. Information is preserved in aggregate, but the merge corrupts the value vectors of retained tokens and breaks the softmax denominator. The resulting attention distribution is biased in ways that compound across layers.
Both approaches accept that you cannot recover what you've compressed away. ResKV challenges that assumption.
How It Works
The core idea: split the KV budget $b$ into a main cache of size $m$ and a residual cache of size $r$:
$$b = m + r$$
The main cache stores exact KV pairs for the top-$m$ important tokens. The residual cache stores compact summaries of the evicted tokens' attention contributions — not the tokens themselves, but their aggregated effect on the softmax.
Shared-Softmax Integration
The key design decision: residual entries participate in the same softmax normalization as main-cache tokens. Standard attention is:
$$\text{Attn}(q, K, V) = \frac{\sum_{i \in \mathcal{M}} e^{q \cdot k_i} v_i + \sum_{j \in \mathcal{E}} e^{q \cdot k_j} v_j}{\sum_{i \in \mathcal{M}} e^{q \cdot k_i} + \sum_{j \in \mathcal{E}} e^{q \cdot k_j}}$$
Eviction methods zero out the $\mathcal{E}$ terms. ResKV approximates them with residual entries $(k_l^{\text{res}}, v_l^{\text{res}})$:
$$\text{Attn}{\text{ResKV}} = \frac{\sum{i \in \mathcal{M}} e^{q \cdot k_i} v_i + \sum_{l=1}^{r} e^{q \cdot k_l^{\text{res}}} v_l^{\text{res}}}{\sum_{i \in \mathcal{M}} e^{q \cdot k_i} + \sum_{l=1}^{r} e^{q \cdot k_l^{\text{res}}}}$$
Both numerator mass (value contribution) and denominator mass (normalization weight) of the evicted set are approximated and restored.
Validation Proxy
Not all layers and heads benefit equally from residual slots. ResKV uses a validation proxy at cache construction time to allocate $r$ unevenly across layers and heads — giving more residual slots to positions where the approximation quality is higher and the eviction loss is larger.
Dynamic Gate
At decode time, the query's attention pattern varies. When attention is sharp (concentrated on a few tokens), residual entries matter less — the main cache already dominates. When attention is diffuse (spread across many tokens), residual contributions are critical.
A dynamic gate $\alpha_q$ scales residual contributions per query based on attention sharpness over the main cache:
$$\alpha_q = \text{gate}(\text{sharpness}(q, \mathcal{M}))$$
High entropy in the main-cache attention distribution increases residual weight.
Show Me The Code
import torch
import torch.nn.functional as F
from dataclasses import dataclass
@dataclass
class ResKVConfig:
total_budget: int # b = m + r
residual_ratio: float # fraction of budget for residual (e.g. 0.2)
class ResKVCache:
def __init__(self, config: ResKVConfig, num_heads: int, head_dim: int):
self.m = int(config.total_budget * (1 - config.residual_ratio))
self.r = config.total_budget - self.m
self.main_k = self.main_v = None # (B, H, m, D) exact tokens
self.res_k = self.res_v = None # (B, H, r, D) eviction summaries
def update(self, new_k, new_v, importance_scores):
"""Prefill: select top-m tokens, cluster evicted into r residual slots."""
if self.main_k is None:
self.main_k, self.main_v = new_k, new_v
B, H, _, D = new_k.shape
self.res_k = torch.zeros(B, H, self.r, D, device=new_k.device, dtype=new_k.dtype)
self.res_v = torch.zeros_like(self.res_k)
return
all_k = torch.cat([self.main_k, new_k], dim=-2)
all_v = torch.cat([self.main_v, new_v], dim=-2)
if all_k.shape[-2] <= self.m:
self.main_k, self.main_v = all_k, all_v
return
# Keep top-m by importance
topk_idx = importance_scores.topk(self.m, dim=-1).indices
evict_mask = self._evict_mask(all_k.shape[-2], topk_idx)
self.main_k = all_k[..., ~evict_mask, :]
self.main_v = all_v[..., ~evict_mask, :]
# Compress evicted tokens into r residual slots
evicted_k = all_k[..., evict_mask, :]
evicted_v = all_v[..., evict_mask, :]
self.res_k, self.res_v = self._cluster(evicted_k, evicted_v)
def attend(self, query):
"""Shared-softmax attention over main + gated residual caches."""
gate = self._gate(query)
k = torch.cat([self.main_k, self.res_k * gate], dim=-2)
v = torch.cat([self.main_v, self.res_v * gate], dim=-2)
scale = query.shape[-1] ** -0.5
w = F.softmax((query @ k.transpose(-2, -1)) * scale, dim=-1)
return w @ v
def _gate(self, query):
"""Dynamic gate: high entropy -> higher residual contribution."""
scale = query.shape[-1] ** -0.5
logits = (query @ self.main_k.transpose(-2, -1)) * scale
probs = F.softmax(logits, dim=-1)
entropy = -(probs * probs.clamp(min=1e-9).log()).sum(dim=-1).mean()
return torch.sigmoid(entropy).item()
def _cluster(self, evicted_k, evicted_v):
"""Chunk-average evicted tokens into r residual slots."""
E = evicted_k.shape[-2]
chunk = max(1, E // self.r)
rk, rv = [], []
for i in range(self.r):
s, e = i * chunk, min((i + 1) * chunk, E)
if s >= E:
rk.append(torch.zeros(*evicted_k.shape[:-2], 1, evicted_k.shape[-1],
device=evicted_k.device, dtype=evicted_k.dtype))
rv.append(torch.zeros_like(rk[-1]))
else:
rk.append(evicted_k[..., s:e, :].mean(dim=-2, keepdim=True))
rv.append(evicted_v[..., s:e, :].mean(dim=-2, keepdim=True))
return torch.cat(rk, dim=-2), torch.cat(rv, dim=-2)
@staticmethod
def _evict_mask(seq_len, topk_idx):
mask = torch.ones(seq_len, dtype=torch.bool, device=topk_idx.device)
mask[topk_idx[0, 0]] = False # simplified; broadcast over batch/heads in practice
return mask
# Drop-in usage
cfg = ResKVConfig(total_budget=512, residual_ratio=0.2) # 410 main + 102 residual
cache = ResKVCache(cfg, num_heads=32, head_dim=128)
# During prefill
cache.update(layer_k, layer_v, importance_scores=cumulative_attn_weights)
# During decode
output = cache.attend(query) # numerator + denominator both restored
Benchmark Results
All comparisons use identical KV budgets. ResKV splits the budget $b = m + r$ while baselines use the full budget for exact tokens only.
LongBench (16 long-context tasks: single/multi-doc QA, summarization, code completion):
| Setting | H2O | SnapKV | PyramidKV | ResKV |
|---|---|---|---|---|
| All 32 configurations tested | baseline | baseline | baseline | best in all 32 |
Performance gap is largest at tight budgets (10-20% retention rate), where eviction methods lose the most information and ResKV's residual recovery provides the greatest advantage.
RULER (synthetic tasks: retrieval, multi-hop aggregation, chain tracing):
| Configurations improved | H2O | SnapKV | ResKV |
|---|---|---|---|
| Out of 64 tested | — | — | 63 / 64 |
RULER tasks explicitly require gathering distributed evidence across the context window — exactly the scenario where dropped token contributions hurt most. ResKV's residual cache directly addresses this failure mode.
Efficiency: Negligible memory overhead (residual slots come from the same budget $b$, not additional allocation). Decode throughput is stable at extended context lengths.
Gotchas & Limitations
Residual cache goes stale during long generation. ResKV builds the residual cache at prefill and freezes it. Tokens generated during decoding are added to the main cache (with further eviction), but the residual cache doesn't update. For very long generations, early evicted tokens' summaries may become stale.
Clustering quality limits recovery. Evicted tokens are summarized into $r$ slots by chunked averaging or k-means. If evicted tokens are semantically diverse, a single cluster centroid poorly represents all of them. Recovery is approximate by design, but heterogeneous eviction sets degrade it further.
No flash-attention kernel integration. The shared-softmax over concatenated main + residual caches doesn't map to standard flash-attention implementations. Current inference requires a custom attention kernel or a two-pass approach, which adds latency.
Layer/head allocation is fixed at prefill. The validation proxy assigns residual slots per layer and head once at cache construction time. It doesn't adapt as the decode proceeds and actual attention patterns evolve.
Try It Today
Paper: arXiv:2607.29591
ResKV is a drop-in upgrade for any eviction-based KV cache. To apply it:
- Define your total budget $b$ and residual ratio (start with 10-20% residual).
- At prefill, replace standard eviction with ResKV's
update()— cluster evicted tokens into residual slots instead of discarding. - At decode, use
attend()for shared-softmax over main + residual. - Optionally enable the validation proxy to allocate residual slots non-uniformly across layers.
The biggest wins will be in tasks with sparse, distributed evidence across long contexts: RAG over large corpora, multi-document reasoning, long-form code context.
What's your experience with KV cache compression in production? Drop a comment.
Sources
- arXiv paper: https://arxiv.org/abs/2607.29591
- arXiv HTML: https://arxiv.org/html/2607.29591
Top comments (0)