How do you implement the KV cache in a transformer decoder?

In this article

1. Before you start
2. Why interviewers ask this
3. Implementing the KV cache from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the PyTorch and machine learning coding round guide; the code builds on attention you may already have written.

Why interviewers ask this

A decoder that is fast but quietly wrong is worse than a slow one, and the KV cache is where that happens. Candidates report being asked to add a KV cache to a given transformer, often next to multi-head attention and grouped-query attention, in senior machine learning coding rounds at AI labs and large tech companies.

A weak answer says "we cache the keys and values so it is faster" and cannot say what shape the cache has, what the mask looks like on a decode step or why queries are not stored. A strong answer writes the append, gets the mask and position offsets right, and names the memory cost. A Staff-level answer adds that the cache stores only the key-value heads under grouped-query attention, and that serving systems manage the cache in blocks.

Implementing the KV cache from first principles

Decode step

A decode step is a pass of a language model that takes the newest token and produces the scores for the following token, so it adds exactly one token to the sequence; generating a reply of n tokens takes n decode steps one after another.

Cache position

The cache position is the index the next token will occupy, equal to the number of tokens already in the cache; it is the value used for the new token's positional encoding and the starting row of its attention mask.

In plain words

Imagine taking minutes of a meeting. A new sentence arrives. You do not rewrite the notes for every earlier sentence. You add one new note. Then you reread the old ones to see what the new sentence refers to. The analogy breaks in one place. The notes are exact, not a summary, so a cached model gives the same output as an uncached one.

In one precise sentence, a KV cache is a per-layer store of the keys and values of every token already processed. It lets a language model skip recomputing them at each decode step. The Hugging Face documentation describes it as storing key-value pairs from the attention layers of previously processed tokens.

Step 1: Prefill fills the cache from the prompt

Prefill is the first pass, in which the whole prompt goes through the model together and every layer stores the keys and values of every prompt token. Input: the prompt, shape (batch, prompt length, model width). Output: a cache per layer holding keys and values of shape (batch, KV heads, prompt length, head width), plus the scores for the first new token.

Prefill can run all prompt tokens in parallel with a causal mask, like training. The decode steps after it cannot, because each needs the token the previous one produced. Shazeer (2019) notes that incremental inference cannot be parallelised across the sequence and is slow because of the memory-bandwidth cost of repeatedly loading the large key and value tensors. You now know where the cache starts; the next step is how it grows.

Step 2: Project only the new token and append it

On each decode step, the layer projects only the newest token into a query, a key and a value, then concatenates the new key and value onto the stored ones. The update rule in the Hugging Face docs is the concatenation of the past keys with the new key, and the same for values.

Input: one token's hidden vector and the layer's cache. Output: a cache that is one position longer. The query is used once and thrown away, because only the newest token's query is needed to compute the representation that gives the following token. That is why it is a KV cache and not a query-key-value (QKV) cache. Without the cache, step t would recompute keys and values for all t tokens; with it, step t computes them for one. The worked example counts that difference. Next, the new query must meet the whole cache correctly.

Step 3: Offset the position and the mask by the cache length

The new token must be treated as sitting at the cache position, both in its positional encoding and in its attention mask. The Hugging Face snapshot for v4.52.3 gives the shapes: the attention weights have shape (new tokens, past plus new tokens), the mask covers past plus new length, and a cache holding 10 tokens means the next cache position is 10.

The bug is to feed the new token as position 0, as a fresh sequence would. Nothing crashes. The shapes are fine, but the new token's query, key and value all carry the wrong position, and the wrong key and value are stored, so every later step reads them too and the output drifts further from the uncached one. In the toy run, the correct offset gives outputs that agree with the uncached forward pass to nine decimal places, and position 0 gives different outputs.¹ Any positional scheme that changes keys must apply the true position before the key is stored, since cached keys are never recomputed.

Step 4: Sample, loop and keep one cache per layer and sequence

After attention, the model outputs scores for the following token; you pick a token (greedy, sampled, or one beam of beam search), append it to the output and start the next decode step with the same caches. The sampled token loops back into the next decode step. Each layer owns its own cache, and each sequence in a batch owns its own cache rows.

Input: the chosen token and every layer's cache. Output: the next token and caches one position longer. Below is the whole layer in PyTorch; the cache stores keys and values before they are repeated for grouped-query attention, so it holds KV heads, not query heads.

import math
import torch

class CachedAttention(torch.nn.Module):
    def __init__(self, d_model, n_q, n_kv, head_dim):
        super().__init__()
        self.n_q, self.n_kv, self.hd = n_q, n_kv, head_dim
        self.wq = torch.nn.Linear(d_model, n_q * head_dim, bias=False)
        self.wk = torch.nn.Linear(d_model, n_kv * head_dim, bias=False)
        self.wv = torch.nn.Linear(d_model, n_kv * head_dim, bias=False)
        self.wo = torch.nn.Linear(n_q * head_dim, d_model, bias=False)

    def forward(self, x, cache=None):
        # x: (batch, new, d_model); cache: None or (K, V), each (batch, n_kv, past, hd)
        b, new, _ = x.shape
        q = self.wq(x).view(b, new, self.n_q, self.hd).transpose(1, 2)
        k = self.wk(x).view(b, new, self.n_kv, self.hd).transpose(1, 2)
        v = self.wv(x).view(b, new, self.n_kv, self.hd).transpose(1, 2)
        if cache is not None:                      # step 2: append
            k = torch.cat([cache[0], k], dim=2)
            v = torch.cat([cache[1], v], dim=2)
        past = k.shape[2] - new                    # the cache position
        g = self.n_q // self.n_kv                  # query heads per KV head
        k_full = k.repeat_interleave(g, dim=1)
        v_full = v.repeat_interleave(g, dim=1)
        scores = q @ k_full.transpose(-2, -1) / math.sqrt(self.hd)
        rows = torch.arange(new, device=x.device).unsqueeze(1) + past   # step 3: offset
        cols = torch.arange(k.shape[2], device=x.device).unsqueeze(0)
        scores = scores.masked_fill(cols > rows, float("-inf"))
        out = torch.softmax(scores, dim=-1) @ v_full
        out = out.transpose(1, 2).reshape(b, new, -1)
        return self.wo(out), (k, v)                # return the grown cache

Call it with the whole prompt and cache=None for prefill, then with one token and the returned cache for each decode step. Add positions to x outside the layer, starting at the cache position. torch.softmax already uses the max-subtraction trick explained in numerically stable softmax; keep it if you write the softmax yourself. For the full model around this layer, see building generative pre-trained transformer (GPT) from scratch.

1 · Prefill2 · Append3 · Attend4 · LoopReplay

Step-by-step description

1. The prompt fills the cache.
2. One new key and value join the cache.
3. One query reads the whole cache.
4. The sampled token starts the next step.

• Prefill fills the KV cache from the prompt.
• Each decode step projects one new token and appends its key and value to the cache.
• One query attends over the whole cache, with the position and mask offset by the cache length.
• The sampled token loops back into the next decode step.

Worked example

A toy model makes the savings countable. It has 2 layers, 4 query heads, 2 KV heads, a head width of 8 and 16-bit numbers (all illustrative). Each token then adds 2 (keys and values) × 2 layers × 2 KV heads × 8 × 2 bytes = 128 bytes to the cache, and 8 tokens fill 1,024 bytes, held as keys of shape (2, 8, 8) per layer.¹

Generating 8 tokens with no cache recomputes keys and values for 1, then 2, up to 8 tokens: 36 token projections in total. With the cache it is 8. At 1,024 tokens the counts are 524,800 against 1,024, a ratio of 512.5.¹ The attention scores per head also fall: 204 entries computed without the cache over 8 steps against 36 with it, because the single query is scored only against the cached keys.¹ The check that matters is exactness: a numpy run with 5 prompt tokens and 3 decode steps matched the full uncached pass to nine decimal places, and its first decode step had attention weights of shape (1, 6).¹

The general formula is bytes = batch × sequence length × layers × 2 × KV heads × head width × bytes per number, which matches the arithmetic above.¹ It needs 0.5 GiB with 8 KV heads and 0.0625 GiB with 1 KV head.¹ The large configuration is illustrative, not a specific released model; read the real layer count and KV heads from a model's config.json before quoting a size.

Text description of the diagram

• One 4,096-token sequence needs 2 GiB with 32 KV heads.
• It needs 0.5 GiB with 8 KV heads and 0.0625 GiB with 1 KV head.

Trade-offs and edge cases

The KV cache trades memory for compute. The Hugging Face docs say the cache can occupy a large part of memory and become a bottleneck for long-context generation. Decoding stays bound by memory bandwidth, since attention still reads the whole cache every step; the cache removes recomputation, not reading.

Grouped-query attention (GQA) shrinks the cache by sharing each KV head among a group of query heads. Ainslie et al. (2023) define it as an intermediate number of key-value heads, more than one but fewer than the query heads, and report quality close to multi-head attention with speed comparable to multi-query attention (MQA), which uses a single key-value head. They report that MQA can lose quality, give a recipe that uptrains an existing multi-head checkpoint to MQA with 5% of the original pre-training compute, and show that GQA uptrained this way reaches quality close to multi-head attention at a speed comparable to MQA. Shazeer's multi-query proposal shares keys and values across all heads.

In serving, the cache of each request is large and changes size. Kwon et al. (2023) report that managing it inefficiently wastes memory through fragmentation and duplication, which limits batch size, and that their block-based vLLM system improves throughput by 2-4 times at the same latency, compared with systems such as FasterTransformer and Orca.

Text description of the diagram

• A sliding-window cache keeps only the last tokens.
• An offloaded cache moves most layers to the CPU and gives the same result as the default.
• A quantized cache is lossy, and a paged cache manages memory in blocks.

The Hugging Face docs describe a sliding-window cache that keeps only the last tokens, an offloaded cache that moves most layers to the central processing unit (CPU) and gives the same result as the default, and a quantized cache that is lossy and, per the v4.47.1 docs, comes at the cost of speed, while a paged cache manages memory in blocks. Padding adds another edge case: the cache position counts stored tokens regardless of padding, so left-padded batches need a mask that hides the pad entries.

Concatenate or preallocate?

The code above uses torch.cat, which allocates a new tensor and copies the whole cache on every step. That is memory traffic proportional to the current length per step, so quadratic over a generation, and the changing sizes fragment allocator memory. The first optimization to name is preallocation: reserve a buffer of the maximum length once and write each new key and value at the cache position, which the "Static Cache" section of the Hugging Face v4.47.1 page describes. The cost is reserving memory for the maximum length, and a mask that hides the unfilled slots.

To add a cache to an existing model, change four things: every attention layer takes and returns its cache; the model returns a list of per-layer caches; position ids start at the cache position; and the generate loop passes the caches back in.

How to answer it in an interview

The spoken structure: say what is stored, then the append, then the offset, then the memory trade-off. Rehearse the answer out loud and write the code once from memory.

Common mistakes

• "I cache queries, keys and values." (wrong) Only the newest token's query is used, so keys and values are stored. (right) "The cache holds keys and values; the query is used once."
• "The new token is position zero." (wrong) That drops the cache length from the position and the mask. (right) "Positions and mask rows start at the cache position."
• "The mask is the new length square." (wrong) The mask must cover past plus new keys. (right) "Weights are (new, past plus new), with row i blocking keys after past plus i."
• "The cache makes generation free." (wrong) Attention still reads the whole cache each step. (right) "It removes recomputation; decoding stays memory-bound."
• "I cache n_q heads." (wrong) Under GQA that wastes memory. (right) "I cache n_kv heads and repeat them per query group."

Follow-up questions

How do you implement the KV cache in PyTorch?

Keep a (K, V) pair per layer, shape (batch, KV heads, length, head width), and torch.cat the new key and value on the length axis each step, as in the code above. Pass the old cache in and return the new one.

How do you handle the KV cache with grouped-query attention?

Store keys and values with only the KV heads, then repeat each KV head for its group of query heads at attention time. The cache shrinks by the ratio of query heads to KV heads; in the large example that is a factor of 4.¹ See multi-head attention without library helpers for the base layer.

What changes in the mask and the positions when you add a KV cache?

The mask gains past extra columns and its rows start at the cache position, so a decode step has a single row with no blocked entries. The new token's positional encoding uses the cache position, not zero.

Why is the KV cache only used for inference?

Training processes a whole sequence in parallel, so there is nothing to reuse across steps, and the Hugging Face docs say caching should only be used for inference. Backpropagation also needs the full computation graph.

How big is the KV cache, and what do you do when it is too big?

Use the formula: batch × length × layers × 2 × KV heads × head width × bytes. Reduce it with fewer KV heads, a sliding window, offloading, quantization or block-based paging, each from the variants table above.

Does the KV cache change the model's output?

No, a plain cache is exact mathematically; in floating point the results agree up to rounding, as the nine-decimal toy check shows. It only reuses values that would be recomputed identically. Sliding-window and quantized caches do change what is attended to or stored, while offloading does not. The causal mask page explains why earlier keys never change.

Sources

1. How caching works (Hugging Face Transformers documentation) — Hugging Face, 26 Oct 2026
2. How caching works, v4.52.3 snapshot (mask shape and cache position) — Hugging Face, 26 Oct 2026
3. Best Practices for Generation with Cache (v4.47.1) — Hugging Face, 26 Oct 2026
4. Best Practices for Generation with Cache (v4.51.1, sliding-window cache) — Hugging Face, 26 Oct 2026
5. Fast Transformer Decoding: One Write-Head is All You Need (arXiv 1911.02150) — arXiv (Shazeer), 6 Nov 2019
6. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (arXiv 2305.13245) — arXiv (Ainslie et al.), 22 May 2023
7. Efficient Memory Management for Large Language Model Serving with PagedAttention (arXiv 2309.06180) — arXiv (Kwon et al.), 12 Sep 2023