KV Cache

Introduction

During autoregressive generation a transformer recomputes attention against the entire sequence at every step, and most of that work is redundant: the keys and values for tokens already processed are deterministic and do not change during the run. The KV cache is the optimization that exploits this redundancy. The rest of this chapter is about what is being cached, why caching it is mathematically safe, how much memory it costs, and the techniques that have grown up around it to keep that cost in check.

As you generate more text, GPU memory usage grows. For a single short request the model weights still dominate, but as context length and batch size grow, the KV cache can exceed the weight footprint and becomes the limiting factor in production serving workloads.

Long-context requests are also more expensive to serve in practice, which shows up indirectly in inference provider pricing. The KV cache is one of several reasons (memory bandwidth, scheduling, batching headroom all come into play), but it is the one that can comfortably scale context length where every prompt token has its keys and values cached, holding GPU memory for the entire duration of the request.

How LLMs Generate Text

Before we can understand what the KV cache solves, we need to understand how language models generate text. The process is called autoregressive decoding, and it works one token at a time.

Formally, the model computes a conditional probability distribution at each step:

The conditional is on every token from 1 to t-1, not just the last one. A transformer does not predict the next token in isolation; it attends to the full history at every step. That is what lets these models stay coherent across thousands of tokens, but it is also why every generation step has to chew through the entire sequence again.

Interactive: Generation Loop

Click Next Token to step through autoregressive generation. Every token in the context gets reprocessed on each step, and the attention matrix gains one row and one column per token.

The Quadratic Problem

Without any optimization, each generation step runs the entire sequence through the transformer. At step 1 we process 1 token. At step 2 we process 2. At step t we process t tokens. The problem compounds at the attention layer.

Inside self-attention, every token attends to every other token. That means computing an n × n attention matrix at each step. But we already computed the attention scores for tokens 1 through t-1 at the previous step. We throw all of that away and start from scratch.

For a 1,000-token response, this is roughly 333 million attention score computations. For 4,096 tokens, about 22 billion. The cubic scaling makes long contexts progressively more punishing.

Interactive: The Quadratic Problem

Compare the work done with and without KV caching. The gap between the two curves widens fast as the sequence gets longer.

Attention Review: Q, K, V

Each token embedding is projected into three vectors via learned weight matrices, then attention is computed as a scaled dot product:

During decoding, Q, K, and V have different shapes. This asymmetry is why caching works:

The Key Observation: K and V Stay Valid

K and V are not immutable in any deep mathematical sense; they happen to stay valid throughout causal decoding because of two things at once: the weights of the model are frozen at inference time, and the causal mask prevents any future token from feeding back into a past one's hidden state. Together those two facts mean that once you have computed K_i and V_i at some layer, no later step in the decode will ever need a different value for them. Look at how K_i and V_i are formed at layer ℓ:

Once we compute K_i and V_i for token i at a given layer, they will never change. The input to the projection at layer ℓ is the hidden state h_i^(ℓ) that token i carries out of the previous layer, and under causal attention that hidden state is itself fixed the moment token i has been processed (no future token can flow back into it). The weight matrices W_K and W_V are fixed too, since training is over. With both factors frozen, K_i = h_i^(ℓ) · W_K is a deterministic value that will be identical every time it is computed.

Q is the only piece that changes at each step. Q_t = h_t^(ℓ) · W_Q is recomputed from the new token's hidden state. K and V for every previously processed token are bitwise identical to what we got when we first ran them through the model.

Interactive: Q vs K/V Roles

Q is freshly computed at every step from the new token. K and V from previous tokens just sit in the cache, untouched. Step through to see the asymmetry.

How KV Cache Works

With the key observation established, the implementation is straightforward. Allocate a buffer to store K and V vectors as we compute them, and reuse them on every subsequent step instead of recomputing.

# Initialize empty cache for each layer
K_cache = []   # grows to shape [t, d_k]
V_cache = []   # grows to shape [t, d_v]

for each new token x_t:
    # Only compute projections for the NEW token
    q_t = x_t @ W_Q     # shape [1, d_k]
    k_t = x_t @ W_K     # shape [1, d_k]
    v_t = x_t @ W_V     # shape [1, d_v]

    # Append to the cache
    K_cache.append(k_t)  # now shape [t, d_k]
    V_cache.append(v_t)  # now shape [t, d_v]

    # Full attention using cached context
    scores  = q_t @ K_cache.T / sqrt(d_k)  # [1, t]
    weights = softmax(scores)               # [1, t]
    output  = weights @ V_cache             # [1, d_v]

Each step now does work proportional to t (one Q against t cached K/V pairs), not t². We compute K_t and V_t once, store them, and never touch them again.

These complexities refer to attention operations only. The MLP and LayerNorm sublayers still cost O(d_model) per token per step, which sums to O(T) per layer per decode step with the cache in place. The cache changes the attention term from cubic to quadratic; it does not make the rest of the transformer free.

Interactive: Decoding Animation

Step through the decode and see what stays and what gets recomputed. K and V land in the cache once per token; Q is computed fresh at every step.

Per-Layer, Per-Head: The Real Structure

So far we have talked about a single attention layer with a single head. Real transformers have many layers and many attention heads. The KV cache exists separately for every layer and every head.

Self-attention is the only operation where token i can see token j. Everything else is pointwise. That is why the KV cache only lives inside attention: it is the only place where past context needs to be retrieved. When using KV cache, the FFN and LayerNorm layers process only the new token, doing a constant amount of work per step.

In multi-head attention, each head has its own W_Q, W_K, and W_V projection matrices. Each head produces its own K and V vectors, which must be cached separately. For a model with L layers and H heads per layer, the total cache is L × H separate K and V matrices:

Interactive: Layer Architecture

Here is where the cache actually lives in the transformer block. Embeddings, LayerNorm, and the feed-forward sublayer touch one token at a time. Attention is the only place that needs every previous K and V, which is why it is the only place the cache exists.

Prefill vs Decode

Every LLM inference request has two distinct phases. Understanding them is essential for reasoning about latency and system design.

During prefill, the model processes all prompt tokens simultaneously, just as it does during training. This is efficient because the GPU can fill its compute units with large matrix operations over the full sequence. During decoding, however, we generate one token at a time. Each decode step reads the entire KV cache but writes only a tiny amount. This pattern makes decoding memory-bandwidth bound rather than compute-bound.

Interactive: Latency Comparison

Watch the difference between prefill and decode phases.

Memory Cost: The Math

The KV cache memory cost can be computed exactly. For standard multi-head attention:

Since H × d_head = d_model, the formula simplifies in the standard case. But for models using Grouped Query Attention (GQA), you replace H with H_kv, the number of KV heads, which is smaller:

Interactive: Memory Calculator

Configure model parameters and see how KV cache memory scales.

The Space-Time Tradeoff

The KV cache is a textbook space-time tradeoff. You spend memory to avoid spending compute. The complexity numbers make this concrete:

This is the same idea as dynamic programming. Instead of recomputing overlapping subproblems, store their results. Computing Fibonacci naively recomputes fib(2) an exponential number of times. Cache the results and each value is computed once. The KV cache applies exactly this principle to attention: each K_i and V_i is computed once and reused in all subsequent steps.