A decode loop with and without a KV cache
By the end: You can write both generation loops in NumPy, show they produce identical output, and count the work the cache saves per step rather than asserting that it helps.
- The naive loop: re-running the whole prefix for every new token
- Why that is quadratic, and in what
- What the cache actually holds, and what it does not
- Prefill as one batched pass, decode as a loop of single-token passes
- Attention with no mask, because the future does not exist yet
- Cache memory arithmetic and why it becomes the binding constraint
New to quantization & inference optimisation? Every number inside a model is stored with a certain precision: how many binary digits it gets, and so how finely it can be told apart from its neighbours. Sixteen bits per number is the usual starting point. Quantization stores them with fewer — often eight, or four — which shrinks the model in memory by the same proportion and speeds it up, because what limits inference is usually how fast numbers can be moved to the processor rather than the arithmetic itself. Why it matters
You can already say that prefill is compute-bound and decode is memory-bound, and that a KV cache is what makes the second one bearable. This section writes both loops over a one-layer model in NumPy, checks they agree token for token, and counts the operations each performs — because the interview question is usually not what a KV cache is but what exactly it saves, and the honest answer is arithmetic rather than adjectives.
A model small enough to fit on a page#
One embedding table, one causal attention layer, one projection back to vocabulary size. It is not a good language model and it is not meant to be; it is the smallest thing with the structure that makes caching work.
import numpy as np
rng = np.random.default_rng(0)
V, d = 64, 32 # vocabulary, model width
E = rng.normal(size=(V, d)) * 0.1 # token embeddings
W_q, W_k, W_v, W_o = (rng.normal(size=(d, d)) * 0.1 for _ in range(4))
W_lm = rng.normal(size=(d, V)) * 0.1 # projection to logits
work = {"proj": 0, "scores": 0} # counters, to measure rather than claim
def softmax(z, axis=-1):
z = z - z.max(axis=axis, keepdims=True)
e = np.exp(z)
return e / e.sum(axis=axis, keepdims=True)The naive loop#
The obvious way to generate is to append the new token to the sequence and run the model over the whole thing again. It is correct, it is what every first implementation does, and the wasted work is easy to miss because each individual pass looks reasonable.
def forward(tokens):
x = E[tokens] # (T, d)
T = x.shape[0]
q, k, v = x @ W_q, x @ W_k, x @ W_v # (T, d) each
work["proj"] += 3 * T
scores = q @ k.T / np.sqrt(d) # (T, T)
work["scores"] += T * T
scores = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, scores)
ctx = softmax(scores) @ v # (T, d)
return (ctx @ W_o) @ W_lm # (T, V)
def generate_naive(prompt, n_new):
tokens = list(prompt)
for _ in range(n_new):
logits = forward(tokens)[-1] # every row computed, one row used
tokens.append(int(np.argmax(logits)))
return tokensTwo things are recomputed on every step. The keys and values for every token already in the sequence, which have not changed and cannot change — the causal mask guarantees that a token's key and value depend only on itself and what came before it, never on what comes after. And the full square score matrix, of which only the last row is ever read.
That prefix invariance is not an incidental property. It is the causal mask, restated, and it is the entire reason a cache is sound rather than an approximation. If you proved it as an assertion when you wrote attention — truncating the sequence leaves earlier outputs unchanged — this is what you were proving it for.
Prefill, then decode#
The cached version splits generation into two phases that run different code. Prefill processes the prompt in one batched pass, exactly as the naive forward does, and keeps the keys and values it computed. Decode then handles one token at a time, computing three vectors rather than three matrices, appending two of them to the cache and attending over everything stored.
def prefill(tokens):
x = E[tokens] # (T, d)
T = x.shape[0]
q, k, v = x @ W_q, x @ W_k, x @ W_v
work["proj"] += 3 * T
scores = q @ k.T / np.sqrt(d) # (T, T)
work["scores"] += T * T
scores = np.where(np.triu(np.ones((T, T), dtype=bool), 1), -np.inf, scores)
ctx = softmax(scores) @ v
logits = ((ctx @ W_o) @ W_lm)[-1] # (V,) - only the last position matters
return {"k": k, "v": v}, logits # keep k and v, discard q
def decode_step(token, cache):
x = E[token] # (d,) - one token, not a sequence
q, k, v = x @ W_q, x @ W_k, x @ W_v # (d,) each
work["proj"] += 3
cache["k"] = np.vstack([cache["k"], k]) # (t, d), one row longer
cache["v"] = np.vstack([cache["v"], v])
# no mask: everything in the cache is in the past, the future does not exist yet
scores = cache["k"] @ q / np.sqrt(d) # (t,) - a vector, not a matrix
work["scores"] += scores.shape[0]
ctx = softmax(scores) @ cache["v"] # (d,)
return (ctx @ W_o) @ W_lm # (V,)The decode step needs no causal mask
This surprises people, and it is worth being able to say why. The mask exists to stop a position attending to tokens after it. During decode the only tokens that exist are the ones already generated, so there is nothing after the current position to mask. The mask reappears in prefill, where many positions are processed at once, and in training, where the whole sequence is present.
def generate_cached(prompt, n_new):
cache, logits = prefill(list(prompt))
tokens = list(prompt)
for _ in range(n_new):
nxt = int(np.argmax(logits))
tokens.append(nxt)
logits = decode_step(nxt, cache)
return tokensThe same output, and the work saved#
A cache is not an approximation, so the first thing to check is that the two loops produce the same tokens. Then count.
prompt, n_new = [5, 9, 2, 41], 60
work = {"proj": 0, "scores": 0}
a = generate_naive(prompt, n_new)
naive = dict(work)
work = {"proj": 0, "scores": 0}
b = generate_cached(prompt, n_new)
cached = dict(work)
print(a == b) # True
print(naive) # {'proj': 6030, 'scores': 85330}
print(cached) # {'proj': 192, 'scores': 2086}The shape of the saving matters more than these particular numbers. At step t the naive loop performs work proportional to t for the projections and to t squared for the scores, so over n steps it is quadratic and cubic respectively. The cached loop performs constant projection work per step and work proportional to t for the scores, so over n steps it is linear and quadratic. Every measurement is one power of n cheaper.
naive total cached total
projections O(n^2) O(n)
score elements O(n^3) O(n^2)
memory held O(n) tokens O(n) cache rowsWhat the cache costs#
The saving is bought with memory, and in a real model that memory is the constraint that decides how many requests a GPU can serve at once. Two vectors per token, per layer, per attention head, at whatever precision the cache is kept in.
bytes = 2 * layers * kv_heads * head_dim * bytes_per_number * tokens * batch
# a 70B-class model: 80 layers, 8 KV heads (grouped-query), head_dim 128, fp16
# 2 * 80 * 8 * 128 * 2 = 327,680 bytes per token, about 0.31 MB
# 8,000 tokens of context -> about 2.5 GB per sequence
# 32 concurrent sequences -> about 80 GB, the whole cardTwo consequences follow that get asked about. The cache grows with the number of tokens, not with how fast they arrive, so a long conversation costs memory for its whole lifetime even while it is idle. And because a prefix produces the same keys and values every time, two requests sharing a system prompt can share those rows — which is what prompt caching sells, and what paged attention makes possible at block granularity.
Where people go wrong
Believing the cache is an optimisation that trades a little accuracy for speed, and then hunting for a quality regression that is not there. It is exact: the same arithmetic on the same numbers, performed once instead of repeatedly. If cached and uncached generation disagree by more than floating point tolerance, there is a bug — commonly a position index that is not advanced, so the rotary embedding is applied at the wrong offset.
If this comes up in an interview
Write a generation loop with a KV cache, and say what it saves.
Prefill runs the prompt through in one batched pass, with the causal mask, and keeps the keys and values. Each decode step then embeds the single new token, computes one query, key and value vector, appends the key and value to the cache, attends over the whole cache with no mask, and produces logits for one position. It saves recomputing every past key and value — constant work per step instead of work proportional to the sequence so far — and it computes one row of attention scores instead of the full square matrix. The output is identical; a cache is exact, not an approximation.
Then they push further
Why is it exact?
Because of the causal mask. A token's key and value depend only on that token and the ones before it, never on what follows, so they cannot change as generation continues. That is the same property as a prefix of a sequence producing unchanged outputs when the rest is truncated.
Why is there no mask in the decode step?
The mask prevents attending to the future, and during decode the future has not been generated. Everything in the cache is already in the past. The mask is needed in prefill and in training, where many positions are processed together.
What does it cost you?
Memory, and it is usually the binding constraint on concurrency. Two vectors per token per layer per KV head. For an 80-layer model with 8 KV heads and head dimension 128 in fp16, about 0.31 MB per token, so 8,000 tokens is 2.5 GB for one sequence and thirty-odd concurrent sequences fill an 80 GB card. Grouped-query attention and paged allocation are both responses to that.
What gets marked up. Anyone can say the cache stores past keys and values. The follow-ups are where it separates: that it is exact and why, that decode needs no mask, and that the cost is memory rather than compute.
Why is the query not cached alongside the key and value?
Because it is used exactly once. The query for position t is compared against every key up to t to produce that position's output, and no later position ever refers to it — later positions bring their own queries. Keys and values, by contrast, are read again at every subsequent step, which is what makes them worth storing.
A naive loop regenerates 500 tokens from a 2,000-token prompt. How much work is wasted?
Each step re-projects the whole prefix, so step t does work proportional to t rather than to one. Summed over 500 steps from a base of 2,000 that is roughly 1.1 million token-projections instead of 500, a factor of about two thousand. The attention scores are worse still: the naive loop builds the full t-by-t matrix each step and reads one row of it, which is cubic in total against the cached version's quadratic.
How does prompt caching relate to the KV cache?
It is the same data, kept between requests rather than within one. A shared prefix produces identical keys and values every time it is processed, because nothing after it can affect them, so a server that has already prefilled that prefix can reuse those rows and skip the prefill work. That is why prompt caching is billed against input tokens and why it only helps when the prefix is genuinely identical from the first token.
Build this: Implement both loops over the one-layer model above and assert the generated token sequences are identical. Instrument both with counters and plot projections and score elements against the number of tokens generated, confirming the quadratic-against-linear and cubic-against-quadratic shapes. Then compute the cache size in bytes for a model you actually use and work out how many concurrent 8,000-token sessions fit on one 80 GB card.
This is one section of Quantization & inference optimisation, 19 hours across 4 sections. Read it in full to get the whole topic in order.
Stop 28 of 50 on the AI Engineer route
0 of 50 topics done · 0h behind you, 603h to go
What you have marked done is stored in this browser only. It does not reach a server, and clearing site data will clear it.