KV-Cache Optimization: Memory-Efficient Inference

Learning Objectives
Notation
Core Intuition
The KV-Cache Memory Problem
PagedAttention (vLLM)
Continuous Batching
KV-Cache Quantization
KV-Cache Eviction & Compression
Chunked Prefill & Disaggregated Serving
MQA & GQA for Cache Reduction
Common Pitfalls
Summary
Exercises

Learning Objectives

Compute KV-cache memory for a given model, batch, and sequence length.
Explain PagedAttention's virtual memory approach and fragmentation elimination.
Compare static vs continuous batching throughput.
Design KV-cache eviction policies that preserve quality.
Analyze memory savings from MQA, GQA, and cache quantization.

Notation

$B$ — batch size, $S$ — sequence length, $L$ — layers
$H$ — heads, $d_k$ — head dimension
KV memory per token: $2 \times L \times H \times d_k \times \text{bytes}$

Core Intuition

During autoregressive generation, the model caches all previous key-value pairs (KV-cache) to avoid recomputation. For a 70B model generating 4096 tokens with batch size 32: the KV-cache alone requires 160GB+ — exceeding the model weights themselves. Managing this memory efficiently (paging, compression, eviction) is THE bottleneck for high-throughput LLM serving.

KV-Cache Optimization (Paged Attention)

Seq len

512

Pre-allocatedPaged blocks

Explore: PagedAttention (vLLM) allocates KV-cache in fixed-size pages on demand — eliminating memory waste from pre-allocated max-length buffers and enabling higher batch sizes.

The KV-Cache Memory Problem

Memory per request:

\text{KV memory} = 2 \times L \times H \times d_k \times S \times \text{dtype\_size}. \tag{1}

Example (LLaMA-70B, BF16):

$L=80$ , $H=64$ , $d_k=128$ , dtype=2 bytes.
Per token: $2 \times 80 \times 64 \times 128 \times 2 = 2.6$ MB.
For 4096 tokens: $2.6 \times 4096 = 10.7$ GB per request.
Batch of 32: $10.7 \times 32 = 342$ GB (just for KV-cache).

The challenge: Model weights (140GB in BF16) + KV-cache (342GB) = 482GB. Need 6+ H100 80GB GPUs JUST for one batch.

Why it matters: KV-cache memory limits maximum batch size → limits throughput → limits cost-efficiency.

PagedAttention (vLLM)

Kwon et al. (2023): Treat KV-cache like virtual memory — allocate in pages:

Problem with naive allocation:

Pre-allocate max-length cache for each request (e.g., 4096 tokens).
Most requests generate fewer tokens → massive internal fragmentation.
Average waste: 60-80% of allocated KV memory is unused.

PagedAttention solution:

Divide KV-cache into fixed-size pages (e.g., 16 tokens per page).
Allocate pages on-demand as generation proceeds.
Pages can be non-contiguous in physical memory (page table maps logical → physical).
When request finishes: reclaim pages immediately.

Benefits:

Near-zero memory waste (no internal fragmentation).
2-4x more requests served simultaneously.
Memory sharing: identical prefixes share KV pages (copy-on-write).

Continuous Batching

Static batching (naive): Wait until all requests in a batch finish → start next batch.

Problem: Fast requests finish early; GPU idle until slowest request completes.

Continuous batching (iteration-level scheduling):

After each generation step: check for completed requests.
Immediately replace completed requests with new ones from the queue.
GPU always has maximum batch size active.

Throughput improvement: 2-5x over static batching (depends on output length variance).

Combined with PagedAttention: vLLM achieves both memory efficiency and scheduling efficiency.

KV-Cache Quantization

Reduce precision of cached keys/values:

FP16 → INT8: 2x memory reduction. Quality impact: minimal (less than 0.5% degradation).

FP16 → INT4: 4x memory reduction. Quality impact: noticeable but acceptable for many applications.

Per-channel quantization:

\text{KV}_{\text{quant}} = \text{round}\left(\frac{\text{KV} - z}{s}\right), \quad s = \frac{\max - \min}{2^b - 1}. \tag{2}

Store scale $s$ and zero-point $z$ per channel.

KIVI (Liu et al., 2024): Key cache in INT2, Value cache in INT2, with per-channel calibration. 4-8x memory savings with less than 1% quality loss.

KV-Cache Eviction & Compression

When cache exceeds budget: Evict tokens that are least important for future generation.

Eviction policies:

1. StreamingLLM (Xiao et al., 2023): Keep attention sinks (first tokens) + recent window:

Keep first 4 tokens (attention sinks — always attended to).
Keep last $W$ tokens (recent context).
Evict everything in between.
Supports infinite-length generation with fixed memory.

2. H2O (Heavy Hitter Oracle): Keep tokens that received highest cumulative attention:

\text{score}(t) = \sum_{\text{recent layers}} \sum_{\text{heads}} \alpha_t, \tag{3}

Keep top- $K$ tokens by score + recent window.

3. Scissorhands: Keep tokens at attention boundaries (where attention patterns change).

4. PyramidKV: Different cache budgets per layer (early layers need less cache).

Chunked Prefill & Disaggregated Serving

Chunked prefill: Process long prompts in chunks rather than all at once:

Avoids one massive prefill compute spike.
Interleaves prefill chunks with generation steps.
Better latency for concurrent requests.

Disaggregated serving (Splitwise, DistServe):

Separate "prefill" and "decode" GPU pools.
Prefill GPUs: optimized for compute (large batch, short sequence).
Decode GPUs: optimized for memory bandwidth (small batch, long cache).
Transfer KV-cache from prefill → decode GPUs after prompt processing.

Benefit: Each GPU type runs at peak efficiency for its workload.

MQA & GQA for Cache Reduction

Multi-Query Attention (MQA): Single K,V shared across all heads:

KV-cache: $1/H$ of multi-head attention.
For $H=64$ : 64x less KV memory.
Quality: slight degradation (2-5% on benchmarks).

Grouped-Query Attention (GQA): $G$ groups share K,V:

KV-cache: $G/H$ of MHA (e.g., $G=8, H=64$ : 8x less).
Quality: near MHA (within 1%).
Used by: LLaMA-2 70B ( $G=8$ ), Mistral ( $G=8$ ).

Memory savings (LLaMA-70B, 4096 tokens, batch 32):

MHA: 342GB KV-cache.
GQA (8 groups): 42.7GB KV-cache.
MQA (1 group): 5.3GB KV-cache.

Common Pitfalls

Pitfall 1. Pre-allocating max sequence length for all requests. If average generation is 200 tokens but max is 4096: 95% of memory is wasted. Use paged allocation (vLLM).

Pitfall 2. KV-cache eviction without keeping attention sinks. The first few tokens typically receive disproportionate attention (attention sink phenomenon). Evicting them causes catastrophic quality collapse.

Pitfall 3. Quantizing KV-cache to INT4 without per-channel calibration. Global quantization causes large errors for outlier channels. Per-channel or per-token quantization is essential.

Summary

KV-cache: Dominates inference memory (often exceeds model weights).
PagedAttention: Virtual memory for KV; eliminates fragmentation; 2-4x efficiency.
Continuous batching: Replace finished requests immediately; 2-5x throughput.
Cache quantization: INT8/INT4 for 2-4x memory savings with minimal quality loss.
Eviction (StreamingLLM, H2O): Fixed-budget attention for infinite generation.
GQA: Architecture-level 8x KV reduction (LLaMA-2, Mistral).

Exercises

Exercise 1. Compute KV-cache memory for Mixtral 8x7B (32 layers, 8 KV heads, dim 128) at batch 64, sequence 8192. Compare with model weight memory.

Exercise 2. For PagedAttention with 16-token pages: compute fragmentation vs naive allocation for a request that generates 137 tokens.

Exercise 3. Design a StreamingLLM configuration for summarizing a 100K-token document: specify sink tokens, window size, and analyze which information is lost.

Exercise 4. Compare throughput (tokens/second) of static vs continuous batching for a workload where output lengths follow exponential distribution with mean 200 tokens.

Exercise 5. For KV-cache INT4 quantization: measure the attention score error introduced by quantizing keys to 4 bits. When does this become problematic?