KV cache and PagedAttention

Why this matters

The single biggest reason inference is fast is the KV cache. The single biggest reason inference is expensive (in memory) is also the KV cache. Every inference engineer either understands the cache deeply or fails at scale.

This lesson opens up the cache: how it is structured, how vLLM's PagedAttention manages it across many concurrent requests, and how prefix caching turns "read this 5000-token system prompt every time" into a one-time cost.

Learning objectives

1. Explain why KV cache reduces decoding from O(T^2) to O(T) per generated token.
2. Compute KV cache memory exactly for a given model and batch.
3. Describe PagedAttention's page-based layout.
4. Use prefix caching for repeated system prompts.
5. Reason about KV-cache quantisation (FP8 / INT8) trade-offs.

1. The setup

During training, a transformer processes T tokens in parallel — one big forward pass.

During autoregressive generation, you produce one new token at a time. Without a cache, each new token would require a full forward pass over the entire sequence so far. To generate N tokens after a prompt of P tokens you would do O((P+N)^2) work — way too expensive.

The KV cache stores the K and V tensors per layer per token as you compute them. At step t, only the new token's Q is computed; the K, V are appended to the cache; attention reads against the entire (cached + new) K, V.

So:

• Prefill (process the prompt): one big forward pass, fills the cache.
• Decode (generate token by token): one tiny forward pass per token, reads the cache.

This is why prompts are processed fast (parallel) and generation is slower per token (sequential, memory-bound).

2. Memory cost — the formula you must memorise

For a model with L layers, h_kv KV heads, head dim d_h, dtype b bytes:

KV memory per token = 2 * L * h_kv * d_h * b
                    = 2 * L * d_kv * b      (where d_kv = h_kv * d_h)

For batch B, sequence T:

total KV memory = 2 * B * T * L * d_kv * b

Worked examples

Llama-3 8B: L=32, h_kv=8, d_h=128, bf16 (b=2):

• per token: 2 * 32 * 8 * 128 * 2 = 131072 bytes ≈ 128 KB
• For T=8192: ~1 GB per request.

Llama-3 70B: L=80, h_kv=8, d_h=128:

• per token: 2 * 80 * 8 * 128 * 2 = 320 KB
• For T=128k, B=1: ~40 GB. The model's weights weigh 140 GB. KV is significant.

DeepSeek-V3 with MLA: KV memory ~10× lower per token thanks to latent compression. This is why MLA is a big deal for long-context serving.

3. Naive implementation

class KVCache:
    def __init__(self, n_layers):
        self.K = [None] * n_layers       # tensors (B, h_kv, T_cache, d_h)
        self.V = [None] * n_layers

    def update(self, layer, k_new, v_new):
        if self.K[layer] is None:
            self.K[layer], self.V[layer] = k_new, v_new
        else:
            self.K[layer] = torch.cat([self.K[layer], k_new], dim=2)
            self.V[layer] = torch.cat([self.V[layer], v_new], dim=2)
        return self.K[layer], self.V[layer]

This works for one request but is wasteful when you batch requests of different lengths — you have to pad to the max length, wasting memory.

4. PagedAttention (vLLM, Kwon et al., 2023)

Inspired by virtual memory in operating systems. Instead of one contiguous tensor per request, slice the cache into fixed-size pages (e.g., 16 tokens per page) and track which pages belong to which request.

Physical KV cache: a big pool of pages.
Per-request page table: list of page indices.
Attention kernel: gather K/V from indirected pages.

Benefits:

• Zero internal fragmentation: a 17-token request uses 2 pages (16 + 1, with 15 wasted), not max-length-padded memory.
• Zero external fragmentation: pages are uniform; allocator is O(1).
• Sharing: multiple requests sharing the same prefix can share pages — this is the basis of prefix caching.
• Throughput: enables concurrent requests of wildly different lengths in the same batch (continuous batching).

Empirically, vLLM achieves ~2-4× the throughput of naive HF generation thanks to PagedAttention + continuous batching.

You don't implement PagedAttention by hand — you use vLLM. But you should be able to explain it.

5. Continuous batching (a.k.a. iteration-level scheduling)

Traditional batching: gather B requests, run them together, return when all are done. The slowest one stalls the rest.

Continuous batching:

• Each scheduling step, the engine looks at all in-flight requests.
• Finished requests are removed; new ones are added.
• Per-step batch composition changes.

vLLM, TGI, SGLang all do this. Combined with PagedAttention, throughput gets close to "ideal" (machine fully utilised regardless of request mix).

6. Prefix caching

If many requests share a long system prompt:

[SYSTEM: You are a helpful assistant ... 2000 tokens]
[USER: Hello]                            ← 5 tokens

Naively each request re-processes the 2000-token prefix. With prefix caching, the engine:

1. Hashes the prefix tokens.
2. Stores the prefix's KV pages keyed by that hash.
3. On the next request with the same prefix, reuses those pages instead of recomputing.

Result: dramatic latency drop on shared-system-prompt apps.

vLLM, SGLang, and TGI all support it. In Anthropic's API it is exposed as cache_control blocks (Lesson 1.1 in course/).

vllm serve meta-llama/Llama-3.1-8B-Instruct --enable-prefix-caching

For 50%+ of API workloads (RAG, agents) this is the single biggest cost lever.

7. KV cache quantisation

bf16 KV is precise but costly. Cheaper alternatives:

• fp8 KV (E4M3): half the memory. 1-2% quality drop on long contexts. vLLM --kv-cache-dtype fp8 (Hopper).
• INT8 KV: similar; broader hardware support.
• INT4 / 2-bit KV: aggressive; only for low-stakes batched serving.

vllm serve meta-llama/Llama-3.1-8B-Instruct --kv-cache-dtype fp8

For long contexts (≥32k) this often unlocks 2× the concurrent users at acceptable quality.

8. Sliding-window cache and attention sinks

For models trained with sliding-window attention (Mistral 7B), the KV cache only needs to hold the last W tokens, not all of them:

KV size per request = 2 * W * L * d_kv * b

For W=4096, the cache is constant size regardless of sequence length. Cheap long-context.

Attention sinks (Xiao et al., 2024 — "StreamingLLM"): keep the first 4-8 tokens and the last W tokens. Improves quality at long generations because the leading tokens act as anchor positions in attention.

9. The big picture for cost / throughput

In modern serving, your throughput limit usually is KV cache memory, not weight memory. With H100 (80 GB) running Llama-3 70B (140 GB across 2 GPUs in TP=2):

• Model weights: ~70 GB per GPU.
• Free for KV: ~10 GB per GPU = ~20 GB total.
• At 320 KB/token, ~62k tokens of KV before OOM.
• That can be B=1, T=62k or B=10, T=6200 or anything in between.

→ Long-context jobs cap the number of concurrent requests.

→ Short-context jobs cap on compute, not memory.

→ FP8 KV doubles your concurrent capacity at long context. That is why every commercial endpoint advertises fp8 inference today.

Hands-on lab (3 hours)

kv_cache_lab.ipynb:

1. Re-implement KV-cached decoding for your nano-GPT (Lesson 2.4). Measure tokens/sec with and without cache for a 256-token generation.
2. Compute KV memory for Qwen2.5-7B-Instruct and meta-llama/Llama-3.1-70B-Instruct at T=32k, B=4, bf16. Verify against nvidia-smi while running vLLM.
3. Spin up vLLM:

   pip install vllm
   vllm serve Qwen/Qwen2.5-7B-Instruct --enable-prefix-caching

   Send the same 2000-token system prompt with 10 different user messages. Measure latency of the 1st vs 2nd-10th request.
4. Now turn off prefix caching. Compare.
5. Toggle --kv-cache-dtype fp8 (Hopper required); compare quality on a small eval and memory.
6. Bonus: implement attention sinks on your nano-GPT and show stable generation past block_size.

Common pitfalls

1. Forgetting to detach() the cache → the autograd graph stays alive across steps, blowing up memory.
2. Concatenating with torch.cat in a Python list — slow at long sequences. Use a pre-allocated buffer or PagedAttention via vLLM.
3. Mixing FP8 KV with a model that wasn't tested for it → noisy outputs at long contexts.
4. Setting --max-model-len too high without considering KV memory → OOM at high concurrency.
5. Sharing prefix-cache across users with different system prompts → information leakage. Always include user/tenant id in the cache key.

Self-check

1. Why does KV caching change generation from O(T^2) to O(T)?
2. What does PagedAttention do that torch.cat-style caches do not?
3. How much memory does a Llama-3 8B KV cache use per 8k context?
4. What does prefix caching speed up?
5. Cost of FP8 KV cache vs bf16?

References

• Kwon et al. (2023), "Efficient Memory Management for Large Language Model Serving with PagedAttention" (vLLM paper).
• Xiao et al. (2024), "Efficient Streaming Language Models with Attention Sinks."
• Pope et al. (2022), "Efficiently Scaling Transformer Inference."
• vLLM blog post.
• HuggingFace TGI docs.
• DeepSeek-AI (2024), "DeepSeek-V2 Technical Report" (MLA).