KV-Cache Optimization Techniques

Learning Objectives
Notation
Core Intuition
PagedAttention (vLLM)
Prefix Caching
KV-Cache Quantization
Token Eviction Strategies
Sliding Window Cache
Multi-Query and Grouped-Query Caching
Low-Rank KV Compression
Cache-Aware Scheduling
Common Pitfalls
Summary
Exercises

Learning Objectives

Explain memory fragmentation in KV-cache and how PagedAttention solves it.
Derive the memory savings from KV quantization (INT4, INT8).
Design token eviction policies based on attention scores.
Analyze the tradeoffs of sliding window vs full cache.
Explain how prefix caching enables efficient shared prompts.

Notation

$P$ — page size (number of tokens per page)
$N_{\text{pages}}$ — total pages allocated
$\text{CacheHit}(r_1, r_2)$ — fraction of shared prefix between requests
$\alpha_{\text{evict}}$ — attention score threshold for eviction

Core Intuition

The KV-cache is the dominant memory consumer during LLM serving (often exceeding model weights for long contexts). Naive allocation reserves maximum-length contiguous memory per request, causing massive waste. Modern systems treat KV-cache like virtual memory: paging, sharing, compression, and intelligent eviction dramatically improve throughput.

KV-Cache Optimizations

Seq len

512

Explore: KV cache dominates inference memory. MQA shares one K,V across all heads; GQA groups heads — both cut cache size while preserving most quality.

PagedAttention (vLLM)

Problem: Pre-allocating contiguous memory for max sequence length wastes 60–80% of GPU memory (requests rarely use full context).

Solution: Store KV-cache in fixed-size pages (blocks), allocated on demand.

Architecture:

Physical memory divided into pages of size $P$ tokens.
Each request has a page table mapping logical KV positions to physical pages.
Pages allocated only when needed; freed when request completes.

Benefits:

Near-zero internal fragmentation (waste $\leq P-1$ tokens per request).
Memory shared across requests with common prefixes.
Dynamic allocation enables much higher batch sizes.

Attention computation: Modified to gather pages from non-contiguous memory:

\text{Attention}(\mathbf{q}_t) = \text{softmax}\left(\frac{\mathbf{q}_t[\mathbf{K}_{\text{page}_1}; \ldots; \mathbf{K}_{\text{page}_n}]^T}{\sqrt{d_k}}\right)[\mathbf{V}_{\text{page}_1}; \ldots; \mathbf{V}_{\text{page}_n}]. \tag{1}

Memory utilization: Increases from ~20–40% (static) to ~95% (paged).

Prefix Caching

Observation: Many requests share common prefixes (system prompts, few-shot examples, shared documents).

Implementation: Hash the token sequence; if the KV-cache for a prefix already exists, reuse it via copy-on-write.

\text{Memory for } N \text{ requests with shared prefix } P_{\text{shared}}:

\text{Paged} = N \times T \times c, \quad \text{Prefix-cached} = P_{\text{shared}} \times c + N \times (T - P_{\text{shared}}) \times c. \tag{2}

Savings: For $N=100$ requests with $P_{\text{shared}} = 2000$ tokens: saves $100 \times 2000 \times c$ → massive for batch serving.

Radix tree: Organize cached prefixes in a radix tree for efficient lookup and longest-prefix matching.

KV-Cache Quantization

Reduce cache precision from FP16 to INT8 or INT4:

\mathbf{K}_{\text{int8}} = \text{round}\left(\frac{\mathbf{K} - z_K}{s_K}\right), \quad \mathbf{K} \approx s_K \cdot \mathbf{K}_{\text{int8}} + z_K. \tag{3}

Per-token or per-channel calibration:

Per-token: one scale/zero-point per token (per row of K/V).
Per-channel: one scale per feature dimension (per column).

Memory reduction:

FP16 → INT8: 2x reduction in KV-cache.
FP16 → INT4: 4x reduction.

Quality impact: INT8 KV-cache typically has negligible quality loss ( $<0.5\%$ on benchmarks). INT4 KV requires careful calibration but is increasingly viable.

KIVI (2024): Quantize K to per-channel INT2 and V to per-token INT2 with residual FP16 for recent tokens. Achieves 10x compression with minimal quality loss.

Token Eviction Strategies

When KV-cache exceeds budget, evict less important tokens.

Strategies:

1. Attention-score based (H2O): Keep tokens with highest cumulative attention scores ("Heavy Hitters"):

\text{importance}(j) = \sum_{t > j}\sum_h A_{tj}^{(h)}. \tag{4}

Keep top- $k$ tokens by importance + recent window.

2. Sliding window + sink tokens (StreamingLLM): Always keep the first few tokens ("attention sinks") + last $w$ tokens. Discard middle.

\text{Cache} = \{1, \ldots, s\} \cup \{t-w+1, \ldots, t\}. \tag{5}

3. Random eviction: Surprisingly competitive baseline — remove random tokens beyond the window.

4. Learned eviction: Train a small network to predict which tokens will be needed.

Sliding Window Cache

Only cache the last $w$ tokens (discard everything beyond):

\mathbf{K}_{\text{cache}} = \mathbf{K}_{t-w+1:t}, \quad |\text{cache}| = w \times 2Ld_k. \tag{6}

Fixed memory: Cache size is bounded regardless of context length.

Effective context: With $L$ layers and window $w$ , information can propagate $L \times w$ tokens through the residual stream (though exponentially attenuated).

Used in: Mistral (alternating sliding window + full attention layers).

Multi-Query and Grouped-Query Caching

MQA: All $H$ heads share one K/V → cache is $1/H$ the size of MHA.

GQA with $G$ groups: Cache is $G/H$ the size.

Tradeoff analysis:

\text{KV memory ratio} = \frac{\text{GQA cache}}{\text{MHA cache}} = \frac{G}{H}. \tag{7}

For LLaMA-2 70B: $H=64, G=8 \Rightarrow 8\times$ reduction.

This is orthogonal to quantization: GQA + INT4 KV = 32x total reduction.

Low-Rank KV Compression

Observation: Across the sequence, K and V matrices often lie in a low-rank subspace.

Approach: Project cached KV through learned low-rank matrices:

\tilde{\mathbf{K}} = \mathbf{K}\mathbf{P}_K, \quad \mathbf{P}_K \in \mathbb{R}^{d_k \times r}, \quad r \ll d_k. \tag{8}

Store $\tilde{\mathbf{K}} \in \mathbb{R}^{T \times r}$ instead of $\mathbf{K} \in \mathbb{R}^{T \times d_k}$ . Compression ratio: $d_k/r$ .

Attention becomes: $\mathbf{q}\mathbf{P}_K\tilde{\mathbf{K}}^T$ (extra $d_k \to r$ projection).

Cache-Aware Scheduling

Batching strategy: Group requests by prefix length to maximize prefix cache reuse.

Preemption: When memory is full, preempt (swap to CPU/disk) the least-recently-used request's KV-cache rather than evicting tokens.

Chunked prefill: Process long prompts in chunks, overlapping prefill of new requests with decode of existing ones — maximizes GPU utilization.

Common Pitfalls

Pitfall 1. Applying token eviction uniformly across layers. Different layers attend to different patterns; important tokens vary by layer. Layer-wise eviction policies perform better.

Pitfall 2. Ignoring "attention sink" tokens. The first few tokens often receive disproportionate attention regardless of content (due to softmax normalization). Evicting them breaks the model.

Pitfall 3. Over-quantizing KV-cache for long-context tasks. INT4 KV works well for short contexts but quality degrades for tasks requiring precise retrieval from early tokens.

Summary

PagedAttention eliminates memory fragmentation → near-100% memory utilization.
Prefix caching reuses shared prompt KV-cache across requests.
KV quantization (INT8/INT4) reduces cache 2–4x with minimal quality loss.
Token eviction (H2O, StreamingLLM) enables unbounded context with fixed memory.
Sliding window provides O(1) cache with bounded receptive field.
GQA + quantization compose multiplicatively: up to 32x cache reduction.

Exercises

Exercise 1. For PagedAttention with page size $P=16$ , compute the maximum wasted memory per request and compare to static allocation with max $T=32768$ .

Exercise 2. Derive the memory savings from prefix caching for $B=64$ requests sharing a 1024-token system prompt.

Exercise 3. Compute the KV-cache size for a 7B model ( $L=32, H=32, d_k=128$ ) at $T=4096$ in (a) FP16, (b) INT8, (c) INT4 with GQA $G=4$ .

Exercise 4. Design a token eviction policy for a streaming chat application where the model must remember user preferences stated at the beginning.

Exercise 5. For the "attention sink" phenomenon: explain mathematically why the first token receives high attention scores even when semantically irrelevant.