The KV Cache & Memory Math

Weights are the part of an LLM's memory footprint that everyone knows about. The part that actually limits how many users you can serve is another one: the KV cache. This chapter shows what it is and how to calculate it on the back of an envelope, which is a question that comes up in nearly every serving interview.

Imagine a note-taker at a long meeting. Each time someone speaks, the note-taker writes down the key points of that statement. When a new question arrives, they answer it by looking over their notes, not by replaying the entire recording. The notes take space, and the longer the meeting, the more space. The KV cache is that set of notes.

In attention, every new token compares itself with all earlier tokens, using a key and a value that each earlier token produced at each layer. Those keys and values for old tokens never change. Recomputing them on every step would repeat a huge amount of work, so the server stores them. Prefill fills the cache for the prompt, and each decode step adds one more entry to it, per layer.

The Llama 3 paper lists 8 key-value heads for its 8B and 70B models, even though they have 32 and 64 query heads. Sharing a KV head among several query heads is called grouped-query attention, and it shrinks the cache. The 8B model has 32 layers and the 70B model has 80, both with a head size of 128. With 16-bit values, 2 bytes each, that gives 2 × 32 × 8 × 128 × 2 = 131,072 bytes, or 128 KiB, per token for the 8B model, and 320 KiB per token for the 70B model.

The code turns that into capacity. It takes an 8B model on one 80 GB GPU and a 70B model on four, keeps 10% of memory free, subtracts the weights, and divides what is left by the KV cache of one sequence at different context lengths. Press Run and read the last column, then the final line.

On one GPU, the 8B model has about 56 GB left after its 16 GB of weights. At a context of 8,192 tokens, one sequence needs 1.07 GB, so about 52 fit. Stretch the context to 32,768 tokens and only 13 fit. At 131,072 tokens, just 3. The 70B model on four GPUs has 147 GB free and shows the same pattern: 54 sequences at 8,192 tokens, 3 at 131,072. The last line shows why grouped-query attention matters. Without it, the 70B model would need 2,560 KiB per token, eight times more, and its capacity would shrink to a fraction.

Use the demo below to try other models and settings. Change the context length, the number of GPUs, or how the numbers are stored, and see how many conversations fit once the weights are loaded.

• -Weights: fixed, and the first thing you subtract.
• -KV cache: grows with every token of every active request, so it varies with load and context length.
• -Activations and overhead: temporary buffers, kept as a margin of safety.

A simple server reserves one large contiguous block per request, sized for the longest possible answer. Most requests use far less, so most of the reservation is wasted, and memory fragments. The vLLM paper by Kwon and colleagues proposed PagedAttention, which stores the cache in small fixed-size pages, like an operating system does with memory, so nothing is reserved ahead of need and pages can be shared. They report 2 to 4 times higher throughput than earlier systems at the same latency. The earlier chapter on serving covers this in more detail, and the chapter on quantization shows how to shrink the cache itself by storing keys and values in fewer bits.