Across which axes is the KV cache stored, per layer, per head, per batch?
The KV cache spans all three axes: per layer, per KV head, and per batch request. Each combination stores its own K and V tensor over the sequence.
Imagine a multistory library where every floor (a transformer layer) has its own card catalog. On each floor, several librarians (the heads) keep separate drawers of cards. And every visitor (a request in the batch) brings their own set of bookmarks. To find a book, you need the right floor, the right librarian's drawer, and your personal bookmark stack. The KV cache works the same way: layers, heads, and requests are all separate axes, and any one of them being missed will make your memory math wrong by an order of magnitude.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
The KV cache is the dominant memory term in modern LLM inference, and getting its shape right is the difference between a serving deployment that holds 32 concurrent requests and one that holds 3. The shape question reduces to one observation: how many independent axes does the cache span?
The answer is three: per layer, per KV head, and per request in the batch. Each axis contributes a multiplicative factor to the total memory, and each compresses under different conditions. Layers are structural and uncompressed in production. Heads compress via the GQA, MQA, and MLA family. Batch requests can share prefixes via PagedAttention but otherwise remain independent.
This deep dive walks each axis, derives the per-token memory formula, shows why the modern architectural shifts (GQA, MLA, paged allocation) are all responses to one axis or another, and works through a sizing example for a 70B-class deployment.
The three axes, named and shaped
The cache stores K and V tensors. Per layer, the conceptual shape is (batch, num_kv_heads, seq_len, head_dim). Across the network there are L such shapes, two each (one K, one V).
Axis 1: layer
Every transformer block has its own attention sublayer with its own W_K, W_V projections. Layer 12 sees a different residual stream from layer 1, so K and V at layer 12 are unrelated to K and V at layer 1. The cache must hold all L copies, this is the L factor in the memory formula.
Axis 2: KV head
Within a single layer, attention has multiple heads. The number of KV heads (which may be smaller than the number of query heads under GQA or MQA) sets how many K, V slices the layer holds. Under MHA, num_kv_heads = num_heads. Under GQA, num_kv_heads is a small group count. Under MQA, it is 1.
Axis 3: batch request
Different requests have different prompts, so they produce different K and V. The cache spans the batch dimension. Prefix sharing via PagedAttention is the one production technique that lets multiple requests reference the same physical pages for a shared prompt prefix.
All three axes are real and multiplicative. Missing one of them moves your memory estimate by an order of magnitude.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- HuggingFace transformers' `past_key_values` is a tuple of length L, each entry holding K and V of shape (batch, H_kv, T, d_h).
- vLLM's PagedAttention treats this cache as paged memory with per-request page tables and prefix sharing.
What an interviewer would ask next. Try answering before peeking at the approach.
QHow does PagedAttention's prefix sharing change the batch-axis story?
Requests sharing the same system prompt point at the same physical pages for the prefix and only allocate new pages for their private suffix. Reference counting handles eviction. The batch axis logically exists but the physical storage partially collapses.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Pretending the cache is per layer only, then sizing memory as if heads or batch share storage. Each axis is genuinely independent and contributes a multiplicative factor.
60 second bullets to scan on the way to the call.
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.