KV cache: reuse the work behind a token

A KV cache retains the key and value vectors computed for earlier tokens in each attention layer. When the next token arrives, the model can use those saved vectors rather than process the whole prefix again.

What stays, what changes

During causal generation, an earlier position cannot depend on later tokens. Its already-computed keys and values therefore remain reusable as the sequence grows. The new token supplies its own query, key and value. Its query attends to the retained keys, and the result mixes the corresponding values.

The cache stores temporary state for a passage. The model’s learned weights can be shared across many passages; each passage may need different cached state.

Continue with one new token

With torch and transformers installed, this example processes a prompt, chooses one token, then passes only that token and the saved cache into the next model call:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
inputs = tokenizer("The capital of France is", return_tensors="pt")

with torch.inference_mode():
    first = model(**inputs, use_cache=True)
    next_id = first.logits[:, -1].argmax(-1, keepdim=True)
    second = model(
        input_ids=next_id,
        past_key_values=first.past_key_values,
        use_cache=True,
    )
    following_id = second.logits[:, -1].argmax(-1).item()

The second call receives one token instead of the entire prompt. following_id predicts the token after it. This example has no padding; batched padded inputs also need the corresponding full attention mask.

Caching removes repeated prefix processing, but the new query still reads earlier keys and values. Ordinary full-attention cache capacity grows linearly with retained tokens at fixed architecture and precision. More simultaneous passages multiply that requirement.

Sharing identical prefixes and managing cache blocks can reduce duplication. Reusing state across separate API requests is a serving feature built on the same principle. Our RAM story follows what happens when many agents need their contexts back.

Related: attention, prefill, prompt caching.