If you are deploying Large Language Models (LLMs) or multi-tenant AI agents in production, you have probably run into unexpected Out-Of-Memory (OOM) errors.
Most engineering teams size their GPUs based on model weights (~140 GB for a 70B parameter model in FP16) and treat the remaining VRAM as extra headroom. In production, that assumption breaks down fast.
The silent killer of LLM deployments isn't the model weights—it's the KV (Key-Value) Cache.
1. What is KV Cache and Why Does It Spill VRAM?
During LLM generation, the model stores the key and value tensors of previously processed tokens in VRAM so it doesn't have to recompute attention over the full sequence at each new step.
While this dramatically accelerates inference, the cache grows linearly with batch size and context length. On long-running agent workloads (16K–32K+ tokens), the KV cache often consumes far more VRAM than the model weights themselves.
The Problem: A single agent session with a 32K context window can eat tens of gigabytes of VRAM before generating a single output token. Multiply that by concurrent sessions, and OOM errors become guaranteed.
2. The Math: Calculating KV Cache Memory Footprint
The KV cache memory footprint scales according to this exact formula:
Memory = 2 × batch_size × seq_len × num_layers × num_kv_heads × head_dim × precision_bytes
(The leading 2 accounts for storing both Key and Value tensors separately).
Real-World Example: Llama 2 70B at 32K Context
Let's plug in the parameters for Llama 2 70B:
- Layers: 80
- KV Heads: 8 (using Grouped-Query Attention)
- Head Dimension: 128
- Precision: FP16 (2 bytes per value)
For 1 single sequence at a 32,768-token context:
Memory = 2 × 1 × 32,768 × 80 × 8 × 128 × 2 bytes
= 10,737,418,240 bytes
≈ 10.74 GB per sequence
A single user eats 10.74 GB of VRAM. If Llama 2 70B used standard Multi-Head Attention (64 heads) instead of GQA (8 heads), that number would soar to over 85 GB per sequence.
3. Core KV Cache Optimization Techniques
To prevent memory starvation, modern serving frameworks combine four main architectural solutions:
| Technique | How It Works | Impact |
|---|---|---|
| PagedAttention | Applies OS virtual memory principles to GPU RAM, dividing cache into non-contiguous blocks via block tables. | Reduces fragmentation waste from 60–80% to <4%. |
| KV Cache Quantization | Reduces cached key/value precision from FP16 down to FP8 or INT4. | Cuts VRAM usage by 50% (FP8) with negligible quality loss. |
| Prefix Caching | Stores and reuses KV cache for shared system prompts, RAG documents, or tool schemas. | Eliminates duplicate prefill compute and memory across requests. |
| Eviction Policies | Evicts low-importance tokens (e.g., StreamingLLM or Heavy-Hitter approaches). | Enables unbounded context lengths without hard OOM limits. |
4. Hands-On Tutorial: Implementing PagedAttention with vLLM
vLLM utilizes PagedAttention as its default memory manager. Here is how to configure and deploy an optimized inference server.
Step 1: Install vLLM
pip install vllm
Step 2: Start the Server with PagedAttention & Optimization Flags
Launch the server with high memory allocation, prefix caching, and FP8 KV cache quantization:
vllm serve meta-llama/Llama-2-70b-chat-hf \
--gpu-memory-utilization 0.90 \
--max-model-len 32768 \
--dtype float16 \
--enable-prefix-caching \
--kv-cache-dtype fp8
-
--gpu-memory-utilization 0.90: Allocates 90% of GPU VRAM for model weights and the paged KV cache pool. -
--enable-prefix-caching: Avoids recomputing identical prompt prefixes across requests. -
--kv-cache-dtype fp8: Halves the KV cache memory footprint per token.
Step 3: Run Concurrent Load Benchmarks
vLLM includes built-in benchmarking scripts to measure throughput under heavy load:
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-2-70b-chat-hf &
python benchmarks/benchmark_serving.py \
--backend vllm \
--model meta-llama/Llama-2-70b-chat-hf \
--num-prompts 100 \
--request-rate 10
Monitor memory usage via nvidia-smi during execution to verify that memory scales cleanly to your ceiling without crashing.
5. Research Benchmarks & Performance Impact
According to the original PagedAttention research paper (Kwon et al., SOSP 2023):
- Memory Waste: Dropped from 60–80% (naive contiguous allocation) down to under 4%.
- Throughput: Achieved 2–4x higher throughput compared to prior state-of-the-art serving systems (like Orca and FasterTransformer) at equivalent latency.
- Production Scale: Real-world deployments (such as LMSYS Chatbot Arena) reduced total required GPUs by 50% while serving 2–3x more traffic.
💡 Summary & Next Steps
Optimizing the KV cache isn't optional for long-context LLM applications—it is mandatory for keeping infrastructure costs under control and avoiding OOM downtime.
👉 Originally published on GPUYard. Check out the original article for more deep dives into LLM infrastructure and MLOps optimizations!
Top comments (0)