The KV Cache Explained: Why LLM Decoding Is Memory-Bound

A modern GPU can perform hundreds of trillions of arithmetic operations every second. Yet when you chat with a language model on it, the words trickle out one at a time. Why would a machine that fast feel slow? The answer has almost nothing to do with arithmetic.

• Huge compute
• Yet one token at a time

Today is episode five: the KV cache. We will see why models remember keys and values instead of recomputing them, derive exactly how much memory that costs, and then show that decoding is limited by memory bandwidth, not by compute. Last time we derived rotary position embeddings and context windows. Now we ask what it costs to actually use that context.

• Cache size formula
• Roofline argument
• Mitigations and code

Imagine writing a story one word at a time, and before every new word you must reread the whole story from the first line, rebuilding your understanding from scratch. That is exactly what a naive language model does. The past has not changed, so recomputing it is pure waste.

• Past tokens never change
• Recomputing is waste

Let us recall attention from the ML course. Each token vector is projected into a query, a key and a value. The output is softmax of Q times K transposed, divided by the square root of the head dimension, then multiplied by V. Each row of the score matrix says how much one token looks at every other token.

• Queries, keys, values
• softmax(QK^T / sqrt d) V

Now the key observation. Generation is causal: token number t only attends to tokens one through t. So the key and value of an earlier token depend only on that token and the tokens before it. Once computed at a given layer, they never change, no matter how many words come after.

• Token t sees tokens 1 to t
• Old keys and values are fixed

So we store them. For every layer and every attention head, we keep the keys and values of all past tokens in memory. This stored collection is the KV cache. At each new step we compute the query, key and value for only the newest token, append its key and value to the cache, and attend over everything stored.

• Store K and V per layer, per head
• Append one entry per step

How much work does this save? Say we generate n tokens. Without a cache, step t reprocesses t tokens, so the total is one plus two plus up to n, which is n times n plus one over two. For n equal to one thousand, that is about five hundred thousand token passes. With a cache we process just one new token per step: one thousand passes. That is five hundred times less.

• No cache: n(n+1)/2 passes
• Cache: n passes
• n = 1000: 500x fewer

This splits inference into two phases. In prefill, the model reads your whole prompt at once. All prompt tokens go through the network in parallel, and their keys and values are written into the cache. The model also produces the distribution for the first new token.

• Whole prompt in parallel
• Fills the cache
• Matrix times matrix

Then comes decode. The model generates one token, feeds it back in, generates the next, and so on. Each step is strictly sequential, because token t plus one needs token t. Each step processes a single token, reading the entire cache and all of the model weights to do it.

• One token per step
• Strictly sequential
• Reads weights and cache

These two phases feel very different. Prefill is where you wait for the first token, often called time to first token. It uses the hardware well because there is lots of parallel work. Decode determines the speed of the stream after that, the tokens per second. As we will see, decode uses the hardware badly, and that is where our story lives.

• Prefill sets time to first token
• Decode sets tokens per second

Now the cost. Let us derive the cache size piece by piece. For one token, in one layer, in one key-value head, we store a key vector and a value vector. Each has length head dimension, which we call d sub h. That is two times d sub h numbers.

• Key vector: d_h numbers
• Value vector: d_h numbers

Multiply by the number of key-value heads, h sub kv, because each head keeps its own pair. Multiply by the number of layers, L, because every layer has its own cache. Now multiply by the sequence length, S, since every token is stored. Multiply by bytes per number, b, and by batch size, B, because each sequence in the batch has its own private cache.

• Heads, layers, length
• Bytes per number
• Batch size

The result: KV cache bytes equals two times L times h sub kv times d sub h times S times b times B. The two is for keys and values. Notice every factor is a plain multiplication. Memory grows linearly in sequence length and linearly in batch size. Remember that: it will matter a lot.

• 2 x L x h_kv x d_h x S x b x B
• Linear in S and in B

Let us plug in numbers for a Llama-style seven billion parameter model. These are assumed, typical values: thirty-two layers, thirty-two attention heads, and in this classic multi-head design thirty-two key-value heads too, with head dimension one hundred twenty-eight. We store in fp16, which is two bytes per number.

• 32 layers, 32 KV heads
• Head dim 128, fp16

First the cost per token. Two times thirty-two layers is sixty-four. Times thirty-two heads is two thousand forty-eight. Times one hundred twenty-eight is two hundred sixty-two thousand one hundred forty-four. Times two bytes is five hundred twenty-four thousand two hundred eighty-eight bytes. That is half a mebibyte for every single token.

• 2 x 32 x 32 x 128 x 2
• = 524,288 bytes per token

Now scale up. A four thousand ninety-six token context at batch size one is four thousand ninety-six times half a mebibyte: two gibibytes. Batch size sixteen means thirty-two gibibytes. The model weights themselves are seven billion parameters times two bytes, about fourteen gigabytes. So the cache at that batch size is more than twice the size of the model.

• 4096 tokens, batch 1: 2 GiB
• Batch 16: 32 GiB
• Weights: about 14 GB

Capacity is only half the problem. The cache must also be read. Each decode step reads every weight and the whole cache to produce a single token. So the question becomes: how fast can the GPU move bytes, compared with how fast it can do math? To answer it we need the idea of arithmetic intensity.

• Bytes must travel to compute
• Two limits: bandwidth, FLOPs

Arithmetic intensity is the number of floating point operations performed per byte loaded from memory. High intensity means the chip is busy computing. Low intensity means it spends its time waiting for data. It is the ratio that decides which limit you hit first.

• FLOPs per byte loaded
• Low means waiting on memory

Compute it for decode. Take one weight matrix W of size d by d, and one token vector x. The product W times x needs d squared multiplies and d squared adds, so about two d squared FLOPs. But we must load all d squared weights, at two bytes each in fp16, which is two d squared bytes. Divide: two d squared over two d squared is one. About one FLOP per byte.

• FLOPs: 2 d squared
• Bytes: 2 d squared
• About 1 FLOP per byte

Why so low? Because every weight is loaded, used for a single multiply and add, and then thrown away. There is no reuse. Compare that with prefill, a matrix times a matrix with, say, a thousand prompt tokens: each loaded weight gets used a thousand times, so intensity is roughly a thousand FLOPs per byte.

• Decode: each weight used once
• Prefill: reused per prompt token

The roofline model puts this on one picture. The horizontal axis is arithmetic intensity, the vertical axis is achievable performance. On the left, performance climbs along a slope equal to memory bandwidth. Then it hits a flat ceiling set by the peak compute of the chip. The corner is called the ridge point.

• Slope: memory bandwidth
• Ceiling: peak compute
• Ridge point in between

Use assumed round numbers for an illustrative accelerator: two hundred teraflops of fp16 compute, and one terabyte per second of memory bandwidth. Real chips differ, so check your own datasheet. The ridge point is two hundred trillion divided by one trillion: two hundred FLOPs per byte. To be compute-bound, a workload needs about two hundred FLOPs for every byte it loads.

• 200 TFLOPs, 1 TB/s assumed
• Ridge point: 200 FLOP/byte

Where does decode land? At about one FLOP per byte, far to the left of two hundred. It sits on the sloped part of the roofline, with performance of one terabyte per second times one FLOP per byte, one teraflop. That is half a percent of the assumed peak. The compute units are idle more than ninety-nine percent of the time, waiting for memory.

• Intensity about 1
• About 0.5 percent of peak
• Memory-bound

This gives a simple rule for speed. Each token requires reading all the model bytes once, so tokens per second is at most bandwidth divided by model bytes. With one terabyte per second and fourteen gigabytes of fp16 weights, that is about seventy-one tokens per second at batch size one. Add the two gigabyte cache from before, and the bytes per step are about sixteen, so roughly sixty-two tokens per second. This is an upper bound, real systems land lower.

• tok/s ≈ bandwidth / bytes
• 1000 GB/s / 14 GB ≈ 71
• With cache: about 62

Notice what this predicts. A GPU with twice the FLOPs but the same bandwidth decodes no faster. Halving the model bytes, with eight-bit weights, nearly doubles speed. And batching is nearly free: if we serve sixteen users, the weights are read once for all of them, so intensity rises toward sixteen and throughput climbs. Only the cache, which is private per user, does not get shared.

• More FLOPs: no help
• Fewer bytes: faster
• Batching shares weight reads

So we want a smaller cache. Look at our formula: the factor h sub kv is a design choice. In multi-query attention, all query heads share a single key-value head. That cuts the cache by the number of heads, thirty-two times in our example, at some risk to quality.

• All queries share one KV head
• Cache shrinks by head count

Grouped-query attention is the middle path, used in many recent open models. Queries are split into groups, and each group shares one key-value head. Take eight key-value heads instead of thirty-two. The per-token cost becomes two times thirty-two times eight times one hundred twenty-eight times two bytes: one hundred twenty-eight kibibytes, a quarter of before. Four thousand ninety-six tokens now take only five hundred twelve mebibytes.

• 8 KV heads instead of 32
• 128 KiB per token
• 4096 tokens: 512 MiB

Two more ideas, by name. Sliding window attention lets each token attend to only the most recent window of tokens, so the cache stops growing past the window size. And cache quantization stores keys and values in eight or even four bits instead of sixteen, shrinking bytes directly. Both trade some accuracy for memory. We will go deeper on serving techniques in a later lesson.

• Sliding window caps cache length
• Quantized cache: 8 or 4 bits

Now let us see it in code with PyTorch and Hugging Face. We load a small causal language model and tokenize a prompt. First comes the slow path: a loop where at every step we feed the entire growing sequence into the model, take the logits of the last position, pick the argmax token, and append it. No state is kept between steps.

• Feed full sequence each step
• Take last-position logits

Now the cached path. We run the prompt once with use cache set to true. The model returns logits and a past key values object. In every later step we pass only the newest token id plus that past key values, and get back an updated cache. Each step now processes one token, and the loop is otherwise identical.

• use_cache=True on the prompt
• Pass only the new token
• Reuse past_key_values

To time them fairly, we generate the same number of tokens in both loops, call torch cuda synchronize before reading the clock because GPU work is asynchronous, and use the same greedy decoding so outputs should match. Do check that the token ids are identical; if they are not, you have a bug, not a speedup.

• Synchronize before the clock
• Greedy so outputs match
• Assert identical tokens

What should you see? Without the cache, time per token rises as the sequence grows, so total time grows roughly quadratically. With the cache, time per token stays nearly flat, so total time grows about linearly. The gap widens with length. I will not quote specific timings, since they depend entirely on your hardware, so run it and plot your own curve.

• No cache: roughly quadratic
• Cache: roughly linear

Finally, compute the memory. Read the config: number of hidden layers, number of key-value heads, and the head dimension, which is hidden size divided by attention heads. Plug into our formula in a tiny function, with bytes per element from the tensor dtype. You can also sum the element sizes of every tensor in past key values and check that your formula matches.

• Read layers, KV heads, head dim
• Formula versus measured tensors

The first mistake is forgetting that the cache grows linearly with both batch and length. A server that fits sixteen short chats comfortably can run out of memory when the same users paste long documents. Plan memory for your worst-case length times your batch size, not for the average conversation.

• Cache scales with length and batch
• Plan for worst case

Second mistake: sizing the cache with attention heads when the model uses grouped-query attention. The formula needs key-value heads, which may be several times fewer. Third: assuming a faster-FLOPs GPU means faster decoding. Check memory bandwidth first. And fourth: using a cache while silently changing the prompt prefix, which invalidates the stored keys and values.

• Use KV heads, not query heads
• Check bandwidth, not FLOPs
• Changed prefix breaks the cache

Let us recap. Past keys and values never change, so we cache them instead of recomputing. Prefill processes the prompt in parallel; decode generates one token at a time. The cache costs two times layers times KV heads times head dimension times length times bytes times batch.

• Cache frozen keys and values
• Prefill versus decode
• Size formula

Decode is a matrix times a vector, at about one FLOP per byte, far below the ridge point of a typical GPU. So tokens per second is roughly bandwidth divided by the bytes read per step. The fixes are fewer bytes: grouped-query attention, sliding windows, quantized caches, and sharing weight reads through batching.

• About 1 FLOP per byte
• tok/s ≈ bandwidth / bytes
• Shrink bytes, share reads

Next time we turn from speed to trust. Why does a model that can read its entire context so fluently still state false things with total confidence? Episode six covers why language models hallucinate, and what actually helps.

• Episode 6: why models hallucinate
• Subscribe for the series