KV Cache

What a KV cache is

A transformer generates text one token at a time. To produce the next token, self-attention compares the current query against the Key and Value vectors of every token seen so far. Those Key and Value vectors depend only on the tokens that already exist, so they do not change once computed. The KV cache is the buffer that holds them, indexed by layer and attention head, so the model can append to it rather than rebuild it on every step.

Without a cache, each decoding step would recompute the Keys and Values for the entire prefix. With a cache, the model computes K and V once per token, writes them into the cache, and at every later step reads them back. Only the single new token needs fresh K and V projections.

• Caches the Key (K) and Value (V) tensors, not the queries, which are different every step.
• Stored per transformer layer and per attention head.
• Read during attention, appended to after each generated token.
• Exists only at inference; training does full parallel attention and needs no incremental cache.

Why it exists: turning O(n squared) into O(n) per step

Generating token number n requires attention over all n prior tokens. If the model recomputed every prior Key and Value at each step, step n would cost work proportional to n, and generating a sequence of length n would cost roughly n squared total. That quadratic blowup makes long generations impractical.

Caching removes the redundancy. Because past Keys and Values never change, the model keeps them and only projects the new token, so each step does work proportional to the current length n through the attention read, with no repeated projection of the prefix. The cost shifts from recomputation to memory: the cache must hold every past K and V.

• No cache: each step re-derives the full prefix, total generation cost scales with n squared.
• With cache: the prefix is reused, the new token's K and V are computed once.
• The tradeoff is compute saved in exchange for memory consumed.

Prefill versus decode

Inference splits into two phases. Prefill processes the entire prompt at once. Because every prompt token is known up front, the model computes all of their Keys and Values in a single parallel pass and fills the cache. This phase is compute-bound and produces the first output token.

Decode then runs one token at a time. Each step reads the existing cache, computes K and V for the single newest token, appends them, and emits the next token. Decode is typically memory-bandwidth-bound because it moves the growing cache through the attention units while doing comparatively little arithmetic per step.

• Prefill: parallel pass over the whole prompt, fills the cache, compute-bound.
• Decode: one token per step, appends one K and one V per layer per head.
• Time to first token reflects prefill; per-token latency reflects decode.

Memory cost: linear in sequence length and batch

The cache must store a Key and a Value for every token, in every layer, for every attention head, for every sequence in the batch. That makes total size grow linearly with sequence length and with batch size. Long contexts and large batches push systems into being memory-bound on the KV cache rather than on model weights.

The standard size estimate multiplies these dimensions together, with a factor of 2 because both K and V are stored. Using the number of key/value heads matters for models with grouped attention, where the K/V head count is smaller than the query head count.

• Grows linearly with both sequence length S and batch size B.
• Independent of the prompt's content; only its length matters.
• For a 70B-class model at BF16 with a 128K context, the cache alone can reach tens of gigabytes per sequence.
• Once weights are loaded, remaining VRAM caps how many tokens and concurrent requests fit.

How decoding appends to the cache

The decode loop is short: project the new token, append its K and V, attend over the full cache, sample the next token, repeat. The pseudocode below sketches the per-step append for a single layer.

# Per-layer decode step (single sequence, simplified)
# cache_k, cache_v: tensors of shape [n_kv_heads, seq_so_far, head_dim]

def decode_step(x, cache_k, cache_v, Wq, Wk, Wv, Wo):
    q = x @ Wq                      # query for the new token only
    k = x @ Wk                      # new token's key
    v = x @ Wv                      # new token's value

    cache_k = append(cache_k, k)    # grow cache by one position
    cache_v = append(cache_v, v)

    scores = softmax(q @ cache_k.T / sqrt(head_dim))
    attn   = scores @ cache_v       # attend over the whole prefix
    return attn @ Wo, cache_k, cache_v

• Only the new token is projected each step; the prefix is read, not recomputed.
• The cache length grows by one position per generated token.
• Causal masking is implicit because the cache holds only past and current tokens.

Optimizations that shrink or manage the cache

Because the cache dominates inference memory, most serving optimizations target it. Multi-query attention (MQA) shares a single Key/Value head across all query heads, cutting the H_kv term to 1. Grouped-query attention (GQA) is the middle ground: query heads are split into groups that each share one K/V head, recovering most of the quality of full multi-head attention while keeping the cache small. Modern open models commonly ship with GQA.

Systems-level techniques attack waste and bandwidth. PagedAttention, the algorithm behind vLLM, stores the cache in non-contiguous blocks like operating-system virtual memory, which nearly eliminates fragmentation and lets requests share identical prefixes. KV-cache quantization stores K and V at lower precision such as 8-bit or 4-bit to cut bytes_per_param. Prompt and prefix caching reuses the cached Keys and Values of a shared prompt prefix across requests so prefill is not repeated.

• MQA: one shared K/V head, smallest cache, some quality cost.
• GQA: grouped K/V heads, interpolates between MQA and full multi-head attention.
• PagedAttention / vLLM: block-based memory management, low fragmentation, prefix sharing.
• Quantization: lower-precision K and V reduce bytes per parameter.
• Prefix / prompt caching: reuse a prefix's cached K and V to skip repeated prefill.