Understanding Transformer Decoding with a KV Cache

Hi, I'm Matsuken, a data scientist at a Japanese technology company.

In this article, I'll explain how decoding works in a Transformer that uses a key-value (KV) cache.

This is an English version of my original Japanese article, which I originally wrote and published on Qiita.

Transformer inference consists of two stages:

Prefill: Process all input tokens—the prompt—in parallel. During this stage, the model computes the keys ( $K$ ) and values ( $V$ ) and stores them in a cache.
Decode: Generate text one token at a time. At each step, the model reads and reuses the previously computed keys and values, then appends one new row—the key and value for the latest token—to the cache.

For an explanation of the first stage, prefill, see my previous article. The stored $K$ and $V$ tensors produced by this process are collectively called the KV cache.

Given an input $X$ , an attention layer computes the queries, keys, and values as follows:

\begin{aligned} Q &= XW_Q \cr K &= XW_K \cr V &= XW_V \end{aligned}

It then produces the output $Y$ :

Y = \operatorname{softmax}\left(\frac{QK^{\mathrm{T}}}{\sqrt{d}} + M\right)V

The process is illustrated below.¹

Here, $n$ is the number of tokens and $d$ is the token representation dimension.

In a decoder-only Transformer, this computation takes place in every Multi-Head Attention block: each block receives $X$ and produces $Y$ .²

Prefill

During the initial pass, the input $X$ normally contains multiple tokens, so it is an $n \times d$ matrix. This is the prefill stage covered in my previous article.

The next stage—decode—is where the KV cache comes into play.

Decode

Prefill processes the prompt and determines the first generated token. In each subsequent decode step, the model receives the token generated in the previous step and autoregressively predicts the next token.

Conceptually, we could append every generated token to the original input sequence and process the entire sequence again. In practice, this would repeatedly perform the same computations. Instead, the model computes only the new query, key, and value vectors for the latest token. It reuses the keys and values for all previous tokens from the KV cache. This reuse is the purpose of the KV cache.

Suppose the prompt is:

[BOS] two kids are playing in a swimming pool with a green colored crocodile.

and the first generated token is “Then.”

Because $n$ tokens have already been processed, the model uses the representation of token $n+1$ to compute three new vectors: $q_{n+1}$ , $k_{n+1}$ , and $v_{n+1}$ .

Next, the model updates $K$ and $V$ . It reuses the entries through position $n$ from the cache and appends $k_{n+1}$ and $v_{n+1}$ , respectively.

The new query $q_{n+1}$ is multiplied by the transpose of the updated key matrix, which consists of the cached keys followed by $k_{n+1}$ . This produces the attention-score vector $s_{n+1}$ .

The model applies softmax to $s_{n+1}$ and multiplies the resulting attention weights by the updated value matrix—the cached values followed by $v_{n+1}$ . The result is $y_{n+1}$ .

We have now obtained the $y_{n+1}$ required by this layer.

Although the explanation above follows the operations step by step, we can also view the computation as part of a conceptual full attention matrix. It has three regions: entries that were computed in previous steps, future positions excluded by the causal mask, and the newly computed row for the current step.

The complete decode-step computation can be summarized as follows:

You might wonder whether $y_{n+1}$ alone is enough, or whether the model also needs $y_0, \ldots, y_n$ . Recall that this step required only $x_{n+1}$ because the earlier keys and values were already cached. Therefore, computing only $y_{n+1}$ is sufficient.

As long as $K$ and $V$ are cached, each subsequent step needs only the representation of the newest token to produce the next-token output. That is what makes decoding with a KV cache so efficient.

The model maintains a separate KV cache for every Transformer layer. If the model has $N$ layers, it therefore has $N$ cache pairs:

(K_{\mathrm{cache}}^{(1)}, V_{\mathrm{cache}}^{(1)}), \ldots, (K_{\mathrm{cache}}^{(N)}, V_{\mathrm{cache}}^{(N)})

This efficiency comes with a trade-off: as the number of tokens—the context-window length—increases, the memory required to store the key and value caches also grows. Several techniques address this issue, but I'll save those for another article.

To keep the diagrams simple, they omit division by the scale factor $\sqrt{d}$ . ↩
To focus on the KV-cache mechanism, the rest of this article represents Multi-Head Attention as a single attention head. ↩