KV Cache Eviction Policies: Balancing Memory and Context

KV Cache Eviction Policies: Balancing Memory and Context

The Key-Value (KV) cache is one of the most important architectural optimizations in modern transformer-based large language models. During autoregressive generation, each new token requires attention over all previous tokens. Recomputing the keys and values for every prior token at every step would be prohibitively expensive. Instead, models cache these tensors so that only the current token's keys and values need to be computed at each step. However, this cache grows linearly with sequence length, and for long contexts or large batch sizes, it quickly becomes the dominant consumer of GPU memory. KV cache eviction policies are the strategies that decide which entries to keep and which to discard when memory pressure forces a choice.

What Is the KV Cache?

In a transformer decoder, the self-attention mechanism computes queries, keys, and values for every token. During generation, the keys and values for previously processed tokens do not change. The KV cache stores these tensors across layers and heads so they can be reused. For a model with L layers, H heads, a head dimension of d, and a sequence length of n, the cache stores approximately 2 * L * H * d * n floating point values per sequence. For a 70B parameter model processing 32K tokens, this can easily exceed 40 GB of memory per sequence.

The fundamental tension is simple: a larger cache means longer effective context and better generation quality, but it also means higher memory usage, lower batch sizes, and slower inference. Eviction policies attempt to resolve this tension by keeping only the most useful entries.

Why Eviction Policies Matter

Without eviction, you face a hard ceiling: once the cache is full, you either stop accepting longer contexts or you crash with out-of-memory errors. Eviction policies let you support effectively longer contexts than would otherwise fit in memory by discarding entries that are unlikely to be referenced again. This matters for several reasons:

• Throughput: Smaller per-sequence cache footprints allow larger batch sizes, directly increasing tokens per second.
• Long-context applications: Document QA, code assistants, and agents all benefit from long context windows, but memory limits make naive caching impractical.
• Cost: Memory is the most expensive resource in inference clusters. Reducing cache size reduces the hardware required to serve a given workload.
• Quality preservation: A well-designed eviction policy can discard 50% or more of cached entries with minimal impact on output quality, while a poor policy degrades the model rapidly.

Common Eviction Strategies

Several families of eviction policies have emerged, each with different trade-offs between implementation complexity, memory savings, and quality retention.

Window-Based Eviction

The simplest approach is a sliding window: keep only the most recent K tokens and discard everything older. This is essentially what models like Mistral use natively. It is trivial to implement and has zero computational overhead beyond managing a circular buffer, but it loses all information beyond the window, which can cause the model to "forget" early instructions or facts.

Attention-Score-Based Eviction

Since attention scores directly measure how much each cached token contributes to the current output, they are a natural signal for eviction. Policies like Heavy-Hitter Oracle (H2O) observe that attention is highly skewed: a small fraction of tokens receive most of the attention weight. By keeping the top-k tokens by cumulative attention score and evicting the rest, you can preserve quality while drastically reducing cache size.

Recency-Weighted Hybrid

Pure attention-score policies can evict tokens that are important but rarely attended to (such as system prompts). Hybrid policies combine attention scores with a recency bonus, ensuring that recent tokens are protected while low-attention older tokens are evicted. StreamingLLM popularized the idea of keeping "attention sinks" — the first few tokens, which absorb disproportionate attention — alongside a sliding window of recent tokens.

Token Merging and Compression

Rather than discarding tokens outright, some approaches merge similar tokens or project groups of tokens into a smaller representation. This preserves more information per unit of memory but adds computational cost and implementation complexity.

Implementing a Window-Based KV Cache

Let's start with a minimal sliding-window cache implementation in PyTorch. This demonstrates the core mechanics before we add attention-aware eviction.

import torch
import torch.nn as nn

class WindowedKVCache:
    def __init__(self, window_size: int, num_layers: int, num_heads: int, head_dim: int):
        self.window_size = window_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.keys = [None] * num_layers
        self.values = [None] * num_layers
        self.position = 0  # tracks absolute position for positional embeddings

    def update(self, layer_idx: int, new_keys: torch.Tensor, new_values: torch.Tensor):
        """
        new_keys/new_values shape: (batch, num_heads, seq_len, head_dim)
        """
        if self.keys[layer_idx] is None:
            self.keys[layer_idx] = new_keys
            self.values[layer_idx] = new_values
        else:
            self.keys[layer_idx] = torch.cat([self.keys[layer_idx], new_keys], dim=2)
            self.values[layer_idx] = torch.cat([self.values[layer_idx], new_values], dim=2)

        # Enforce window by slicing the oldest entries
        if self.keys[layer_idx].shape[2] > self.window_size:
            self.keys[layer_idx] = self.keys[layer_idx][:, :, -self.window_size:, :]
            self.values[layer_idx] = self.values[layer_idx][:, :, -self.window_size:, :]

        self.position += new_keys.shape[2]
        return self.keys[layer_idx], self.values[layer_idx]

    def reset(self):
        self.keys = [None] * self.num_layers
        self.values = [None] * self.num_layers
        self.position = 0

This implementation is intentionally simple. In production, you would pre-allocate a fixed buffer and use circular indexing to avoid the cost of torch.cat and slicing on every step. But the logic above captures the essential idea: append new entries, then trim to the window size.

Implementing Attention-Score-Based Eviction

Now let's implement a more sophisticated policy that uses attention scores to decide which tokens to evict. We will track a running sum of attention weights per cached position and evict the lowest-scoring positions when the cache exceeds a budget.

import torch

class AttentionBasedKVCache:
    def __init__(self, max_size: int, num_layers: int, num_heads: int, head_dim: int,
                 keep_recent: int = 64, keep_initial: int = 4):
        """
        max_size: maximum number of cached tokens
        keep_recent: always preserve the most recent N tokens
        keep_initial: always preserve the first N tokens (attention sinks)
        """
        self.max_size = max_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.keep_recent = keep_recent
        self.keep_initial = keep_initial

        self.keys = [None] * num_layers
        self.values = [None] * num_layers
        self.attention_scores = [None] * num_layers  # (batch, num_heads, cached_len)

    def update(self, layer_idx: int, new_keys, new_values, attn_weights=None):
        """
        attn_weights: (batch, num_heads, new_seq_len, cached_len)
        If provided, accumulate attention scores for eviction decisions.
        """
        if self.keys[layer_idx] is None:
            self.keys[layer_idx] = new_keys
            self.values[layer_idx] = new_values
            self.attention_scores[layer_idx] = torch.zeros(
                new_keys.shape[0], self.num_heads, new_keys.shape[2],
                device=new_keys.device, dtype=torch.float32
            )
        else:
            self.keys[layer_idx] = torch.cat([self.keys[layer_idx], new_keys], dim=2)
            self.values[layer_idx] = torch.cat([self.values[layer_idx], new_values], dim=2)
            new_scores = torch.zeros(
                new_keys.shape[0], self.num_heads, new_keys.shape[2],
                device=new_keys.device, dtype=torch.float32
            )
            self.attention_scores[layer_idx] = torch.cat(
                [self.attention_scores[layer_idx], new_scores], dim=2
            )

        # Accumulate attention weights into the score tracker
        if attn_weights is not None:
            # Sum across query positions and average across heads
            contribution = attn_weights.sum(dim=2).mean(dim=1)  # (batch, cached_len)
            self.attention_scores[layer_idx] += contribution.unsqueeze(1)

        # Evict if over budget
        current_len = self.keys[layer_idx].shape[2]
        if current_len > self.max_size:
            self._evict(layer_idx, current_len - self.max_size)

        return self.keys[layer_idx], self.values[layer_idx]

    def _evict(self, layer_idx: int, num_to_evict: int):
        cached_len = self.keys[layer_idx].shape[2]
        # Build a mask of protected positions
        protected = torch.zeros(cached_len, dtype=torch.bool, device=self.keys[layer_idx].device)
        protected[:self.keep_initial] = True
        protected[-self.keep_recent:] = True

        # Average attention scores across heads and batch
        scores = self.attention_scores[layer_idx].mean(dim=(0, 1))  # (cached_len,)
        scores[protected] = float('inf')  # never evict protected tokens

        # Find the lowest-scoring unprotected positions
        evict_indices = torch.argsort(scores)[:num_to_evict]
        keep_mask = torch.ones(cached_len, dtype=torch.bool, device=scores.device)
        keep_mask[evict_indices] = False

        self.keys[layer_idx] = self.keys[layer_idx][:, :, keep_mask, :]
        self.values[layer_idx] = self.values[layer_idx][:, :, keep_mask, :]
        self.attention_scores[layer_idx] = self.attention_scores[layer_idx][:, :, keep_mask]

This implementation demonstrates the key ideas behind H2O-style eviction: accumulate attention scores, protect critical tokens (attention sinks and recent context), and evict the lowest-scoring positions. In a real deployment, you would also need to handle the position IDs carefully, since evicting tokens creates gaps in the sequence that positional embeddings must account for.

Integrating Eviction Into a Generation Loop

The cache must be integrated into the actual attention computation. Here is a simplified generation loop that uses the attention-based cache:

import torch
import torch.nn.functional as F

def generate_with_eviction(model, input_ids, max_new_tokens, cache, tokenizer):
    device = input_ids.device
    generated = input_ids.clone()

    # Prefill phase: process the entire prompt
    with torch.no_grad():
        outputs = model(input_ids, use_cache=False)
        logits = outputs.logits[:, -1, :]
        # Store initial KV pairs from the model's internal layers
        # (In practice, you would hook into the model's attention layers)

    next_token = torch.argmax(logits, dim=-1, keepdim=True)
    generated = torch.cat([generated, next_token], dim=1)

    # Decode phase: one token at a time
    for step in range(max_new_tokens - 1):
        with torch.no_grad():
            # Compute attention with current cache
            # The model produces new keys/values and attention weights
            outputs = model(next_token, past_key_values=cache, output_attentions=True)
            logits = outputs.logits[:, -1, :]

            # Extract attention weights from the last layer for eviction scoring
            attn_weights = outputs.attentions[-1]  # (batch, heads, 1, cached_len)

            # Update cache with attention-aware eviction
            # (This would be done inside the model's forward pass in practice)
            for layer_idx in range(cache.num_layers):
                cache.update(
                    layer_idx,
                    outputs.past_key_values[layer_idx][0],
                    outputs.past_key_values[layer_idx][1],
                    attn_weights=attn_weights if layer_idx == cache.num_layers - 1 else None
                )

        next_token = torch.argmax(logits, dim=-1, keepdim=True)
        generated = torch.cat([generated, next_token], dim=1)

        if next_token.item() == tokenizer.eos_token_id:
            break

    return generated

Note that in a real implementation, the eviction logic typically lives inside the attention layer itself, not in the outer generation loop. Frameworks like Hugging Face Transformers and vLLM provide hooks for custom cache classes that integrate directly with the model's forward pass.

Best Practices

Choosing and tuning an eviction policy requires careful consideration of your specific workload. Here are practical guidelines:

• Always protect attention sinks. The first few tokens of a sequence receive disproportionate attention regardless of their semantic relevance. Evicting them causes catastrophic quality degradation. Keep at least the first 4 tokens permanently.
• Reserve a recent window. Even with attention-score-based eviction, always protect the most recent 32 to 128 tokens. The model relies heavily on immediate context for coherent generation.
• Profile before optimizing. Measure actual cache memory usage and attention distributions for your specific model and workload before choosing a policy. Different models have very different attention patterns.
• Use per-layer budgets. Early layers tend to have more uniform attention, while later layers are more concentrated. Consider different eviction budgets per layer rather than a single global setting.
• Validate quality rigorously. Eviction can silently degrade output quality in ways that are not obvious from perplexity alone. Use task-specific benchmarks and human evaluation, not just loss metrics.
• Consider the interaction with quantization. KV cache quantization (for example, to INT8 or FP8) reduces memory per token and can be combined with eviction. The two techniques are complementary, not alternatives.
• Pre-allocate buffers. Dynamic tensor resizing via torch.cat is expensive. In production, pre-allocate the maximum cache size and use index pointers to track valid entries.
• Handle position IDs correctly. When tokens are evicted, the remaining tokens may have non-contiguous position IDs. Some implementations renumber positions, while others use relative position encodings that are robust to gaps. Choose an approach compatible with your model's positional encoding scheme.

Choosing the Right Policy

The best eviction policy depends on your constraints. If your primary goal is maximum throughput and your contexts are short enough that a window covers the relevant information, a simple sliding window is hard to beat for its simplicity and speed. If you need long-context quality with limited memory, attention-score-based eviction with sink protection is the current state of the art. For the most demanding applications, token merging or learned compression policies can squeeze out additional quality at the cost of implementation complexity and compute overhead.

It is also worth noting that eviction is not always necessary. If your model and hardware can support the full cache for your target context length, the highest quality comes from keeping everything. Eviction is a tool for when memory is the binding constraint, which is the common case in production inference but not universal.

Conclusion

KV cache eviction policies are a critical tool for making long-context LLM inference practical under real-world memory constraints. By intelligently discarding cached entries that contribute least to the model's attention computation, you can dramatically reduce memory usage while preserving generation quality. The key is to balance recency, attention importance, and the protection of structural tokens like attention sinks. Start with a simple sliding window if you need a quick win, then move to attention-score-based eviction with sink protection when you need to push context lengths further. Whatever policy you choose, always validate quality on your actual workloads and profile memory usage end to end, because the interaction between eviction, batching, quantization, and model architecture is subtle and workload-dependent.