Skip to content
datarekha

What is a KV cache and how does it speed up LLM inference?

The short answer

A KV cache stores the key and value tensors for previously processed tokens in each transformer layer, so autoregressive decoding computes only the new token's projections instead of rebuilding the whole prefix. This reduces per-step attention work from quadratic to linear in context length, while cache memory grows with layers, sequence length, batch size, and KV-head count.

How to think about it

Yes. During autoregressive decoding—the model generates one token, feeds it back, and repeats—a KV cache stores the key and value tensors for earlier tokens at every transformer layer. The next step computes the new token’s projections and attends to the stored tensors instead of rebuilding the entire prefix, so attention work per decode step grows linearly with context rather than quadratically.

The price is memory. The cache grows with the number of layers, active requests, and tokens in each request.

Why it works

Attention creates three representations for each token:

  • A query, or what this token is looking for.
  • A key, or what each token can be matched on.
  • A value, or the information returned when a key matches.

For a new token at position t, attention is conceptually:

softmax(q_t K^T / sqrt(d_head)) V

Here, q_t is the new token’s query, K contains keys from the context, V contains the corresponding values, and d_head is the size of one attention head.

A decoder-only LLM uses causal attention, meaning a token can attend to itself and earlier tokens but not to future tokens. That restriction creates the important invariant: once the model has processed the first 4,096 tokens, a later token cannot change the earlier tokens’ hidden states. Their keys and values are still valid.

So the model stores them.

It does not need to store old queries. A past query has already produced its attention output, and no future step will ask that old query another question. The next step needs one new query and the keys and values for all available positions.

Conceptually, one layer’s decode step looks like this:

# Conceptual pseudocode; tensor shapes are omitted
q_new = W_q @ x_new
k_new = W_k @ x_new
v_new = W_v @ x_new

K_cache.append(k_new)
V_cache.append(v_new)

scores = q_new @ transpose(K_cache) / sqrt(d_head)
x_new = softmax(scores) @ V_cache

Production kernels usually write into preallocated or paged memory rather than repeatedly copying arrays with append. The mathematical operation is the same.

What “quadratic to linear” really means

Suppose a prompt contains 4,096 tokens and the model is about to generate token 4,097.

Without a KV cache, a naïve decoder reruns the transformer over all 4,097 positions. For one attention head, that full pass conceptually considers a 4,097 × 4,097 score matrix: 16,785,409 entries before applying the causal mask. Only roughly half are valid causal pairs, but the work still grows as O(T²) with context length T.

With a KV cache, the model computes the new token’s query, key, and value once. Its query then compares against the 4,097 available keys. That is one row of scores, or O(T) work for the attention part of this decode step.

This is a large reduction in repeated work. It is not a 4,097-times speedup for the whole model, because the model still performs projections, feed-forward layers, memory reads, and other operations. The point is that it no longer recomputes every previous token’s attention state.

The first pass over the prompt is called prefill. It processes the prompt, usually in parallel, and fills the cache. The one-token-at-a-time phase is called decode.

KV caching does not make decode constant-time. If the prompt has 4,096 tokens and the model generates 1,000 more, each new token still reads a growing context, from roughly 4,100 keys to roughly 5,100 keys. Across those 1,000 steps, that is about 4.6 million query-key comparisons per attention head. The per-step cost is linear in the current context, and the total decode attention work is still quadratic in the number of generated tokens.

The memory cost

The cache is made of activations, not model weights. A useful approximation for its memory is:

B × L × T × 2 × H_KV × D_head × bytes_per_element

Here:

  • B is the number of active sequences.
  • L is the number of transformer layers.
  • T is the number of cached tokens, including prompt and generated tokens.
  • The 2 accounts for both keys and values.
  • H_KV is the number of key-value heads.
  • D_head is the head dimension.

Consider one 32-layer model with 32 key-value heads, a head dimension of 128, a 4,096-token context, and FP16 cache entries, which use 2 bytes each:

1 × 32 × 4,096 × 2 × 32 × 128 × 2

That equals 2,147,483,648 bytes, or 2 GiB, for one sequence. A batch of eight equally long sequences would need about 16 GiB for KV cache alone, before model weights, temporary buffers, CUDA workspace, and allocator overhead.

This is why a model can fit in GPU memory for one request and still fail under production traffic.

Modern models often use grouped-query attention, or GQA, where several query heads share fewer key-value heads. If the same example uses eight KV heads instead of 32, the cache falls to 512 MiB per sequence. Multi-query attention, or MQA, goes further by using one KV head, though the architecture and quality trade-offs depend on the model.

The production pattern

A serving system normally does four things:

  1. It runs prefill once for the prompt and stores K and V at every layer.
  2. For each generated token, it computes only the new token’s layer states and appends its K and V.
  3. It reads the cache when calculating attention for the next token.
  4. It discards the cache when the request ends, unless a compatible prefix is deliberately being reused.

The cache belongs to a sequence. It must not accidentally be shared between unrelated users. Systems serving variable-length requests often allocate the cache in fixed-size blocks or pages. That reduces fragmentation and lets a batch grow and shrink without reserving the maximum context length for every request.

Reusing the KV cache for an identical prompt prefix is a separate optimisation often called prefix caching. It can skip some prefill work, but it requires the same tokenized prefix and compatible model, adapter, positional-encoding, and runtime settings. A normal per-request KV cache only avoids recomputing tokens within that request.

The senior-level trade-off

KV caching trades arithmetic for memory and memory bandwidth.

At a short context, this trade is usually excellent. At a very long context, every generated token still has to read the entire cache. Decode can become memory-bandwidth-bound: the GPU spends more time moving K and V from memory than doing multiplication. A cache therefore reduces redundant computation but does not magically make long-context generation cheap.

There are several ways to control the cost:

  • Use GQA or MQA so fewer KV heads are stored.
  • Use a lower-precision KV format, such as FP8 or an integer format, when the runtime and model tolerate it.
  • Limit maximum context or generated tokens.
  • Use paged allocation so finished or short requests do not strand large memory regions.
  • Reuse compatible prompt prefixes when many requests share the same system prompt.
  • Batch requests carefully, because a larger batch improves hardware utilisation but multiplies cache memory.

KV caching is mainly an inference technique for autoregressive models. It is usually not useful during ordinary training, where teacher-forced sequences are processed in parallel and retaining a cache across training steps would consume memory without avoiding the same kind of repeated generation. It is also not the relevant optimisation for an encoder-only model that processes an input once rather than generating token after token.

A failure mode to recognise

A common production symptom is a GPU out-of-memory error that appears only after traffic increases or users send long prompts. The weights fit, and a short request succeeds, but several long generations exhaust memory. The usual cause is the cache formula: active sequences, not merely model size, determine the additional memory requirement.

Another serious bug is a stale or cross-request cache. The first visible symptom may be a response that starts normally but suddenly mentions a name, instruction, or document from a previous request. The cache must be keyed to the correct sequence and model configuration, then cleared or safely branched when that sequence ends or forks.

If enabling caching produces no latency improvement, inspect whether the system is actually re-running prefill for every token, transferring the cache between CPU and GPU, padding every request to an unnecessarily large length, or spending most of its time reading a very large cache. Separate prefill and decode timings; otherwise a fast decode path can be hidden by a slow prompt pass.

What they’ll ask next

Does a KV cache reduce memory usage?

No. It usually increases memory usage because it stores intermediate tensors that would otherwise be recomputed. It reduces computation and often latency, but it can lower the maximum batch size or context length that fits on a GPU.

Why cache keys and values but not queries?

The current token’s query is used once to read the context. Earlier queries have already produced their outputs and are never needed again. Earlier keys and values, however, are read by every later token, so caching them removes repeated work.

How would you reduce KV-cache pressure in a serving system?

I would first measure tokens per request, active concurrency, and cache bytes per sequence. Then I would consider GQA or MQA, lower-precision KV storage, paged allocation, strict context limits, prompt-prefix reuse, and batching policies. The right choice depends on whether the bottleneck is GPU memory capacity, memory bandwidth, prefill compute, or decode latency.

Say this in the interview

“A KV cache stores each layer’s keys and values for tokens already processed, so autoregressive decoding computes only the new token and attends to the stored context instead of rebuilding the prefix; that makes per-step attention linear in context length, at the cost of cache memory that grows with tokens and concurrency.”

Learn it properly KV cache & continuous batching

Keep practising

All NLP & LLMs questions

Explore further