arthrop0d

Brief articles on varied topics.

Transformer Key-Value Cache During Generation

artificial intelligenceperformance

A transformer key-value (KV) cache stores the attention keys and values computed for tokens that have already been processed. During autoregressive generation—producing one token, then using it to produce the next—the model reuses those stored tensors instead of recomputing them for the whole preceding text.

The problem it solves

A language model generates text one token at a time. To choose the next token, it looks at the input prompt and all tokens it has generated so far; this sequence is the context.

For example, generating the third word requires the model to process the prompt and the first two generated words. Generating the fourth requires the prompt and all three earlier words. Without a cache, an implementation could run the transformer over the entire sequence again at every step:

  1. Process the prompt and generate token 1.
  2. Process the prompt plus token 1 and generate token 2.
  3. Process the prompt plus tokens 1 and 2 and generate token 3.
  4. Repeat.

The earlier part of the sequence has not changed, so most of that work is redundant. The KV cache preserves the intermediate attention data that remains valid when a new token is appended.

What attention needs

Each transformer layer uses self-attention, a calculation in which every token can use information from other tokens in the context. For each input token, a layer produces three vectors:

The layer compares the new token's query with the keys of earlier tokens. These comparisons produce attention weights, and the layer uses those weights to form a weighted combination of the corresponding values.

For a newly generated token, the model therefore needs:

The keys and values for earlier tokens are exactly what the KV cache stores. It normally does not store earlier queries, because queries from earlier generation steps are no longer needed to calculate the next token's attention.

How generation uses the cache

Generation has two distinct phases.

1. Prefill the prompt

The model first processes the supplied prompt, often called the prefill phase. It runs the prompt through every transformer layer and computes keys and values for every prompt position. Those tensors are placed in the cache, organized by layer and usually by attention head.

The model then uses the final prompt position to produce the first next-token prediction.

2. Decode one token at a time

For each generated token, the model runs the new token through the layers. At each layer it computes a query, key, and value for that one new position:

for token in generated_tokens:
    for layer in transformer_layers:
        q, k, v = layer.project(token)
        key_cache[layer].append(k)
        value_cache[layer].append(v)
        token = layer.attend(q, key_cache[layer], value_cache[layer])
    token = choose_next_token(token)

The names in this pseudocode describe the operations, not a universal library API. The important sequence is that the new key and value are appended, then the new query attends to the complete cached sequence. The cache now includes the new token and can be used at the next step.

Because a transformer uses a causal mask, a token is prevented from seeing future tokens. Adding a later token therefore cannot change the earlier tokens' keys or values. That is the property that makes caching correct: the earlier attention inputs are stable.

Position information, such as rotary positional embeddings, is also applied consistently for each token position. A cached key represents its original position, so it remains usable when later positions are appended.

The cache avoids recomputing the earlier layers' representations and their K/V projections. It does not eliminate the need to compare the new query with all earlier keys.

Why memory grows with context length

A cache keeps two tensors—K and V—for every token, every transformer layer, and every attention head or key/value head. If the context gets twice as long, it contains roughly twice as many cached numbers.

A simplified memory estimate is:

cache bytes ≈ 2 × layers × tokens × KV heads × head dimension × bytes per number

The first 2 accounts for keys and values. For example, with 32 layers, 8,192 tokens, eight key/value heads, a head dimension of 128, and two-byte values such as FP16:

2 × 32 × 8192 × 8 × 128 × 2 ≈ 1 GiB

That is for one sequence. Batching several requests multiplies the requirement, and longer context limits multiply it again. A model using a separate key/value head for every attention head uses more memory than one using multi-query attention or grouped-query attention, designs that share K/V heads across several query heads.

The KV cache is therefore often a major part of inference memory, alongside the model's weights and temporary working space. Systems may reduce its cost with lower-precision cache values, shared K/V heads, or paged allocation that manages cache blocks efficiently. These techniques change how the cache is stored; they do not change its basic purpose.

A cache makes decoding substantially faster by removing repeated prefix computation, but it does not make attention independent of context length. At every new step, the new query still attends to an increasingly long list of cached keys and values. The cache trades repeated computation for persistent memory that grows with the number of tokens in the context.

← All articles