arthrop0d

Brief articles on varied topics.

Transformer KV Cache During Text Generation

artificial intelligenceperformance

A transformer KV cache stores the attention keys and values already computed for earlier tokens, separately for each transformer layer. During autoregressive generation—producing one token at a time—it lets the model reuse those tensors instead of calculating them again for every new token.

That trade-off is central to inference performance: the cache reduces repeated computation, but its size grows with the context and can become the main consumer of GPU memory.

The problem: each token depends on everything before it

In autoregressive generation, a model predicts the next token from the tokens it has already seen. For example, after processing The server returned, it predicts the next token, then uses the result to predict the token after that.

A transformer layer uses self-attention, a calculation that lets each token combine information from other tokens in the same sequence. Attention represents each token with three vectors:

To generate the next token, the new token's query is compared with the keys for the previous tokens. The resulting weights are used to combine their values.

Without a cache, the model would run the whole sequence through every layer again after each generated token. Generating token 101 would recalculate the representations for tokens 1 through 100; generating token 102 would recalculate them again, along with token 101. The output is correct, but most of the work produces information the model already computed on the previous step.

The cache exists to turn “recalculate the entire prefix” into “calculate only the new token and look up the prefix’s attention data.”

What the cache contains

For every transformer layer, the cache normally holds two tensors:

It does not usually store the model's complete hidden state for every token, nor does it store the generated text itself. It stores the intermediate attention inputs that later tokens need.

A simplified cache shape is:

[layer, batch item, KV head, sequence position, head dimension]

There are separate key and value tensors, so the cache contains two such sets. A head is one independent attention subspace inside a layer; the head dimension is the number of numeric elements in that subspace. Real implementations may arrange or flatten these dimensions differently, but the contents are equivalent.

The cache is generally kept in the model's inference device memory, often GPU memory, and commonly uses a reduced numeric format such as FP16 or BF16. It is associated with a particular sequence: entries from one conversation or request cannot be reused for an unrelated prompt unless the token prefix and all relevant model settings match.

How reuse works during generation

Generation has two useful phases.

1. Prompt processing, or prefill

The model first processes the input prompt, sometimes called the prefill phase. It can process the prompt's tokens together using the transformer’s parallel operations. At each layer, it computes the keys and values for those tokens and places them in the cache.

For a prompt containing 1,000 tokens, the cache now contains 1,000 positions at every layer. The model also produces the information needed to choose the first new token.

2. Token-by-token decoding

During decoding, the model processes one newly generated token at a time. At each layer it:

  1. Computes the new token's hidden representation.
  2. Projects that representation into a new query, key, and value.
  3. Compares the new query with all cached keys.
  4. Combines the cached values according to those attention weights.
  5. Appends the new key and value to the cache.

The new query needs to attend to the earlier context, so the attention operation still reads all cached positions. But the model does not need to recreate their keys and values. It computes only the new position's projections and attention output.

The new token's key and value are appended because future tokens may need to attend to it. The query is normally used once for the current step and is not retained.

This is why a serving log may distinguish prompt tokens or prefill from generated tokens or decode. It is also why a setting named use_cache, past_key_values, or similar appears in transformer inference systems: it controls whether these prior key/value tensors are retained and supplied to later steps. Exact APIs vary by framework.

Why context length becomes a memory problem

The cache grows linearly with the number of tokens in the context. A useful approximation for its size is:

bytes = 2 × layers × KV_heads × head_dimension × tokens × bytes_per_element

The first 2 accounts for keys and values. Batch size multiplies the result, as does the number of simultaneously active requests.

For example, a model with 32 layers, 32 key/value heads, 128 elements per head, an 8,192-token context, and two bytes per element needs approximately:

2 × 32 × 32 × 128 × 8192 × 2 bytes ≈ 4 GiB

That is cache memory for one sequence, before counting the model weights, temporary computation, and runtime overhead. A larger context or a larger batch can exhaust GPU memory even when the model weights fit comfortably.

Some models use multi-query attention or grouped-query attention. These designs use fewer key/value heads than query heads, so they reduce the cache size while retaining many query heads. The formula uses the smaller number of KV heads, which is why this architectural choice matters to serving cost.

The cache avoids repeated computation, but it does not make attention free: every new query still has to read the cached history. As the context grows, that read and the associated attention work also grow. In practice, long-context inference is therefore constrained by both memory capacity and memory bandwidth.

The KV cache is best understood as a deliberate exchange: store the past attention state once, reuse it for every subsequent token, and pay for that storage for as long as the request remains alive.

← All articles