If an interviewer asks "why does the KV cache grow, and what does it cost me?", the answer is a formula you can write on the whiteboard: 2 × layers × KV heads × head dimension × tokens × bytes. This post builds that answer, plus the other questions that cluster around LLM attention and the KV cache, and shows the numbers from a script I ran on 2026-10-08.
These are practice questions I wrote for this site. They aren't leaked and they aren't from any particular company. Each has a strong answer outline, the follow-up an interviewer tends to ask next, and a common wrong answer. For the systems side (RAG, agents, evals) there's a separate list in AI engineer interview questions on RAG, agents and evals, and for prompting there's prompt engineering interview questions with answers.
What is scaled dot-product attention, and why divide by sqrt(d_k)?
The operation, from Attention Is All You Need (Vaswani et al., 2017):
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V
Each query is compared with every key by dot product, the scores become weights through a softmax, and the output is a weighted sum of the values.
Strong answer outline
- Say what Q, K and V are: linear projections of the same hidden states (in self-attention), with the causal mask hiding future positions in a decoder.
- The division matters because dot products of d_k-dimensional vectors grow with d_k. If the components of q and k are independent with mean 0 and variance 1, the dot product has variance d_k, so its standard deviation is sqrt(d_k). Large scores push the softmax toward a one-hot output, where gradients are tiny. Dividing by sqrt(d_k) brings the variance back to about 1. The paper makes this argument itself.
- Say it's a variance argument about initialisation-scale inputs, not a law about trained weights.
I checked the claim instead of just repeating it. The script below samples random normal vectors and measures both the standard deviation of the dot product and how peaked the softmax is over 64 random keys:
d_k= 16 std(dot)= 3.94 sqrt(d_k)= 4.00 mean max-softmax unscaled=0.547 scaled=0.103
d_k= 64 std(dot)= 8.12 sqrt(d_k)= 8.00 mean max-softmax unscaled=0.800 scaled=0.105
d_k= 256 std(dot)= 16.10 sqrt(d_k)= 16.00 mean max-softmax unscaled=0.886 scaled=0.104
d_k= 1024 std(dot)= 32.20 sqrt(d_k)= 32.00 mean max-softmax unscaled=0.955 scaled=0.108
Without scaling, at d_k = 1024 the single largest weight already averages 0.955 of the probability mass. With scaling it stays near 0.10, and with 64 keys the uniform value would be 0.0156, so there's still some peak but the softmax isn't saturated.
Follow-up an interviewer asks: What would go wrong if you used d_k instead of sqrt(d_k)? Scores would shrink as d_k grows, so the softmax would flatten toward uniform and the model would struggle to focus. sqrt(d_k) is the value that keeps the variance constant.
Common wrong answer: "It's for numerical stability so we don't overflow." Softmax implementations already subtract the max for stability. The scaling is about keeping the softmax out of its saturated, low-gradient region.
Why is attention O(n^2), and is that the memory or the compute?
Strong answer outline
- The score matrix
Q K^Tis n by n, so compute for a full forward pass is quadratic in sequence length n. - A naive implementation also writes that n by n matrix to GPU memory. Per head, per layer, in fp16 that is the table below (computed, not estimated):
| n (tokens) | Score entries | fp16 size per head per layer |
|---|---|---|
| 1,024 | 1,048,576 | 2 MiB |
| 4,096 | 16,777,216 | 32 MiB |
| 32,768 | 1,073,741,824 | 2,048 MiB |
| 131,072 | 17,179,869,184 | 32,768 MiB |
- FlashAttention (Dao et al., 2022) computes exact attention with tiling so the full matrix is never written to slow GPU memory. It cuts memory traffic and activation memory, but the arithmetic is still quadratic.
- During autoregressive decoding the per-token cost is linear in the cached length, because one new query attends to n cached keys.
Follow-up: So does a 1M-token context cost a thousand times more than 32k? For prefill compute, roughly 1,000 times more in the attention part (the ratio of squares), though other parts of the network scale linearly, so the whole-model number is lower. For decoding, the cache grows about 30 times, which is a different cost. Say which phase you mean.
Common wrong answer: "FlashAttention makes attention linear." It makes it memory-efficient and faster in wall-clock terms. It's still exact attention, so still quadratic compute.
What is the difference between MHA, MQA and GQA?
Strong answer outline
- MHA (multi-head attention): each query head has its own key and value head. Biggest cache.
- MQA (multi-query attention, Shazeer 2019): all query heads share one key head and one value head. The paper's motivation is incremental decoding: loading the large keys and values tensors on every step is a memory-bandwidth cost, and sharing them shrinks it.
- GQA (grouped-query attention, Ainslie et al. 2023): an intermediate, with one KV head per group of query heads. The paper also describes converting an existing multi-head checkpoint to GQA with a small amount of further training, which is why many open models use it.
- The cache size scales with the number of KV heads, not query heads. Query heads still cost compute.
I pulled the config files from public mirrors of the original model repos and computed the real shapes:
| Model | Layers | Query heads | KV heads | Head dim |
|---|---|---|---|---|
| Llama 2 7B | 32 | 32 | 32 (MHA) | 128 |
| Mistral 7B v0.1 | 32 | 32 | 8 (GQA) | 128 |
| Llama 3 8B | 32 | 32 | 8 (GQA) | 128 |
| Llama 2 70B | 80 | 64 | 8 (GQA) | 128 |
Head dimension is hidden_size / num_attention_heads (4096/32 and 8192/64). The configs came from NousResearch/* and mistralai/Mistral-7B-v0.1 on Hugging Face, fetched 2026-10-08. The Llama ones are third-party mirrors, so I used them only for shape fields. Check Meta's own model card if you cite them.
Follow-up: If GQA saves so much memory, why not always use MQA? Fewer KV heads means less capacity to attend differently per head, and the GQA paper reports quality between MHA and MQA, with GQA close to MHA. It's a trade-off you measure, not a free lunch.
Common wrong answer: "GQA reduces the number of query heads." No. Query heads stay; KV heads are shared.
What is the KV cache and how much memory does it use?
Strong answer outline
- In decoding, the keys and values of earlier tokens don't change. Cache them per layer; for each new token compute one new K and V and attend over the cache.
- The formula, in bytes:
kv_bytes = 2 * layers * kv_heads * head_dim * seq_len * batch * bytes_per_element
The 2 is one K tensor and one V tensor. For fp16 or bf16, bytes_per_element is 2.
Here is the script output (the full script is in the next section so you can run it):
model per token KiB 4k x1 32k x1 4k x32
Llama-2-7B (MHA) 512.0 2.00 16.00 64.00
Mistral-7B v0.1 (GQA) 128.0 0.50 4.00 16.00
Llama-3-8B (GQA) 128.0 0.50 4.00 16.00
Llama-2-70B (GQA) 320.0 1.25 10.00 40.00
70B-shaped model, 4k tokens, batch 1: MHA=10.00 GiB GQA(8)=1.25 GiB MQA(1)=0.16 GiB
Units are GiB; "4k x32" means 4,096 tokens at batch 32. Two numbers to remember: Llama 2 7B pays 512 KiB of cache per token, and GQA cuts that to a quarter on the 8-KV-head 7B/8B shapes. A 70B-shaped model would need 10 GiB for one 4k sequence if it used full MHA, versus 1.25 GiB with 8 KV heads.
Then connect it to serving. With the parameter counts derived from the config (6.74B and 8.03B), fp16 weights are 12.6 GiB and 15.0 GiB. On an 80 GB GPU (74.5 GiB), the space left over holds about 31 full 4k sequences of Llama 2 7B cache versus about 119 for Llama 3 8B, ignoring activations and runtime overhead. So the cache, not the weights, caps your batch size.
Follow-up: What are your options if the cache doesn't fit? Lower precision for the cache (8-bit KV is supported in some serving stacks; measure quality), fewer or shorter sequences, GQA/MQA models, sliding-window attention (Mistral 7B v0.1's config sets a 4096-token sliding window), paging the cache (below), or offloading. I haven't benchmarked KV quantisation quality here, so treat it as a lever to test, not a claim.
Common wrong answer: "The cache stores the attention scores." It stores K and V, not scores. Scores are recomputed for each new query.
What is the difference between prefill and decode?
Strong answer outline
- Prefill processes the whole prompt in one pass. All prompt tokens are computed in parallel, which keeps the GPU busy with large matrix multiplications. It determines time to first token.
- Decode generates one token per step per sequence. Each step reads all the weights and the sequence's cache to produce a single token's worth of compute. That ratio of bytes moved to arithmetic makes decoding typically memory-bandwidth-bound, which is the problem Shazeer's MQA paper targets.
- Consequences: long prompts hurt time to first token; long outputs hurt total latency linearly; batching more sequences raises throughput because the weight read is shared across them, until cache memory runs out.
Follow-up: Why does batching help decode but not much for a single user? Weights are read once per step regardless of batch size, so more sequences per step means more tokens per weight read. A single user gets the latency of one sequence, regardless.
Common wrong answer: "Decode is slow because attention is quadratic." Per decode step, attention is linear in cached length. The slowness is mostly moving weights and cache through memory.
What do PagedAttention and continuous batching solve?
Concept level is what interviews usually want.
Strong answer outline
- The problem: requests have unpredictable output lengths. If you reserve one contiguous max-length slab of cache per request, you waste memory on fragmentation and over-reservation, which lowers the batch size.
- PagedAttention (vLLM, Kwon et al., SOSP 2023) borrows virtual-memory paging: the cache is split into fixed-size blocks, stored non-contiguously, and a block table maps each sequence to its blocks. Memory is allocated as tokens are generated. Blocks can also be shared, for example across samples from the same prompt.
- Continuous batching (iteration-level scheduling, described in the Orca paper, OSDI 2022): instead of waiting for every request in a batch to finish, the scheduler adds and removes requests at each generation step, so a short request leaves and a new one enters immediately.
- Together they raise throughput by keeping the GPU busy and fitting more sequences in memory.
Follow-up: What does prefix caching add? If many requests share a long prefix (a system prompt), the cache blocks for that prefix can be reused, saving prefill compute. This is the mechanism behind the prompt caching discounts that APIs offer; see context caching explained for the pricing side.
Common wrong answer: "PagedAttention makes attention faster." Its main win is memory management, which enables larger batches. The paper's gains are in serving throughput.
Run the numbers yourself
This is the script I used for the table above, trimmed to the cache calculator. I ran the full version on Python 3 with NumPy on 2026-10-08 and the output matches the table.
def kv_bytes(layers, kv_heads, head_dim, seq_len, batch, bytes_per=2):
return 2 * layers * kv_heads * head_dim * seq_len * batch * bytes_per
configs = {
"Llama-2-7B (MHA)": (32, 32, 128),
"Mistral-7B v0.1 (GQA)": (32, 8, 128),
"Llama-3-8B (GQA)": (32, 8, 128),
"Llama-2-70B (GQA)": (80, 8, 128),
}
for name, (layers, kv_heads, head_dim) in configs.items():
per_token = kv_bytes(layers, kv_heads, head_dim, 1, 1)
gib_4k = kv_bytes(layers, kv_heads, head_dim, 4096, 1) / 2**30
print(f"{name:24s} {per_token/1024:7.1f} KiB/token {gib_4k:5.2f} GiB at 4k")
Swap in your own model's config.json values: num_hidden_layers, num_key_value_heads, and hidden_size / num_attention_heads. If num_key_value_heads is missing, the model uses MHA and the KV count equals num_attention_heads.
What the numbers leave out: activations, framework overhead, fragmentation (the thing paging fixes), and any cache quantisation. Treat the output as a lower bound for planning, then measure on your serving stack. For converting memory into API spend, see the LLM token counting and context management guide.
How to prepare for this round
- Be able to derive the cache formula from scratch and plug in a real model's config. That single skill covers most of the memory questions in this post.
- Memorise one number for contrast, such as "MHA 7B pays 512 KiB per token, GQA with 8 KV heads pays 128".
- Practise saying which phase you're talking about. Most wrong answers in this area come from mixing up prefill and decode.
- Know where each claim comes from: the transformer paper for the scaling, Shazeer and Ainslie for MQA/GQA, the vLLM paper for paging. Interviewers notice when you can cite the source.
- The how LLMs work lesson and the context engineering lesson are good background if the vocabulary above is new.
What I could not verify
I did not benchmark any serving engine, so I make no claims about tokens per second. I didn't measure KV-cache quantisation quality. The Llama configs came from third-party mirrors. The vLLM and Orca descriptions are from the papers' abstracts and my reading, at concept level only.
