Production-Grade KV Cache Optimization: Memory Budgeting, PagedAttention, Quantization & Offloading

1. Why Autoregressive Generation Needs KV Cache

In Transformer attention, each token projects to Query, Key, Value vectors. The current token's Query must dot-product with all historical Keys to get attention weights, then weighted-sum historical Values. Without caching, every new token would recompute all prior Keys and Values, turning quadratic complexity into cubic as sequence grows. Since historical Keys/Values never change, caching them avoids pure compute waste.

1.1 KV Cache Mechanism

After forward pass, each token's Key and Value are stored in GPU memory. During decoding, only the new token's Q/K/V are computed; new K/V are appended to the cache. Query is not cached because it's only used for the current step. The cache is a tuple past_key_values (Hugging Face) holding per-layer K/V tensors.

1.2 Prefill vs Decode Phases

Prefill : Processes the entire prompt in parallel, computes all K/V, initializes KV cache. Bottleneck is compute.

Decode : Generates tokens one by one, each step computes Q/K/V for the single new token, reuses cached K/V. Bottleneck is memory bandwidth.

import torch
with torch.no_grad():
    past_key_values = None
    generated = input_ids
    for _ in range(max_new):
        if past_key_values is None:
            inputs = generated  # Prefill: full prompt
        else:
            inputs = generated[:, -1:]  # Decode: last token only
        out = model(input_ids=inputs, past_key_values=past_key_values, use_cache=True)
        past_key_values = out.past_key_values
        logits = out.logits[:, -1, :]
        next_token = logits.argmax(dim=-1, keepdim=True)
        if next_token.item() == tokenizer.eos_token_id:
            break
        generated = torch.cat([generated, next_token], dim=-1)

2. Memory Budgeting Before Deployment

KV cache often exceeds model weights in long-context/high-concurrency scenarios. The per-token KV memory formula:

Per-token KV bytes = 2 × L × H_kv × D × P

2 : one Key + one Value

L : Transformer layers

H_kv : KV heads (GQA reduces this vs MHA)

D : head dimension

P : bytes per element (BF16/FP16=2, FP8=1, NVFP4/INT4=0.5)

Total GPU memory = Model weights + Peak KV cache (per-request KV × concurrent requests) + Runtime overhead (10-20% headroom).

2.1 Example: Llama-2 7B (32 layers, 32 KV heads, head_dim=128, BF16)

2 × 32 × 32 × 128 × 2 = 524,288 bytes ≈ 512 KB per token

4K context → ~2 GB per request

32K context → ~16 GB per request (exceeds 13-14 GB model weights)

Conclusion: short context → weights dominate; long context → KV cache dominates.

3. PagedAttention: OS-Style Paging for KV Cache

3.1 Problem: Static Contiguous Allocation

Early frameworks reserved a contiguous block per request sized for max context. Most requests use far less, causing 60-80% memory waste (vLLM paper).

3.2 PagedAttention Mechanism

KV cache split into fixed-size blocks (pages). Logical sequence maps via page table to disjoint physical pages, allocated on demand and freed after request. Eliminates fragmentation.

3.3 Prefix Caching & Copy-on-Write

Shared prefixes (system prompts, retrieved docs) map to same physical pages. On divergence, copy-on-write creates new pages. Dramatically reduces memory and avoids redundant prefill.

3.4 Production Deployment (vLLM Example)

vllm serve meta-llama/Llama-2-7b-chat-hf \
  --gpu-memory-utilization 0.90 \
  --max-model-len 32768 \
  --block-size 16 \
  --enable-prefix-caching

Key params: --gpu-memory-utilization (leave headroom), --max-model-len, --block-size (default 16 tokens/page), --enable-prefix-caching (on by default in vLLM v1). Monitor: KV cache usage, prefix cache hit rate, throughput, TTFT, TPOT, memory spikes.

4. KV Cache Quantization: Reducing Storage Bitwidth

Quantization compresses KV cache independently of model weights. Two benefits: (1) lower memory footprint → longer context/higher concurrency; (2) reduced decode memory bandwidth → faster token generation.

4.1 Main Schemes

FP8 : 1 byte/element, halves memory vs BF16. Negligible quality loss for chat/RAG. Supported on H100/H200. Default production choice.

NVFP4 : 4-bit block floating point with FP8 scaling factor. Halves memory again vs FP8. Requires Blackwell (B100/B200/GB200). NVIDIA reports near-doubled context budget, up to 3× TTFT improvement, <1% quality drop.

Warning: quantization adds numerical error; long sequences accumulate error. Must validate on real business data.

4.2 FP8 Config (vLLM)

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Meta-Llama-3-8B \
  --enable-prefix-caching \
  --kv-cache-dtype fp8

4.3 Priority Order

PagedAttention + Prefix Caching (no precision change, stable throughput gain)

FP8 KV quantization (H100/H200, halves memory)

NVFP4 on Blackwell for long-context

Validate on long-doc retrieval, multi-hop QA, chain-of-thought

5. Hierarchical Production Optimization Landscape

5.1 Architecture-Level: GQA & MLA

MHA : each query head has own K/V heads → largest KV.

GQA : groups of query heads share K/V heads → reduces KV heads. Used in Qwen2.5, Llama 3, Mistral.

MLA : stores low-rank latent vectors instead of raw K/V, reconstructs at decode. DeepSeek-V2/V3, 8-25× compression. No runtime overhead, no quantization error.

5.2 Memory Management: PagedAttention + Prefix Caching

Solves fragmentation and redundant prefill. First stable optimization layer.

5.3 Hierarchical KV Offload (Cold/Hot Tiering)

L0 GPU HBM : hot KV (active generation)

L1 CPU DRAM : warm KV (recent history)

L2 NVMe SSD : cold KV (early conversation)

New KV starts in HBM; as attention heat drops, pages demote asynchronously. On access, promote back. Used by Baidu Bage AttentionStore, LMCache. Trade-off: latency jitter from data movement.

5.4 KV Eviction Policies

When offload not desired, evict low-importance tokens: H2O, SnapKV, LU-KV, StreamingLLM. Score tokens by attention weights, keep critical + recent window. Risk: loses rare but crucial details (code analysis, precise reasoning). Must validate per use-case.

5.5 Recommended Rollout Sequence

Enable PagedAttention + Prefix Caching

Add FP8 KV quantization if memory tight

For ultra-long Agent contexts, add hierarchical offload

On Blackwell, enable NVFP4

Last resort: KV eviction / sliding window (lossy)

6. Summary Checklist

PagedAttention & Prefix Caching configured correctly for engine version

Memory budget, peak concurrency, max context calculated

Business has stable shared prefixes to benefit from prefix caching

Root causes of throughput issues and OOM: missing these three. Choose GQA/MLA models, enable PagedAttention, then layer quantization/offload per hardware and workload.