KV-cache optimization: Question bank — AI-ML deep dive

Definition The symbols this page uses (read this first)

Everything below is built from a handful of letters. Here is each one, in plain words, before we ever use it. The picture to hold in mind is the strip of past tokens in figure 1.

Token — one chunk of text (a word or word-piece). The model generates them one at a time, left to right.
$L$ — the number of tokens in the sequence so far (its length).
$x_{t}$ — the vector of numbers ("embedding") that represents the token sitting at position $t$ .
$Q$ (query), $K$ (key), $V$ (value) — three vectors the model builds from each token. Think of a token walking into a library: its query is "what am I looking for?", its key is "here's my label so others can find me", its value is "here's the content I'll hand over if someone picks me". The cache stores $K$ and $V$ only.
$W_{Q}, W_{K}, W_{V}$ — three fixed weight matrices (grids of numbers) learned during training. To get a token's query/key/value you multiply its embedding by the matching matrix: $Q_{t} = W_{Q} x_{t}$ , $K_{t} = W_{K} x_{t}$ , $V_{t} = W_{V} x_{t}$ . Because the $W$ matrices are frozen once training ends, $K_{t}$ and $V_{t}$ for a fixed token never change — that fixedness is the whole reason caching works.
$d$ — the model dimension: how many numbers are in each token's embedding vector.
head — attention is split into several parallel copies called heads, each looking at a different slice of the vector. $h$ = the number of heads.
$d_{k}$ — the head dimension: how many numbers each single head works with (so $d \approx h \cdot d_{k}$ ).

Keep this box open while you work — most traps below are really about which of these changes each step and which stays fixed.

Look at figure 1 above: the fixed grey blocks on the left are the cached $K, V$ of past tokens; the bright block on the right is the brand-new token, which needs a fresh query but only appends one key and one value.

True or false — justify

Answer T/F, then give the reason. A bare "True" scores zero — the reasoning is the point.

The KV-cache changes the mathematical output of generation compared to no-cache.

False. The cache is a pure memoisation trick — it stores identical numbers instead of recomputing them. Output is bit-identical (up to floating-point non-determinism); only speed and memory differ.

KV-caching reduces the total amount of memory the model uses during inference.

False. It increases memory: you now store

K, V

for every past token. It trades extra memory for saved compute — the opposite of a memory saving.

KV-caching gives a speedup during training as well as generation.

False. Training uses teacher forcing — all

L

tokens are present at once, so each

K_{i}, V_{i}

is computed exactly once anyway. There is nothing to reuse across steps, so the cache only helps autoregressive generation.

The keys and values for token 5 stay constant while we generate tokens 6, 7, 8, …

True. The weight matrices

W_{K}, W_{V}

are frozen after training and token 5's embedding

x_{5}

is fixed, so

K_{5} = W_{K} x_{5}

and

V_{5} = W_{V} x_{5}

never change. That fixedness is exactly why they are cacheable.

Multi-Query Attention reduces the number of model parameters as much as it reduces cache size.

False. MQA shrinks the cache by a factor of

h

(the number of heads), but the parameter saving is tiny — only the

K, V

projection matrices shrink, and those are a small slice of the whole model. The win is memory-bandwidth at inference, not parameter count.

With a KV-cache, the per-token generation cost becomes independent of how many tokens came before.

Mostly false. The projection cost per token is now constant (

O (d^{2})

, where

d

is the model dimension), but the attention score step still costs

O (t \cdot d)

because the new query

Q_{t}

must dot against all

t

cached keys. Latency still grows with context length, just linearly instead of quadratically overall.

A sliding-window cache of size

W

makes memory constant regardless of sequence length.

True. You only keep the last

W

keys/values, so cache memory is capped at

O (W)

. The cost is that tokens beyond

W

steps back are literally invisible — long-range dependencies are dropped.

Grouped-Query Attention is a special case that contains both MQA and standard attention as endpoints.

True. GQA groups

h

query heads into

g

KV-groups. Setting

g = h

recovers standard multi-head attention; setting

g = 1

recovers MQA. GQA interpolates between them.

Spot the error

Each line states a plausible-sounding claim with a hidden flaw. Name the flaw.

"Let's cache

Q

too — it also needs a matrix multiply, so we'll save even more compute."

The flaw: each new token generates a fresh query

Q_{t}

that never gets reused, because the query encodes "what this new token is looking for". Only

K, V

(what past tokens offer) are reused. Caching

Q

saves nothing.

"Since the cache already holds every key, token

t

can attend to all of them freely."

Missing the causal mask. Token

t

must only see positions

\leq t

. If future keys ever sit in the cache (e.g. in some batched/padded setups), attending to them leaks information the model shouldn't have.

"Doubling the number of heads

h

doubles the compute but leaves cache size the same."

Wrong for standard attention: cache size is

2 L h d_{k} \times

precision (with

h

= heads,

d_{k}

= head dimension), so it scales linearly with $h$ . Doubling heads doubles the cache too. (This is precisely why MQA/GQA exist.)

"MQA and GQA are lossless — they just reorganise the same numbers."

False. Sharing

K, V

across query heads removes representational capacity, so the model must be trained or fine-tuned for it and typically loses a small amount of quality (≈1–2% perplexity for MQA). It is an approximation, not a rewrite.

"FlashAttention and KV-caching solve the same bottleneck, so using both is redundant."

They attack different bottlenecks. KV-cache removes redundant

K, V

projection recompute across steps; FlashAttention removes the memory cost of materialising the

O (L^{2})

score matrix within a step. They compose.

"To save memory we can store the cache in int8 with zero downside."

The flaw is "zero downside". Cache quantization does shrink memory a lot, but low-precision

K, V

inject rounding error into every attention score — usually acceptable, but not free. It's a tradeoff, not a freebie.

"Beam search with

B

beams needs the same cache as greedy decoding."

Wrong: each of the

B

beams is a distinct hypothesis with its own past tokens, so you need

B

separate KV-caches (or a shared-prefix cache with per-beam tails). Memory scales with the beam width.

Why questions

The reasoning is the whole answer here.

Why does the naive (no-cache) projection cost grow as

O (L^{2} d^{2})

rather than

O (L d^{2})

At step

t

the naive model treats the sequence as fresh and recomputes projections for all

t

tokens, costing

O (t d^{2})

. Summing

\sum_{t = 1}^{L} t \approx L^{2} /2

gives the quadratic. The cache turns each step's cost into a constant

O (d^{2})

, so the sum becomes linear.

Why is it

K

and

V

that get cached, and not the attention output?

K, V

depend only on past tokens and are fixed, so they are reusable inputs. The attention output depends on the current

Q_{t}

, which is new every step — so the output must be recomputed and there is nothing stable to cache there.

Why does the KV-cache size scale with the number of layers?

Every transformer layer has its own independent attention with its own

K, V

projections. Caching happens per layer, so total cache = per-layer cache

\times n_{layers}

. This is why a per-token cache jumps from ~48 KB (one layer) to ~4.7 MB (all layers) in a GPT-3-scale model.

Why does MQA help inference latency more than raw FLOP count suggests?

Autoregressive generation is memory-bandwidth bound, not compute bound — the bottleneck is loading the cache from memory each step. MQA shrinks the cache by

h \times

, so far less data crosses the memory bus per token, giving a bigger real-world speedup than the FLOP savings alone imply.

Why can't we just increase batch size indefinitely to amortise generation cost?

Because each sequence in the batch needs its own full KV-cache, and cache memory grows with (batch × length × layers × heads). You quickly hit the GPU memory ceiling — the cache, not the weights, becomes the limiting resource for large batches. See deployment considerations.

Why does the causal mask still matter even after caching, given the cache holds only past tokens?

In practice caches are often batched and padded, and some pipelines temporarily hold placeholder or future entries. The mask is the guarantee that token

t

mathematically attends only to

\leq t

, keeping generation causal regardless of what physically sits in the buffer.

Edge cases

Boundary conditions people forget. Push each case to its extreme.

At the very first token (

t = 1

), what does the KV-cache contain and does caching help yet?

The cache is empty before token 1; you compute

K_{1}, V_{1}

and store them. There is nothing prior to reuse, so caching saves nothing on step 1 — the payoff begins from step 2 onward.

For a sequence of length

L = 1

(single-token generation), what is the cache speedup?

Exactly

1 \times

— no speedup. The speedup factor is

\approx L

, so for

L = 1

there is nothing to amortise. Caching only pays off as the sequence grows.

What happens to a sliding-window cache when the model needs a dependency older than the window

W

That information is simply gone — the key/value has been evicted. The model can no longer attend to it, so any answer requiring that far-back context degrades. This is the fundamental accuracy cost of bounded-memory caching.

In the MQA cache-ratio

\frac{1}{h}

, what does the limit

h \to 1

mean physically?

With

h = 1

there is only one head, so standard attention already is single-KV attention and MQA gives no reduction (

\frac{1}{1} = 1

). The benefit of MQA grows only as head count grows.

In GQA, what does the group count

g = h

correspond to, and what does

g = 1

correspond to?

g = h

means every query head keeps its own KV — that is ordinary multi-head attention with no cache saving.

g = 1

means all heads share a single KV — that is full MQA with maximum saving. GQA lives strictly between these.

What limits the benefit of KV-caching as context length grows extremely large (say 1M tokens)?

Two limits collide: cache memory grows linearly and can exhaust the GPU, and the attention-score step still costs

O (t d)

per token, so per-token latency keeps climbing. This is why very-long-context systems reach for sparse attention or windowing on top of caching.

For a prompt that is processed all at once (the "prefill" phase) before generation, does caching change that step?

The prefill computes all prompt tokens'

K, V

in one parallel pass — like training, there's no reuse to exploit within prefill. Caching's role there is simply to store those

K, V

so the subsequent one-token-at-a-time decode phase can reuse them.

Recall Quick self-test

Which of $Q$ , $K$ , $V$ is recomputed every generation step? ::: Only $Q$ — it encodes what the new token seeks; $K$ and $V$ are cached because they belong to fixed past tokens. Does the cache change the output numbers? ::: No — it only avoids recomputing identical values, so results are (bit-)identical. What resource does MQA/GQA save, and what does it cost? ::: Saves KV-cache memory and memory bandwidth by a factor up to $h$ ; costs a small amount of model quality.

Related vault topics for definitions used above: 6.1.1-transformer-architecture, 6.1.4-multi-head-attention, 6.1.11-sparse-attention, 6.3.1-model-quantization, 5.4.2-beam-search, 6.2.3-inference-optimization, 7.2.1-llm-deployment, and the parent KV-cache optimization.