DL S066 · KV Cache Deep Dive

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 66 / 80 · Module M11 · ~2 hours

The story

Here's a puzzle. You fine-tuned a 7B Llama in Session 62. You load it on an A10 (24 GB), and generation runs at ~40 tok/s. Great. You push context length to 8k and suddenly you can only fit one request at a time, and adding a second user OOMs. Meanwhile your friend on an A100 (80 GB) with the same model handles 16 concurrent users at 80 tok/s per user. What did they do that you didn't?

They didn't do anything. The KV cache did. The KV cache is the single biggest reason LLM inference is memory-bound instead of compute-bound, the reason batching helps so much, and the reason PagedAttention (Kwon et al., SOSP 2023) got published as a full systems paper rather than a two-page workshop note. If you understand nothing else about inference, understand the KV cache: how big it is, why it grows linearly with sequence length, and why every clever inference trick invented between 2019 and 2025 — MQA, GQA, quantization, paging, prefix sharing, MLA — exists to shrink it.

Today we derive the memory formula from scratch, walk through the numbers for Llama-2-7B, Llama-3-70B, Llama-3.1-405B, and DeepSeek-V3, and write a working KV cache in ~50 lines of PyTorch. Tomorrow (S067) we'll shrink it with quantization; by S071 we'll page it. But first: the puzzle.

A little more of the story: how we got here

The KV cache is embarrassingly obvious in hindsight, which is why the original Attention Is All You Need paper (Vaswani et al., 2017) doesn't mention it. Vaswani's team was thinking about training, where every token in a sequence is visible at once and there's nothing to cache. It took the GPT-2 crowd (Radford et al., 2019) actually shipping an autoregressive text generator to notice: at decode time, K and V for token 5 do not change when you emit token 6. Recomputing them is O(T²) work for no reason.

The fix went into transformers as past_key_values around 2019 and quietly turned every LLM demo from "cute research artifact" into something you could put behind a web endpoint. But the cache also invented a new problem: memory pressure that scales with context length instead of parameter count. By 2023, when Anthropic's Claude 2 shipped 100k context and OpenAI teased GPT-4-turbo at 128k, the KV cache — not the weights — had become the dominant term in most serving budgets. Kwon and coauthors wrote PagedAttention specifically to defrag it. In 2024, DeepSeek-V2's Multi-head Latent Attention (MLA) attacked the problem from the model-architecture side, compressing the entire cache down to a shared low-rank latent — a 5–10× memory win. By 2025, every serious open-weights release either uses GQA aggressively (Llama 4, Gemma 3, Qwen 2.5) or invents its own KV shrinker (DeepSeek-V3, Jamba's SSM hybrid).

So when you read "KV cache" think the memory-bandwidth budget of the entire LLM industry. It's not a technicality; it's the single line item that decides whether your model fits on a $2/hr GPU or needs a$30/hr one.

Prerequisites

• Session 044 — Q, K, V projections and multi-head attention.
• Session 052 — causal masking and autoregressive decoding.
• Session 060 — parameter counts for Llama-family models.

1 · Why a cache exists at all

Recall autoregressive decoding: to produce token $t$, we run attention over tokens $1..t-1$. To produce token $t+1$, we run attention over tokens $1..t$. Naively, that means at step $t+1$ we recompute the K and V projections for every prior token, even though nothing about them has changed.

The insight is embarrassingly simple: cache K and V from prior steps, only compute Q for the new token. Attention becomes $O(t)$ per step instead of $O(t^2)$. Total generation cost drops from $O(T^3)$ to $O(T^2)$.

But that cache has to live somewhere. And on modern GPUs, "somewhere" is HBM (High Bandwidth Memory), and HBM is the bottleneck.

Picture-in-words: the cache as a scrolling notebook

Imagine you're a courtroom stenographer transcribing an infinite trial. Every new spoken word (query) you have to compare against every prior word in the record to decide what it means. Without a notebook, you'd have to re-listen to the whole trial for each new word — O(T²) per word, O(T³) for the whole day. The KV cache is your notebook: you write each word down once (one K and one V vector per token per layer), and every new word just skims what's already on the page.

The catch: the notebook grows one page per word, per layer. A 32-layer model taking a 4k-token deposition ends up with 128k "pages" of stenographer notes. Those pages have to live in your desk drawer (HBM). When the drawer fills up, you can't accept new clients. That's an OOM.

2 · The memory formula, derived

Take a Transformer with:

• $L$ layers
• $n_h$ attention heads, each of dim $d_{head}$ (so model dim $d = n_h \cdot d_{head}$)
• Sequence length $T$
• Batch size $B$
• Precision bytes per element (2 for fp16/bf16, 1 for int8, 0.5 for int4)

At each layer, we store both K and V. Each is shape (B, n_h, T, d_head). So per layer:

$\text{KV per layer} = 2 \cdot B \cdot n_h \cdot T \cdot d_{head} \cdot \text{bytes}$

Across all layers:

$\boxed{\text{KV total} = 2 \cdot L \cdot B \cdot n_h \cdot T \cdot d_{head} \cdot \text{bytes}}$

Since $n_h \cdot d_{head} = d$, we can also write it as $2 \cdot L \cdot B \cdot T \cdot d \cdot \text{bytes}$. Same thing.

Tiny worked example you can do in your head

Before the real models, do it on a toy. Suppose $L=2$ layers, $n_h=2$ heads, $d_{head}=4$, sequence $T=3$, batch 1, fp16 (2 bytes):

$2 \cdot 2 \cdot 1 \cdot 2 \cdot 3 \cdot 4 \cdot 2 = 96 \text{ bytes.}$

Write out the actual tensors. Layer 0's K cache is a $2 \times 3 \times 4$ matrix — 24 fp16 numbers. Layer 0's V is another 24. Layer 1 the same. Total 96 numbers × 2 bytes = 96 bytes. That $2 \cdot L$ prefactor is not magic: it's just "K and V, at every layer." Everything after this is scaling that toy up to production model dimensions.

Worked example: Llama-2-7B

Config: $L=32$, $d=4096$, $n_h=32$, $d_{head}=128$, fp16 (2 bytes). Single request $B=1$, context $T=4096$.

$\text{KV} = 2 \cdot 32 \cdot 1 \cdot 4096 \cdot 4096 \cdot 2 = 2{,}147{,}483{,}648 \text{ bytes} = 2.0 \text{ GiB}$

Per token: $2 \cdot 32 \cdot 4096 \cdot 2 = 524{,}288$ bytes $= 0.5$ MiB per token.

So on a 24 GB A10 holding the 14 GB fp16 weights, you have $\approx 10$ GB left for KV. That's $10 / 0.0005 = 20{,}000$ tokens of total KV budget. Sounds like a lot until you realise 8 concurrent users at 2500 tokens each is your ceiling.

Worked example: Llama-3-70B

$L=80$, $d=8192$, $n_h=64$, $d_{head}=128$, but GQA with $n_{kv}=8$ (this is critical, see §4).

Per token, with GQA the formula becomes $2 \cdot L \cdot n_{kv} \cdot d_{head} \cdot \text{bytes} = 2 \cdot 80 \cdot 8 \cdot 128 \cdot 2 = 327{,}680$ bytes = 0.31 MiB.

Without GQA it would be $2 \cdot 80 \cdot 64 \cdot 128 \cdot 2 = 2.6$ MiB per token — 8× larger. GQA is the entire reason Llama-3-70B is serveable.

3 · Why prefill is compute-bound and decode is memory-bound

Two phases:

• Prefill: you have a $T_p$-token prompt. You run the model once over all $T_p$ tokens in parallel. Compute cost $\sim T_p \cdot 2N$ FLOPs where $N$ is param count. Memory reads: model weights, once. Ratio FLOPs/bytes is high → compute-bound.
• Decode: you generate token by token. Each step, you read all model weights (14 GB for 7B) and the full KV cache (grows with each token) just to produce one token. Compute cost $\sim 2N$ FLOPs; memory read $\sim N \cdot \text{bytes} + \text{KV}$. Ratio is terrible → memory-bound.

On an A100 (312 TFLOPs fp16, 2 TB/s HBM), the arithmetic intensity break-even is ~156 FLOPs/byte. Prefill hits it easily. Single-request decode is at ~2 FLOPs/byte. You are using 1% of your GPU's compute.

4 · MHA vs MQA vs GQA

Standard multi-head attention (MHA): $n_h$ query heads, $n_h$ key heads, $n_h$ value heads. KV cache scales with $n_h$.

Multi-Query Attention (MQA) — Shazeer, 2019: keep $n_h$ query heads but only one shared K and one shared V. KV cache shrinks by $n_h$×. Quality drops slightly.

Grouped-Query Attention (GQA) — Ainslie et al., 2023: split query heads into $G$ groups, each group shares one K and one V. So $n_{kv} = G$. Interpolates between MHA ($G = n_h$) and MQA ($G = 1$).

Llama-2 (7B, 13B) uses MHA. Llama-2-70B and all Llama-3 use GQA with $n_{kv}=8$. Mistral uses GQA. Gemma uses MQA in some sizes.

How to tell from a HuggingFace config: look for num_key_value_heads. If it equals num_attention_heads, it's MHA. If it's 1, MQA. Otherwise GQA.

5 · Code: a KV-cached attention forward pass

Let's write it. First the naive (no-cache) version so you can see the difference:

import torch
import torch.nn.functional as F
 
def attention_no_cache(x, W_q, W_k, W_v, W_o, n_heads):
    B, T, D = x.shape
    d_head = D // n_heads
    q = (x @ W_q).view(B, T, n_heads, d_head).transpose(1, 2)  # (B,H,T,d)
    k = (x @ W_k).view(B, T, n_heads, d_head).transpose(1, 2)
    v = (x @ W_v).view(B, T, n_heads, d_head).transpose(1, 2)
    mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
    scores = (q @ k.transpose(-2, -1)) / (d_head ** 0.5)
    scores = scores.masked_fill(~mask, float("-inf"))
    attn = F.softmax(scores, dim=-1)
    out = (attn @ v).transpose(1, 2).reshape(B, T, D)
    return out @ W_o

Every generation step, x is the entire history so far and we recompute all K, V for all tokens. Wasteful.

Now the cached version. Key change: x is just the new token(s), and we append its K, V to a rolling cache.

class KVCache:
    def __init__(self, max_batch, max_seq, n_kv_heads, d_head, dtype=torch.float16, device="cuda"):
        shape = (max_batch, n_kv_heads, max_seq, d_head)
        self.k = torch.zeros(shape, dtype=dtype, device=device)
        self.v = torch.zeros(shape, dtype=dtype, device=device)
        self.len = 0  # how many tokens are populated
 
    def append(self, k_new, v_new):
        # k_new, v_new: (B, n_kv, T_new, d_head)
        T_new = k_new.shape[2]
        self.k[:, :, self.len:self.len + T_new] = k_new
        self.v[:, :, self.len:self.len + T_new] = v_new
        self.len += T_new
        return self.k[:, :, :self.len], self.v[:, :, :self.len]

Notice the preallocation to max_seq. In production you rarely know this in advance — that's exactly the problem PagedAttention (S071) solves.

The cached attention:

def attention_cached(x_new, W_q, W_k, W_v, W_o, n_heads, n_kv_heads, cache):
    B, T_new, D = x_new.shape
    d_head = D // n_heads
    q = (x_new @ W_q).view(B, T_new, n_heads, d_head).transpose(1, 2)
    k_new = (x_new @ W_k).view(B, T_new, n_kv_heads, d_head).transpose(1, 2)
    v_new = (x_new @ W_v).view(B, T_new, n_kv_heads, d_head).transpose(1, 2)
 
    k, v = cache.append(k_new, v_new)  # k,v: (B, n_kv, T_total, d_head)
 
    # GQA: expand kv heads to match query heads (view, not copy)
    if n_kv_heads != n_heads:
        rep = n_heads // n_kv_heads
        k = k.repeat_interleave(rep, dim=1)
        v = v.repeat_interleave(rep, dim=1)
 
    scores = (q @ k.transpose(-2, -1)) / (d_head ** 0.5)
    # causal mask: q at positions [len - T_new, len), k at [0, len)
    T_total = cache.len
    q_pos = torch.arange(T_total - T_new, T_total, device=q.device)[:, None]
    k_pos = torch.arange(T_total, device=q.device)[None, :]
    mask = (q_pos >= k_pos)
    scores = scores.masked_fill(~mask, float("-inf"))
    attn = F.softmax(scores, dim=-1)
    out = (attn @ v).transpose(1, 2).reshape(B, T_new, D)
    return out @ W_o

During prefill: x_new is the whole prompt, cache.len starts at 0. During decode: x_new is a single token, T_new = 1.

6 · Estimating max batch size on your GPU

The napkin math every serving engineer does daily:

GPU memory available for KV = total_HBM − weight_bytes − activations − CUDA_overhead
                           ≈ total_HBM − weight_bytes − ~2 GiB
 
max_tokens_across_all_requests = KV_budget / per_token_kv_bytes
 
max_batch_size ≈ max_tokens / avg_seq_len

Llama-3-8B on A100 (80 GB), fp16, 4k context:

• Weights: $8 \cdot 10^9 \cdot 2 = 16$ GiB
• Overhead: ~4 GiB
• KV budget: ~60 GiB
• Per-token KV (GQA, $n_{kv}=8$, $L=32$, $d_{head}=128$): $2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 128$ KiB
• Max tokens: $60 \text{ GiB} / 128 \text{ KiB} \approx 490{,}000$
• Max batch at 4k avg: $490{,}000 / 4096 \approx 120$

That's why vLLM benchmarks quote absurd numbers like "1000 req/min on one A100" — the KV cache is small enough to batch a lot.

Worked example: Llama-3.1-405B on 8×H100

The 405B model is the interesting stress test because a single H100 (80 GB) can't even hold the weights. You have to tensor-parallel across 8 GPUs.

• Weights: $405 \cdot 10^9 \cdot 2 = 810$ GiB fp16, or ~101 GiB per GPU at TP=8.
• $L=126$, $d=16384$, $n_h=128$, $d_{head}=128$, GQA with $n_{kv}=8$.
• Per-token KV (globally, before sharding): $2 \cdot 126 \cdot 8 \cdot 128 \cdot 2 = 516{,}096$ bytes ≈ 504 KiB.
• With tensor parallelism, KV heads shard across GPUs too, so per-GPU per-token: 504 KiB / 8 = 63 KiB.
• Per GPU KV budget: 80 GiB − 101 GiB … wait, that's negative.

See the problem? 405B doesn't fit even just weights on 8 H100s at fp16 — you need fp8 (Meta ships an official fp8 checkpoint) or 16 H100s. Meta's own blog post on serving Llama 3.1 405B walks through this: they used fp8 weights + fp8 KV cache to fit on a single 8×H100 node.

In fp8: weights are ~405 GiB / 8 = 50 GiB per GPU, leaving ~25 GiB for KV per GPU. At 63 KiB per-token per-GPU (fp16) → 32 KiB fp8, that's 25 GiB / 32 KiB ≈ 800k tokens of headroom per GPU. Suddenly 32 concurrent 8k-context users is trivial. This is why the whole industry moved to fp8 KV in 2024–2025.

Worked example: DeepSeek-V3 with MLA

DeepSeek-V2 (May 2024) and V3 (Dec 2024) do something completely different. Instead of storing K and V, they store a compressed latent of dimension $d_c \approx 512$ and reconstruct K, V from it on the fly with a learned matrix.

From the DeepSeek-V3 tech report (§2.1): per token, MLA stores only $d_c + d_{rope} \approx 512 + 64 = 576$ elements per layer. Compare to GQA-Llama-3 at $2 \cdot 8 \cdot 128 = 2048$ elements per layer. That's a 3.5× additional compression on top of what GQA already did.

DeepSeek-V3 has 61 transformer layers. Per-token KV: $61 \cdot 576 \cdot 2 \text{ (bf16)} = 70{,}272$ bytes ≈ 68 KiB. A 671B-parameter model with per-token KV smaller than 8B Llama-3's. That is the entire pitch of MLA.

The tradeoff: MLA needs one extra matmul per layer to decompress. Prefill gets a bit slower; decode is dominated by memory reads anyway, so it wins net-net. Expect MLA-style compression to become standard in 2026 open releases.

Worked example: fp8 KV cache math

Since Hopper (H100) shipped in 2023 with native fp8 tensor cores, you can store the cache in float8_e4m3 (4 bits exponent, 3 bits mantissa) at half the memory of bf16. NVIDIA's TensorRT-LLM 0.14 release notes show fp8 KV cache is now default for H100 deployments.

Rule of thumb: fp8 KV loses ~0.05 in MMLU vs bf16 KV — well below noise floor. fp8 weights + KV together is the current 2025 sweet spot. Quality holds; memory halves; throughput doubles because bandwidth halves too.

Further reading — the KV memory arms race:

• Ainslie et al., GQA (2023) — the technique in half of production LLMs today.
• DeepSeek-AI, DeepSeek-V2 technical report (2024) — MLA in §2.
• Meta AI, The Llama 3 Herd of Models (2024) — §5.4 on inference optimization.
• NVIDIA, FP8 KV cache in TensorRT-LLM — practical config knobs.

Llama-2-7B (MHA!) on A10 (24 GB), fp16, 4k context:

• Weights: 14 GiB
• Overhead: ~2 GiB
• KV budget: ~8 GiB
• Per-token KV (MHA, $L=32$, $d=4096$): $2 \cdot 32 \cdot 4096 \cdot 2 = 512$ KiB
• Max tokens: $8 \text{ GiB} / 512 \text{ KiB} = 16{,}000$
• Max batch at 4k avg: 4.

Four users. On a $1.5/hr GPU. This is why quantization (S067) and paging (S071) are not optional.

7 · Mermaid: prefill vs decode

8 · Pitfalls

9 · KV eviction and sparse attention — when the cache is too big to keep

Even with GQA and fp8, a 1M-context Gemini-style request has a KV cache that dwarfs the weights. Two families of tricks have emerged in 2024–2025 to keep only the useful KV entries around.

9.1 Attention Sinks and StreamingLLM

Xiao et al. (Attention Sinks, ICLR 2024) observed something weird: if you evict the oldest tokens in a sliding-window KV cache, quality collapses even though those tokens seem irrelevant. Why? Because the softmax normaliser needs a "dumping ground" — a few tokens that always get attention mass so the rest can be legitimately small. In practice, the first 2–4 tokens of any generation are these sinks.

Fix: keep the first $k$ (sink) tokens and the last $W$ (recent) tokens; evict the middle. Now you have unbounded context at fixed memory. This is exactly how attn_sinks=True in vLLM works.

9.2 H2O and heavy-hitter eviction

Zhang et al. (H2O: Heavy-Hitter Oracle, NeurIPS 2023) go further: track which past tokens actually receive attention mass across recent decode steps, and evict the ones that don't. Empirically only ~20% of KV entries do real work; the rest can be dropped with <1% quality loss.

9.3 Sparse attention comes back

Longformer (Beltagy 2020) and BigBird (Zaheer 2020) pre-designed which tokens attend to which. Their descendants — Mamba (Gu 2023), RWKV-6, Jamba (AI21 2024, Mamba+Transformer hybrid), Griffin (Google DeepMind 2024) — replace attention with a state-space model or linear recurrence. The KV cache disappears entirely; a fixed-size hidden state replaces it. Serving cost becomes truly O(1) per token.

2025 status: pure SSMs still lose ~2 MMLU points to attention transformers of the same size on hard reasoning tasks. Hybrids (Jamba, Zamba, Griffin) close the gap and are gaining share for long-context serving. If you're serving 100k+ context in 2026, evaluate a hybrid before you evaluate more RAM.

9.4 Prefix caching — sharing KV across requests

The cheapest KV entry is the one you don't recompute. If two users send "You are a helpful assistant." as their system prompt, why compute K/V for those tokens twice? SGLang's RadixAttention (Zheng et al., 2024) organises the KV cache as a radix tree keyed by token prefixes; identical prefixes share physical KV blocks with copy-on-write for divergent suffixes.

vLLM added --enable-prefix-caching as a first-class flag in v0.5 (mid-2024) and made it default in v1 (Jan 2025). On chat workloads with long system prompts, this alone is a 30–50% throughput uplift with zero quality change. Read the vLLM V1 blog post if you serve chat at scale — it changed the default scheduler in ways that break some older kernels.

10 · The 2025 KV cache scoreboard

Snapshot of what's shipping in production as of mid-2025. Numbers are per-token KV in bytes for a single request, fp16/bf16 unless noted.

Read top to bottom: 6 years of research compressed the per-token KV footprint by roughly 10×. The next order of magnitude is coming from architectural changes (MLA, SSM hybrids), not precision.

Further reading — the frontier:

• Xiao et al., StreamingLLM (ICLR 2024).
• Zhang et al., H2O: Heavy-Hitter Oracle.
• AI21 Labs, Jamba technical report (2024).
• Gu & Dao, Mamba (2023).
• Zheng et al., SGLang / RadixAttention (2024).

11 · Try it yourself

Open a notebook. Do these in order:

1. Read a config, predict the KV size. Pick three models from HuggingFace (meta-llama/Llama-3-8B, mistralai/Mistral-7B-v0.3, Qwen/Qwen2.5-7B). Load their configs, compute per-token KV bytes for bf16, and rank them small to large. Then check against the scoreboard above.

2. Instrument the naive version. Take the attention_no_cache function above; wrap it in a generation loop that produces 128 tokens; time it with torch.cuda.synchronize(). Now do the cached version. Confirm the naive version scales quadratically and the cached one linearly.

3. Break RoPE on purpose. In the cached version, replace position_ids = torch.tensor([cache.len - 1]) with position_ids = torch.tensor([0]). Watch the output degrade. Fix it and diff. That's what pitfall #1 looks like from the driver's seat.

4. Estimate a serving quote. Your friend wants to serve Qwen2.5-32B on 2×A100 (80GB) at 8k context with 20 concurrent users. Do they fit? Show the arithmetic.

Recall — before you close the tab

Stretch: derive the FLOPs/byte ratio for decode of Llama-3-8B on an A100. Show the arithmetic. When does batching push you into the compute-bound regime?

In your own words: the KV cache is __________________________________________________________.

Spaced review: revisit S044 Transformer block, S052 GPT decoder, and S045 positional encodings — RoPE + KV-cache interaction is where most bugs live.

Next session (S067): we shrink the weights and the cache with INT8 and INT4 quantization. GPTQ, AWQ, and why bitsandbytes is the "just make it fit" button.

Bring back tomorrow:

• The KV-size formula (memorised).
• The distinction between prefill (compute-bound) and decode (memory-bound).
• Your napkin estimate for one model+GPU pair you actually care about.

Previous: ← DL S065 · Next: DL S067 →