KV Cache › Worked Example: Tracing a 4-Token Generation

To make the abstract formulas concrete, let's trace through a minimal numerical example from start to finish. We will use a toy model with $d_{\text{model}} = 4$, one attention head ($H = 1$), and $d_k = d_v = 4$. The prompt is three tokens: "The cat sat." We will then generate one additional token and track every vector and matrix involved.

Suppose the three prompt tokens produce input embeddings (after the embedding layer and any positional encoding):

$$x_1 = [1.0, 0.5, 0.2, 0.8], \quad x_2 = [0.3, 1.0, 0.6, 0.1], \quad x_3 = [0.7, 0.4, 1.0, 0.3]$$

For the projection matrices, we use simplified illustrative weights. In practice these are dense learned matrices, but here we treat them as identity-like for clarity:

$$W_K = W_V = I_4 \quad \text{(identity matrix)}$$ $$W_Q = 0.5 \cdot I_4$$

This means the key and value for each token are just that token's embedding, and the query is half the embedding. The simplification does not change the structural argument at all.

We compute the keys and values for all three tokens simultaneously:

$$\begin{aligned} k_1 &= x_1 W_K = [1.0, 0.5, 0.2, 0.8] \\ k_2 &= x_2 W_K = [0.3, 1.0, 0.6, 0.1] \\ k_3 &= x_3 W_K = [0.7, 0.4, 1.0, 0.3] \end{aligned}$$ $$\begin{aligned} v_1 &= x_1 W_V = [1.0, 0.5, 0.2, 0.8] \\ v_2 &= x_2 W_V = [0.3, 1.0, 0.6, 0.1] \\ v_3 &= x_3 W_V = [0.7, 0.4, 1.0, 0.3] \end{aligned}$$

We store these in the cache immediately:

$$K_{\text{cache}} = \begin{bmatrix} 1.0 & 0.5 & 0.2 & 0.8 \\ 0.3 & 1.0 & 0.6 & 0.1 \\ 0.7 & 0.4 & 1.0 & 0.3 \end{bmatrix}, \quad V_{\text{cache}} = \begin{bmatrix} 1.0 & 0.5 & 0.2 & 0.8 \\ 0.3 & 1.0 & 0.6 & 0.1 \\ 0.7 & 0.4 & 1.0 & 0.3 \end{bmatrix}$$

(In this toy example the keys and values are equal since $W_K = W_V = I$; in real models they differ entirely.)

Suppose the model samples token 4 from the distribution produced after the prefill phase, and that token 4 has embedding $x_4 = [0.5, 0.9, 0.3, 0.6]$. To predict token 5, we need to compute attention from position 4 over all previous positions.

Step 1: Compute projections for token 4 only.

$$q_4 = x_4 W_Q = 0.5 \cdot [0.5, 0.9, 0.3, 0.6] = [0.25, 0.45, 0.15, 0.30]$$ $$k_4 = x_4 W_K = [0.5, 0.9, 0.3, 0.6], \quad v_4 = x_4 W_V = [0.5, 0.9, 0.3, 0.6]$$

The key insight is that we compute exactly three vectors ($q_4$, $k_4$, $v_4$) for this step, not $4 \times 3 = 12$ vectors. The keys and values for tokens 1 through 3 remain exactly as computed during prefill.

Step 2: Append $k_4$ and $v_4$ to the cache.

$$K_{\text{cache}} \leftarrow \begin{bmatrix} 1.0 & 0.5 & 0.2 & 0.8 \\ 0.3 & 1.0 & 0.6 & 0.1 \\ 0.7 & 0.4 & 1.0 & 0.3 \\ 0.5 & 0.9 & 0.3 & 0.6 \end{bmatrix}$$

Step 3: Compute attention scores for token 4's query against all four cached keys.

The raw dot products are:

$$\begin{aligned} s_1 &= q_4 \cdot k_1 = (0.25)(1.0) + (0.45)(0.5) + (0.15)(0.2) + (0.30)(0.8) = 0.25 + 0.225 + 0.03 + 0.24 = 0.745 \\ s_2 &= q_4 \cdot k_2 = (0.25)(0.3) + (0.45)(1.0) + (0.15)(0.6) + (0.30)(0.1) = 0.075 + 0.45 + 0.09 + 0.03 = 0.645 \\ s_3 &= q_4 \cdot k_3 = (0.25)(0.7) + (0.45)(0.4) + (0.15)(1.0) + (0.30)(0.3) = 0.175 + 0.18 + 0.15 + 0.09 = 0.595 \\ s_4 &= q_4 \cdot k_4 = (0.25)(0.5) + (0.45)(0.9) + (0.15)(0.3) + (0.30)(0.6) = 0.125 + 0.405 + 0.045 + 0.18 = 0.755 \end{aligned}$$

Step 4: Scale by $\sqrt{d_k} = \sqrt{4} = 2.0$.

$$\hat{s} = [0.745 / 2, \ 0.645 / 2, \ 0.595 / 2, \ 0.755 / 2] = [0.3725, \ 0.3225, \ 0.2975, \ 0.3775]$$

Step 5: Apply softmax to get attention weights.

First compute the exponentials:

$$e^{0.3725} \approx 1.451, \quad e^{0.3225} \approx 1.380, \quad e^{0.2975} \approx 1.346, \quad e^{0.3775} \approx 1.459$$

Sum: $1.451 + 1.380 + 1.346 + 1.459 = 5.636$

Normalized weights:

$$a = [1.451/5.636, \ 1.380/5.636, \ 1.346/5.636, \ 1.459/5.636] \approx [0.257, \ 0.245, \ 0.239, \ 0.259]$$

Notice that token 4's own key (position 4) receives slightly more weight than the others, which is typical for recent tokens in causal attention.

Step 6: Compute the attention output as a weighted sum of cached values.

$$o_4 = 0.257 \cdot v_1 + 0.245 \cdot v_2 + 0.239 \cdot v_3 + 0.259 \cdot v_4$$ $$\begin{aligned} o_4[0] &= 0.257(1.0) + 0.245(0.3) + 0.239(0.7) + 0.259(0.5) = 0.257 + 0.073 + 0.167 + 0.130 = 0.627 \\ o_4[1] &= 0.257(0.5) + 0.245(1.0) + 0.239(0.4) + 0.259(0.9) = 0.129 + 0.245 + 0.096 + 0.233 = 0.703 \\ o_4[2] &= 0.257(0.2) + 0.245(0.6) + 0.239(1.0) + 0.259(0.3) = 0.051 + 0.147 + 0.239 + 0.078 = 0.515 \\ o_4[3] &= 0.257(0.8) + 0.245(0.1) + 0.239(0.3) + 0.259(0.6) = 0.206 + 0.025 + 0.072 + 0.155 = 0.458 \end{aligned}$$

So $o_4 \approx [0.627, 0.703, 0.515, 0.458]$.

What did we compute in this decode step? Three vector projections (for token 4 only), four dot products, a softmax over four values, and four weighted additions. Without the cache, we would have recomputed all projections for tokens 1, 2, and 3 as well. That is 9 additional projection operations for a sequence of length 4. For a sequence of length 100, the savings would be 297 additional projections avoided. For length 1000, the savings grow to 2997 projections avoided in that single step alone.