JA EN
LearnInference & Serving
·FREE·7 min read

The KV Cache from Scratch — The Heart of Fast Inference

An LLM emits one token at a time. Written naively, every single token costs a full recomputation of the whole sequence — a spectacular waste. The keys and values of past tokens never change again, and that one fact drops an entire order of magnitude. What you pay instead is memory, in an amount you can work out yourself, and that is why batch size and context length hit a ceiling.

ModalitytextTaskinference

Generation only moves one token at a time

When a language model writes, the text does not arrive all at once. It predicts the next token, appends it to the input, and predicts the next one. That loop is autoregressive generation, and it is stubbornly sequential.

A 100-token reply therefore means running the model forward 100 times — and there is a trap that every naive implementation falls into.

The naive version recomputes everything, every time

The most obvious implementation goes like this. Feed in the prompt, get a token. Append that token to the prompt and feed the whole thing back in to get the next one. Append, feed the whole thing in again, repeat.

It works. But producing token tt means processing a sequence of length tt from scratch, and as the attention article showed, attention alone costs O(t2)O(t^2) for that. Summed from t=1t=1 to nn, the whole generation lands at O(n3)O(n^3). The longer the context, the more violently that waste compounds.

FIG 1Naive recomputation grows with the cube of the sequence length; with a cache it grows with the square. One step of order is the difference between a plausible wait and an impossible one

Past keys and values never change again

Where exactly is the waste? Self-attention builds a Query, a Key and a Value from each token, scores queries against keys, and mixes the values. Now recall the causal mask: a generative model must not look at future tokens, so each position attends only to itself and what came before.

That has a consequence. The K and V of token 3 are unaffected by tokens 1 and 2 being recomputed — and more importantly, adding tokens 4, 5 or 100 later does not change token 3's K and V by a single bit. Each position's K and V depend only on that position and what precedes it.

If a value never changes and you keep recomputing it, there is one obvious answer: keep it and reuse it. That is the KV cache.

Generating one new token then requires only:

Attention per step drops to O(t)O(t), and the whole generation to O(n2)O(n^2) — a full order of magnitude below O(n3)O(n^3).

def decode_step(x_new, cache, layer):
    q = x_new @ layer.Wq                     # only for the one new token
    k = x_new @ layer.Wk
    v = x_new @ layer.Wv
    cache.k = concat(cache.k, k)             # append, never recompute
    cache.v = concat(cache.v, v)
    w = softmax(q @ cache.k.T / sqrt(d_k))   # against the whole history
    return w @ cache.v

Q is not cached because Q is "what am I asking right now" — only the new token's query is ever needed.

Two phases with opposite personalities: prefill and decode

The moment you add a cache, inference splits into two phases that behave nothing alike.

Prefill runs the entire prompt through once, producing the K and V of every position and filling the cache. All tokens are processed in parallel, so the arithmetic units are busy. This is what determines the delay before the first character appears, and it grows with prompt length.

Decode advances one token at a time. With only one new token per step, the matrix multiplies become vector-by-matrix rather than matrix-by-matrix, and the amount of data read is wildly out of proportion to the arithmetic performed — the weights and the whole KV cache must be pulled from memory every step. Decode is therefore memory-bandwidth bound, not compute bound.

Knowing this asymmetry keeps you from being puzzled by "GPU utilisation is low but generation is slow". The arithmetic units are idle; the memory bus is saturated. In practice the first is measured as TTFT (time to first token) and the second as time per output token, separately. Saying "it's slow" without distinguishing them tells you nothing about what to fix.

There is a quieter consequence too. The KV cache is per-request state. For as long as a conversation continues, that cache occupies memory on a specific GPU. Route the next request to another node the way you would with a stateless web server, and the node without the cache has to redo prefill. LLM serving sits badly with naive load balancing precisely because it carries this state.

Sizing the cache yourself

The speed is bought with memory, and the price is a multiplication.

M=2×L×n×Hkv×dh×b×BM = 2 \times L \times n \times H_{kv} \times d_h \times b \times B
(1)

That product is a tally, not a theory. Said in words: for every layer of the model, and for every token already sitting in the context, the server is holding on to two vectors — one key and one value — and it holds a separate set of those for every request currently in flight. The formula does nothing but count them.

Symbol by symbol: the leading 22 is for K and V. LL is the number of layers (each keeps its own K and V). nn is the context length, i.e. how many tokens are cached. HkvH_{kv} is the number of KV heads and dhd_h the dimension per head, so Hkv×dhH_{kv} \times d_h is the size of one layer's K. bb is bytes per element (2 for fp16). BB is how many requests are in flight — the batch size.

What this says, plainly, is that the cache grows linearly in both context length and batch size.

Put numbers in. With 32 layers, Hkv×dh=4096H_{kv} \times d_h = 4096, fp16, a 4096-token context and a batch of one:

2×32×4096×4096×2 bytes=2 GiB2 \times 32 \times 4096 \times 4096 \times 2\ \text{bytes} = 2\ \text{GiB}

Read the left-hand side as a shopping list rather than as algebra: a key and a value, in each of 32 layers, for each of 4096 tokens, each of them 4096 numbers wide, at two bytes per number — which says the memory is not being spent on anything clever, only on keeping what the model is obliged to keep.

Two gigabytes for a single request. A batch of 8 makes it 16 GiB; doubling the context makes it 32. Whatever memory the weights left behind, this formula eats. The ceiling on concurrency and context length is nothing more than the consequence of that arithmetic.

There are ways to push back. The common one is GQA (grouped-query attention), which keeps many query heads but shares a smaller number of K/V heads: HkvH_{kv} shrinks directly, and the cache shrinks with it. The other is storing the cache at lower precision — reducing bb — which is exactly the subject of the quantization article.

Fragmentation, the other enemy

A naive server reserves one contiguous block per request, sized for the maximum context length. A reply that ends after 200 tokens still holds a 4096-token room, most of it never used. Borrowing the idea behind OS virtual memory — cutting the cache into fixed-size blocks and handing out only what is needed — is why PagedAttention, as implemented in vLLM, became standard. It has a useful side effect: requests that share the beginning of a prompt can share those blocks.

How this shows up on the job

The KV cache becomes your problem when you serve an LLM yourself: deciding how many GPUs to buy, setting an SLO for concurrent users, or investigating latency complaints. Even if you only call an API, it is what explains the shape of the pricing and the latency.

The knobs. In vLLM: --max-model-len (context ceiling), --max-num-seqs (concurrency), --gpu-memory-utilization (the fraction of memory weights and cache may occupy), --kv-cache-dtype (cache precision), --enable-prefix-caching (reuse KV for shared prefixes). In Hugging Face transformers, use_cache and past_key_values are the same concept in the raw.

The most common accident is "it started, then died with OOM". Fitting the weights and fitting the cache for your real concurrency and context length are different questions. Divide the free memory by 2LHkvdhb2 L H_{kv} d_h b and you get the total number of tokens you can hold at once. Doing that one division up front removes one night-time incident.

Do not break your own prefix cache. Putting the current timestamp or a request ID at the front of the system prompt makes every request's prefix unique, so KV that could have been shared is rebuilt each time. Put variable content at the end; prefill load changes on that alone.

Look for slowness on the correct side. Slow to the first token means prefill, which means prompt length. Slow, dragging output means decode, which means memory bandwidth and cache reads — shortening the prompt will not help. The first is fixed by trimming input and sharing prefixes, the second by revisiting batching and cache precision.

The interview and design-review question is "what doubles when you double the context?" KV cache memory doubles — it is linear. Prefill's attention cost quadruples — it is quadratic. Plenty of answers conflate the two, and simply separating them shows you understand the machine.

Summary

Next, the other way to reclaim that memory: holding weights and caches at lower precision — quantization.

Comments

Sign in to comment