{"title":"11.7 KV Cache","description":"","content":"# 11.7 KV Cache\n\n## The Problem\n\nEvery time a language model generates a token, it runs a forward pass through the entire transformer. At each attention layer, the model computes three matrices from the input: queries (Q), keys (K), and values (V). The query represents \"what am I looking for?\", the keys represent \"what do I contain?\", and the values represent \"what information do I provide?\". The attention mechanism matches queries against keys to figure out which tokens to pay attention to, then uses values to produce the output.\n\nHere is the issue. During autoregressive generation, the model produces one token at a time. To generate the 1001st token, it needs to attend to all 1000 previous tokens. That means computing K and V matrices for every single previous token, even though those tokens have not changed since the last step. The keys and values for token 1 are exactly the same whether you are generating token 2 or token 1001. Recomputing them every time is pure waste.\n\nThe numbers get ugly fast. For a 1000-token sequence, generating each new token means recomputing 1000 sets of K and V projections across every attention layer. For a 32-layer model, that is 32,000 matrix multiplications that produce identical results to what you computed one step ago. Multiply that by every token you generate, and you are spending the vast majority of your compute budget redoing work.\n\nThe KV cache solves this by storing the key and value matrices from previous tokens so they can be reused. It is one of the most important optimizations in LLM inference, and understanding how it works (and how it consumes memory) is essential for anyone deploying models in production.\n\n**What you will learn:**\n\n- What the KV cache stores and why it eliminates redundant computation\n- How to calculate KV cache memory for any model configuration\n- Why KV cache (not model weights) is often the memory bottleneck in serving\n- How PagedAttention (vLLM) solves the memory fragmentation problem\n- How Grouped-Query Attention (GQA) reduces KV cache size at the architecture level\n- Compression and eviction strategies for extremely long contexts\n\n---\n\n## What the KV Cache Actually Stores\n\nTo understand the KV cache, you need to understand what happens inside a single attention layer during generation.\n\nWhen the model processes a token, it projects the token's hidden state into three vectors: Q (query), K (key), and V (value). During the initial prompt processing (called the \"prefill\" phase), the model computes Q, K, and V for every token in the input simultaneously. This is efficient because it happens in one batched matrix multiplication.\n\nBut during generation (the \"decode\" phase), you are producing one token at a time. At each step, only the new token's Q vector changes. The K and V vectors for all previous tokens are identical to what they were in the last step. Without caching, the model recomputes all of them anyway.\n\nThe KV cache stores the K and V matrices for every token that has already been processed, at every attention layer. When generating a new token, the model only computes Q, K, and V for that one new token, appends the new K and V to the cache, and then runs attention using the new Q against the full cached K and V.\n\n\n```mermaid\nflowchart TD\n subgraph Without[\"Without KV Cache (Step N)\"]\n direction TB\n I1[\"Input: all N tokens\"]:::red --> C1[\"Compute Q, K, V
for ALL N tokens\"]:::red\n C1 --> A1[\"Run attention\"]:::red\n A1 --> O1[\"Output token N+1\"]:::orange\n end\n\n subgraph With[\"With KV Cache (Step N)\"]\n direction TB\n I2[\"Input: token N only\"]:::green --> C2[\"Compute Q, K, V
for token N only\"]:::green\n CACHE[\"KV Cache
K₁..K_{N-1}, V₁..V_{N-1}\"]:::purple --> CONCAT[\"Append K_N, V_N
to cache\"]:::teal\n C2 --> CONCAT\n CONCAT --> A2[\"Run attention:
Q_N × [K₁..K_N]\"]:::green\n A2 --> O2[\"Output token N+1\"]:::orange\n end\n\n classDef red fill:#ff8787,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef purple fill:#9775fa,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n```\n\n\nThe left side recomputes everything from scratch at every step. The right side computes only what is new, reusing everything else from the cache. The computational savings are enormous. Instead of O(N^2) total computation over a sequence of length N, the decode phase drops to O(N) per step (one query attending to N cached keys).\n\nThis is why KV caching is not optional in practice. Without it, generating a 500-token response from a 1000-token prompt would require roughly 750,000 redundant K and V computations across the generation. With it, those computations happen exactly once.\n\n---\n\n## Memory Math\n\nThe KV cache sounds like a free win: save computation by trading it for memory. But memory is exactly the constraint you run into when serving LLMs. Understanding how much memory the KV cache consumes is critical for capacity planning.\n\nThe formula is:\n\n**KV cache memory = 2 x num_layers x num_kv_heads x head_dim x sequence_length x bytes_per_element**\n\nThe \"2\" accounts for storing both K and V. Let's break down each term:\n\n- **num_layers:** The number of transformer layers (attention blocks) in the model. Each layer maintains its own independent KV cache.\n- **num_kv_heads:** The number of key-value heads. In standard multi-head attention, this equals the number of attention heads. In GQA (which we will cover shortly), this is smaller.\n- **head_dim:** The dimensionality of each attention head. Typically `hidden_dim / num_heads`.\n- **sequence_length:** How many tokens are in the cache. This grows as the model generates.\n- **bytes_per_element:** 2 bytes for FP16/BF16, 1 byte for INT8.\n\n### Concrete Example: Llama 3 8B\n\nLlama 3 8B uses Grouped-Query Attention with these parameters:\n\n- 32 layers\n- 8 KV heads (GQA, not the full 32 attention heads)\n- 128 head dimension\n- FP16 (2 bytes per element)\n\n\n```python\ndef kv_cache_memory_bytes(\n num_layers: int,\n num_kv_heads: int,\n head_dim: int,\n sequence_length: int,\n bytes_per_element: int = 2, # FP16\n) -> int:\n \"\"\"Calculate KV cache memory in bytes for a single request.\"\"\"\n return (\n 2 # K and V\n * num_layers\n * num_kv_heads\n * head_dim\n * sequence_length\n * bytes_per_element\n )\n\ndef kv_cache_memory_mb(\n num_layers: int,\n num_kv_heads: int,\n head_dim: int,\n sequence_length: int,\n bytes_per_element: int = 2,\n) -> float:\n \"\"\"Calculate KV cache memory in megabytes for a single request.\"\"\"\n total_bytes = kv_cache_memory_bytes(\n num_layers, num_kv_heads, head_dim, sequence_length, bytes_per_element\n )\n return total_bytes / (1024 ** 2)\n\n# Llama 3 8B configuration\nllama3_8b = {\n \"num_layers\": 32,\n \"num_kv_heads\": 8, # GQA: 8 KV heads, not 32\n \"head_dim\": 128,\n}\n\nfor seq_len in [2048, 8192, 32768, 131072]:\n mem = kv_cache_memory_mb(**llama3_8b, sequence_length=seq_len)\n print(f\"Seq length {seq_len:>7,}: {mem:>8.1f} MB ({mem / 1024:.2f} GB)\")\n\n# Seq length 2,048: 128.0 MB (0.12 GB)\n# Seq length 8,192: 512.0 MB (0.50 GB)\n# Seq length 32,768: 2,048.0 MB (2.00 GB)\n# Seq length 131,072: 8,192.0 MB (8.00 GB)\n```\n\n\nHere is the same data in a table for quick reference:\n\n\n| Sequence Length | KV Cache per Request | KV Cache for 10 Requests | KV Cache for 100 Requests |\n|-----------------|---------------------|--------------------------|---------------------------|\n| 2,048 tokens | 128 MB | 1.25 GB | 12.5 GB |\n| 8,192 tokens | 512 MB | 5.0 GB | 50.0 GB |\n| 32,768 tokens | 2.0 GB | 20.0 GB | 200.0 GB |\n| 131,072 tokens | 8.0 GB | 80.0 GB | 800.0 GB |\n\n\nThese numbers are for a single request. The moment you start serving multiple concurrent users, multiply accordingly. That last column is where things get alarming.\n\n---\n\n## The Memory Problem\n\nLook at the table above one more time. A single Llama 3 8B request at 32K tokens needs 2 GB of KV cache. The model weights themselves (in FP16) are about 16 GB. So for a single long-context request, the KV cache is already 12.5% the size of the model. Serve 8 concurrent users at that context length, and the KV cache (16 GB) matches the model weights.\n\nThis flips most people's intuition. They assume model weights are the memory bottleneck. For short conversations, that is true. But for production serving with moderate batch sizes and reasonable context lengths, the KV cache dominates.\n\n\n```mermaid\nflowchart LR\n subgraph GPU[\"GPU Memory (80 GB A100)\"]\n direction TB\n W[\"Model Weights
16 GB (fixed)\"]:::primary\n KV[\"KV Cache
grows with batch size
and sequence length\"]:::red\n ACT[\"Activations &
Overhead
~2 GB\"]:::orange\n FREE[\"Free Memory
(determines max
batch size)\"]:::green\n end\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef red fill:#ff8787,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nOn an 80 GB A100 running Llama 3 8B in FP16:\n\n- Model weights: 16 GB (fixed)\n- Activations and overhead: ~2 GB (roughly fixed)\n- Available for KV cache: ~62 GB\n\nThat 62 GB determines how many concurrent requests you can serve:\n\n\n```python\navailable_kv_memory_gb = 62 # after model weights and overhead\n\n# How many concurrent requests at different context lengths?\nfor seq_len in [2048, 8192, 32768]:\n per_request_gb = kv_cache_memory_mb(**llama3_8b, sequence_length=seq_len) / 1024\n max_concurrent = int(available_kv_memory_gb / per_request_gb)\n print(f\"Seq length {seq_len:>6,}: {per_request_gb:.2f} GB/req → {max_concurrent} concurrent requests\")\n\n# Seq length 2,048: 0.12 GB/req → 496 concurrent requests\n# Seq length 8,192: 0.50 GB/req → 124 concurrent requests\n# Seq length 32,768: 2.00 GB/req → 31 concurrent requests\n```\n\n\nAt 2K context, you can handle hundreds of concurrent requests. At 32K, you are down to 31. At 128K (which Llama 3 supports), you can serve only 7 concurrent users on the same hardware. This is why KV cache optimization is not a nice-to-have. It directly determines your serving throughput and cost per request.\n\nThe problem gets worse in practice because traditional KV cache implementations pre-allocate memory for the maximum possible sequence length, even if most requests are much shorter. A request that only uses 500 tokens still reserves space for 32K tokens if that is your configured maximum. This fragmentation wastes massive amounts of GPU memory.\n\nWhich brings us to the single most impactful optimization in modern LLM serving.\n\n---\n\n## PagedAttention: vLLM's Innovation\n\nIf the KV cache memory problem sounds familiar, it should. Operating systems solved the exact same problem for regular RAM decades ago with virtual memory and paging. Programs request memory, the OS allocates it in fixed-size pages, and those pages do not need to be physically contiguous. This eliminates fragmentation and allows the OS to use memory far more efficiently.\n\nPagedAttention applies this same idea to the KV cache. Instead of pre-allocating a single contiguous block of GPU memory for each request's full context length, it divides the KV cache into small, fixed-size \"pages\" (called blocks in vLLM, typically 16 tokens each). Pages are allocated on demand as the sequence grows, and they do not need to be contiguous in GPU memory.\n\n\n```mermaid\nflowchart TD\n subgraph Traditional[\"Traditional KV Cache Allocation\"]\n direction LR\n R1[\"Request 1
(500 tokens used)\"]:::primary --> A1[\"Pre-allocated: 32K tokens
■■□□□□□□□□□□□□□□
94% wasted\"]:::red\n R2[\"Request 2
(2K tokens used)\"]:::primary --> A2[\"Pre-allocated: 32K tokens
■■■■■■□□□□□□□□□□
88% wasted\"]:::red\n R3[\"Request 3
(100 tokens used)\"]:::primary --> A3[\"Pre-allocated: 32K tokens
■□□□□□□□□□□□□□□□
99% wasted\"]:::red\n end\n\n subgraph Paged[\"PagedAttention Allocation\"]\n direction LR\n P1[\"Request 1
(500 tokens)\"]:::primary --> B1[\"32 pages allocated
■■■■■■■■■■■■■■■■
■■■■■■■■■■■■■■■■
No waste\"]:::green\n P2[\"Request 2
(2K tokens)\"]:::primary --> B2[\"128 pages allocated
Pages scattered in memory
No waste\"]:::green\n P3[\"Request 3
(100 tokens)\"]:::primary --> B3[\"7 pages allocated
■■■■■■■
No waste\"]:::green\n end\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef red fill:#ff8787,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nIn the traditional approach, three requests with a 32K max context pre-allocate 96K tokens worth of KV cache memory, but only use about 2,600 tokens combined. That is 97% waste. With PagedAttention, the same three requests allocate only the pages they actually need, and freed pages are immediately available for new requests.\n\nThe practical impact is dramatic. The original PagedAttention paper reported 2-4x throughput improvement over the HuggingFace Transformers baseline, primarily by fitting more concurrent requests into the same GPU memory. This is not a minor optimization. It is the reason vLLM became the default serving framework for LLM inference.\n\nPagedAttention also enables another important feature: KV cache sharing. When multiple requests share the same system prompt (common in production), vLLM can share those KV cache pages across requests using copy-on-write semantics, similar to how operating systems share memory pages between forked processes. If 100 requests all use the same 1000-token system prompt, the KV cache for those tokens is stored once, not 100 times.\n\nYou do not need to configure PagedAttention manually. It is the default memory management strategy in vLLM:\n\n\n```shell\n# PagedAttention is enabled by default in vLLM\npython -m vllm.entrypoints.openai.api_server \\\n --model meta-llama/Llama-3.1-8B-Instruct \\\n --gpu-memory-utilization 0.90 \\\n --port 8000\n```\n\n\nThe `--gpu-memory-utilization` flag controls what percentage of GPU memory vLLM uses for the model weights plus the KV cache page pool. Setting it to 0.90 means vLLM will use 90% of VRAM, leaving 10% as a buffer for CUDA overhead.\n\n---\n\n## Multi-Query Attention (MQA) and Grouped-Query Attention (GQA)\n\nPagedAttention reduces memory waste from fragmentation. But there is another approach that reduces the KV cache size itself at the architecture level: sharing key-value heads across attention heads.\n\nIn standard Multi-Head Attention (MHA), each attention head has its own set of Q, K, and V projections. If a model has 32 attention heads, it has 32 sets of keys and 32 sets of values per layer. That is the full KV cache.\n\nMulti-Query Attention (MQA), introduced by Google in 2019, takes this to the extreme: all attention heads share a single set of K and V. The 32 query heads each still produce different queries, but they all attend to the same keys and values. This cuts the KV cache to 1/32 of the original size.\n\nThe problem is that MQA can hurt quality. Sharing a single K and V across all heads is a pretty aggressive compression. Some tasks degrade noticeably.\n\nGrouped-Query Attention (GQA) splits the difference. Instead of one shared KV head (MQA) or 32 independent ones (MHA), GQA groups the attention heads and gives each group its own K and V. Llama 3, for example, has 32 query heads grouped into 8 KV groups. Each group of 4 query heads shares one set of K and V.\n\n\n```mermaid\nflowchart TD\n subgraph MHA[\"Multi-Head Attention (MHA)\"]\n direction LR\n Q1_1[\"Q₁\"]:::primary\n Q1_2[\"Q₂\"]:::primary\n Q1_3[\"Q₃\"]:::primary\n Q1_4[\"Q₄\"]:::primary\n K1_1[\"K₁\"]:::orange\n K1_2[\"K₂\"]:::orange\n K1_3[\"K₃\"]:::orange\n K1_4[\"K₄\"]:::orange\n\n Q1_1 --- K1_1\n Q1_2 --- K1_2\n Q1_3 --- K1_3\n Q1_4 --- K1_4\n end\n\n subgraph GQA[\"Grouped-Query Attention (GQA)\"]\n direction LR\n Q2_1[\"Q₁\"]:::primary\n Q2_2[\"Q₂\"]:::primary\n Q2_3[\"Q₃\"]:::primary\n Q2_4[\"Q₄\"]:::primary\n K2_1[\"K₁\"]:::orange\n K2_2[\"K₂\"]:::orange\n\n Q2_1 --- K2_1\n Q2_2 --- K2_1\n Q2_3 --- K2_2\n Q2_4 --- K2_2\n end\n\n subgraph MQA[\"Multi-Query Attention (MQA)\"]\n direction LR\n Q3_1[\"Q₁\"]:::primary\n Q3_2[\"Q₂\"]:::primary\n Q3_3[\"Q₃\"]:::primary\n Q3_4[\"Q₄\"]:::primary\n K3_1[\"K₁\"]:::orange\n\n Q3_1 --- K3_1\n Q3_2 --- K3_1\n Q3_3 --- K3_1\n Q3_4 --- K3_1\n end\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n```\n\n\nThe diagram shows 4 query heads for simplicity, but the same principle scales to 32 or more. MHA uses 4 KV heads, GQA uses 2, and MQA uses 1. In each case, the V heads follow the same pattern as the K heads (omitted from the diagram for clarity).\n\nThe memory savings are straightforward to calculate:\n\n\n```python\n# KV cache comparison: MHA vs GQA vs MQA for Llama 3 8B config\nconfigs = {\n \"MHA (32 KV heads)\": {\"num_layers\": 32, \"num_kv_heads\": 32, \"head_dim\": 128},\n \"GQA (8 KV heads)\": {\"num_layers\": 32, \"num_kv_heads\": 8, \"head_dim\": 128},\n \"MQA (1 KV head)\": {\"num_layers\": 32, \"num_kv_heads\": 1, \"head_dim\": 128},\n}\n\nseq_len = 8192\n\nprint(f\"KV cache at {seq_len:,} tokens:\\n\")\nfor name, config in configs.items():\n mem = kv_cache_memory_mb(**config, sequence_length=seq_len)\n print(f\" {name}: {mem:>8.0f} MB ({mem / 1024:.2f} GB)\")\n\n# KV cache at 8,192 tokens:\n#\n# MHA (32 KV heads): 2048 MB (2.00 GB)\n# GQA (8 KV heads): 512 MB (0.50 GB)\n# MQA (1 KV head): 64 MB (0.06 GB)\n```\n\n\nGQA gives Llama 3 a 4x reduction in KV cache compared to full MHA, with negligible quality loss. MQA would give a 32x reduction, but at a meaningful quality cost. GQA is the sweet spot that most modern models have converged on, including Llama 3, Mistral, and Gemma 2.\n\n\n| Architecture | KV Heads (for 32 Q heads) | KV Cache per Request (8K) | Quality Impact | Used By |\n|-------------|---------------------------|---------------------------|----------------|---------|\n| MHA | 32 | 2.0 GB | None (baseline) | GPT-2, older models |\n| GQA | 8 | 512 MB | Negligible | Llama 3, Mistral, Gemma 2 |\n| GQA | 4 | 256 MB | Very small | Some smaller models |\n| MQA | 1 | 64 MB | Noticeable on some tasks | PaLM, Falcon |\n\n\nThis is an architectural decision made during model training, not something you configure at serving time. But it directly affects your inference costs. When choosing between models for deployment, the number of KV heads is a factor worth checking. A model with fewer KV heads is cheaper to serve at the same context length and batch size.\n\n---\n\n## KV Cache Compression and Eviction\n\nEven with GQA and PagedAttention, extremely long contexts (100K+ tokens) or high concurrency can still exhaust GPU memory. When that happens, there are several strategies to reduce KV cache memory further, each with its own trade-off.\n\n### Quantizing the KV Cache\n\nJust as you can quantize model weights from FP16 to INT8 or INT4, you can quantize the KV cache values. This halves (INT8) or quarters (INT4) the cache memory with a small accuracy trade-off.\n\n\n```python\n# vLLM supports KV cache quantization\n# Enable FP8 KV cache for ~50% memory reduction\n# Launch with:\n# python -m vllm.entrypoints.openai.api_server \\\n# --model meta-llama/Llama-3.1-8B-Instruct \\\n# --kv-cache-dtype fp8 \\\n# --port 8000\n\n# Memory comparison\nfor dtype, bytes_per in [(\"FP16\", 2), (\"FP8\", 1)]:\n mem = kv_cache_memory_mb(\n **llama3_8b,\n sequence_length=32768,\n bytes_per_element=bytes_per,\n )\n print(f\"KV cache ({dtype}) at 32K tokens: {mem:.0f} MB ({mem / 1024:.1f} GB)\")\n\n# KV cache (FP16) at 32K tokens: 2048 MB (2.0 GB)\n# KV cache (FP8) at 32K tokens: 1024 MB (1.0 GB)\n```\n\n\nFP8 KV cache quantization is becoming standard in production deployments. The quality impact is minimal for most tasks, and the memory savings directly translate to more concurrent requests or longer context support.\n\n### Token Eviction\n\nNot all tokens in the KV cache are equally important. Research has shown that attention scores follow a power-law distribution: a small number of tokens receive the majority of attention, while most tokens are rarely attended to. Token eviction strategies exploit this by removing low-attention tokens from the cache.\n\nThe simplest eviction policy is to keep the most recent tokens (a sliding window) plus a fixed number of \"anchor\" tokens from the beginning of the sequence (like the system prompt). More sophisticated approaches track attention scores over time and evict tokens that consistently receive low attention.\n\n### Sliding Window Attention\n\nMistral pioneered sliding window attention, where each layer only attends to the last W tokens (typically 4,096) instead of the full sequence. This caps the KV cache at a fixed size regardless of total context length.\n\n\n```mermaid\nflowchart LR\n subgraph Full[\"Full Attention (Standard)\"]\n direction LR\n T1_1[\"Token 1\"]:::teal\n T1_2[\"Token 2\"]:::teal\n T1_3[\"...\"]:::teal\n T1_4[\"Token 8K\"]:::teal\n T1_5[\"Token 8001\"]:::primary\n T1_1 --> T1_5\n T1_2 --> T1_5\n T1_3 --> T1_5\n T1_4 --> T1_5\n end\n\n subgraph SW[\"Sliding Window (W=4096)\"]\n direction LR\n T2_1[\"Token 1\"]:::red\n T2_2[\"Token 2\"]:::red\n T2_3[\"...\"]:::red\n T2_4[\"Token 4905\"]:::teal\n T2_5[\"...\"]:::teal\n T2_6[\"Token 8K\"]:::teal\n T2_7[\"Token 8001\"]:::primary\n T2_4 --> T2_7\n T2_5 --> T2_7\n T2_6 --> T2_7\n end\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef teal fill:#38d9a9,stroke:#000,color:#000\n classDef red fill:#ff8787,stroke:#000,color:#000\n```\n\n\nIn full attention, token 8001 attends to all 8000 previous tokens. With a sliding window of 4096, it only attends to the most recent 4096 tokens. Earlier tokens (shown in red) are evicted from the cache. Information from those earlier tokens can still influence the output indirectly because they affected the hidden states of intermediate tokens, but direct attention to them is lost.\n\n### Choosing a Strategy\n\n\n| Strategy | Memory Savings | Quality Impact | Complexity | Best For |\n|----------|---------------|----------------|------------|----------|\n| KV cache quantization (FP8) | 2x | Minimal | Low (one flag) | Always a good default |\n| Sliding window | Caps at window size | Moderate for long-range tasks | Built into model | Models that support it (Mistral) |\n| Token eviction | Variable (2-4x) | Depends on policy | Medium | Research, long documents |\n| KV cache offloading (CPU) | Frees GPU memory | Adds latency | Medium | Very long contexts, latency-tolerant |\n\n\nFor most production deployments, start with FP8 KV cache quantization (it is nearly free in terms of quality) and use PagedAttention via vLLM. Only reach for eviction or sliding window strategies if you are pushing context lengths beyond what your GPU memory can handle.\n\n---\n\n## Exercise\n\nPut your understanding to the test with these tasks:\n\n**Task 1: Calculate KV cache memory for Llama 3 8B.**\n\nUsing the formula from this chapter, calculate the KV cache memory per request at 2K, 8K, 32K, and 128K tokens. Then calculate how many concurrent requests an 80 GB A100 can handle at each context length (assume 16 GB for model weights and 2 GB for overhead).\n\n**Task 2: Compare MHA vs GQA memory usage.**\n\nModify the calculation to compare what Llama 3 8B's KV cache would look like if it used full MHA (32 KV heads) instead of GQA (8 KV heads). At 32K context length, how many more concurrent requests does GQA enable on the same hardware?\n\n**Task 3: Deploy with vLLM and measure actual memory.**\n\nDeploy Llama 3.1 8B Instruct with vLLM on a GPU. Send requests at increasing context lengths (1K, 4K, 8K, 16K) and use `nvidia-smi` to observe actual GPU memory usage. Compare the measured memory growth against your theoretical calculations. How close are they?\n\n**Task 4: Test FP8 KV cache quantization.**\n\nDeploy the same model twice: once with default FP16 KV cache and once with `--kv-cache-dtype fp8`. Run the same benchmark (50 prompts, 8K context) on both configurations. Compare throughput, memory usage, and output quality (use an LLM-as-judge to score response pairs).\n\n---\n\n## Summary\n\nHere are the key takeaways from this chapter:\n\n- **The KV cache stores key and value matrices from previous tokens** so they are not recomputed during autoregressive generation. Without it, every new token would redundantly recompute K and V for all preceding tokens.\n- **KV cache memory = 2 x layers x KV heads x head_dim x sequence_length x bytes_per_element.** For Llama 3 8B (GQA with 8 KV heads), that is 128 MB per request at 2K tokens and 2 GB per request at 32K tokens.\n- **KV cache, not model weights, is often the memory bottleneck.** At moderate batch sizes and context lengths, the KV cache can consume more GPU memory than the model itself.\n- **PagedAttention (vLLM) eliminates memory fragmentation** by allocating KV cache in small pages on demand, like OS virtual memory. This is the key innovation that gives vLLM 2-4x throughput improvements.\n- **GQA reduces KV cache by sharing KV heads across query head groups.** Llama 3 uses 8 KV heads instead of 32, cutting KV cache to 25% of what full MHA would require.\n- **FP8 KV cache quantization halves memory with minimal quality loss** and is the easiest optimization to enable in production (one vLLM flag).\n\nImportant numbers to remember: Llama 3 8B KV cache is ~0.5 GB per request at 8K context, ~2 GB at 32K. On an 80 GB A100, that means roughly 124 concurrent 8K requests or 31 concurrent 32K requests.\n\nKV cache optimization reduces memory pressure during inference, letting you serve more users on the same hardware. But the KV cache only stores intermediate attention states. The model weights themselves, typically 16 GB for an 8B model in FP16, are a separate memory cost. What if you could shrink those weights too? That is exactly what quantization does, and it is the topic of the next chapter.\n\n---\n\n## References\n\n- [Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)](undefined)\n- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)](undefined)\n- [Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)](undefined)\n- [vLLM Documentation: KV Cache Management](undefined)\n- [Efficient Streaming Language Models with Attention Sinks (Xiao et al., 2023)](undefined)\n- [Mistral 7B (Jiang et al., 2023) - Sliding Window Attention](undefined)","pageType":"ai-engineering"}

11.7 KV Cache

Ashish Pratap Singh

Get Premium