KV Cache by hand

Table of Contents

• Motivation
• What is the KV Cache?
• Setup
• Scenario 1: Generation WITHOUT KV Cache
• Scenario 2: Generation WITH KV Cache
• The Compute vs. Memory Trade-off
• Code
• Appendix A: Worked example with 2 layers

Motivation

The Key-Value (KV) cache is an important optimization in Large Language Model (LLM) inference. Working through the math helps with:

• Building intuition: Understand why autoregressive generation (predicting the next word) contains massive amounts of redundant math without caching.
• Understanding the compute-memory trade-off: The cache turns each step's projections from a matrix-matrix multiply into a matrix-vector multiply. That's far fewer FLOPs — but it also means you're now loading far more memory per unit of compute, which is exactly what makes decoding memory-bandwidth-bound. And the cache itself competes with the weights for VRAM.
• Debugging Performance: When estimating serving costs, knowing the size and behavior of the KV cache is essential for calculating maximum batch sizes and throughput.

This article is written with the assistance of AI.

---

What is the KV Cache?

LLMs generate text autoregressively, meaning they predict the sequence one token at a time. To predict token N, the model needs to look at all previous tokens from 1 to N-1.

In the Transformer attention mechanism, every token is projected into a Query (Q), Key (K), and Value (V) vector.

• The Query represents what the current token is "looking for".
• The Key represents what a token "contains".
• The Value represents the actual "content" the token will contribute to the output.

When generating a new token, the model calculates a new Query, Key, and Value for that specific token. However, its new Query needs to attend to the Keys and Values of all past tokens.

If we don't use a KV cache, we have to recalculate the Keys and Values for every single previous token at every single generation step. The KV cache simply stores the Keys and Values of previous tokens in memory so we only ever have to compute the projection for the newest token.

Quick question: why do we not cache Q? Ans: because of causal masking, a token's query is used exactly once, at the step that token is processed, and is never needed again. Keys and values, by contrast, are read at every subsequent step.

---

Setup

We will walk through the attention calculation for the decode step that predicts the 4th token.

Assume our embedding dimension is $d = 2$ . We have three weight matrices for our attention head: $W_q$ , $W_k$ , $W_v$ .

Let our weight matrices be:

$$W_q = \begin{pmatrix} 1 & 0 \newline 0 & 1 \end{pmatrix}, \quad W_k = \begin{pmatrix} 1 & 1 \newline 0 & 1 \end{pmatrix}, \quad W_v = \begin{pmatrix} 2 & 0 \newline 0 & 2 \end{pmatrix}$$

Assume our sequence currently consists of 2 tokens, $x_1$ and $x_2$ , and we just generated the newest token $x_3$ :

• $x_1 = [1, 0]$
• $x_2 = [0, 1]$
• $x_3 = [1, 1]$ (The new token)

We now need to perform the attention calculation for $x_3$ to predict the next token, $x_4$

---

Scenario 1: Generation WITHOUT KV Cache

Without a KV cache, the model has amnesia. It only knows that the current sequence is three tokens long. To calculate attention, it must process the entire sequence from scratch.

Let our input matrix X be the stack of all three tokens:

$$X = \begin{pmatrix} 1 & 0 \newline 0 & 1 \newline 1 & 1 \end{pmatrix}$$

Step 1: Calculate Q, K, and V for the whole sequence

We must multiply the entire input sequence by our weight matrices.

Calculate Q:

$$Q = X \times W_q = \begin{pmatrix} 1 & 0 \newline 0 & 1 \newline 1 & 1 \end{pmatrix} \times \begin{pmatrix} 1 & 0 \newline 0 & 1 \end{pmatrix} = \begin{pmatrix} 1 & 0 \newline 0 & 1 \newline 1 & 1 \end{pmatrix}$$

Calculate K:

$$K = X \times W_k = \begin{pmatrix} 1 & 0 \newline 0 & 1 \newline 1 & 1 \end{pmatrix} \times \begin{pmatrix} 1 & 1 \newline 0 & 1 \end{pmatrix} = \begin{pmatrix} 1 & 1 \newline 0 & 1 \newline 1 & 2 \end{pmatrix}$$

Calculate V:

$$V = X \times W_v = \begin{pmatrix} 1 & 0 \newline 0 & 1 \newline 1 & 1 \end{pmatrix} \times \begin{pmatrix} 2 & 0 \newline 0 & 2 \end{pmatrix} = \begin{pmatrix} 2 & 0 \newline 0 & 2 \newline 2 & 2 \end{pmatrix}$$

Notice what just happened. The first two rows of K and V are exactly the same as they were during the previous generation step. We wasted valuable GPU FLOPs recalculating them.

Step 2: Compute Attention for the new token

To predict token 4, we only care about the attention output of token 3. We take the 3rd row of Q ( $q_3$ ) and multiply it by the transposed K matrix.

(Two caveats about that "only." First, we can drop rows 1 and 2 of Q here because attention is causal — each token attends only to itself and the tokens before it, so position 3's output is unaffected by anything at positions 4 onward, and positions 1 and 2 produce outputs we already have. In the prompt phase, where all positions are computed at once, a causal mask sets the scores for future positions to −∞ to enforce this. Second, we can discard those rows only because our toy example has a single layer. In a real model, layer 2's attention at position 3 needs layer 1's output at every position — so without a cache you must recompute all N positions at all L layers, and only then throw away everything but the last row. The waste is far worse than this example makes it look. This is worked through in Appendix A.)

$$q_3 = \begin{pmatrix} 1 & 1 \end{pmatrix}$$

Unnormalized Attention Scores:

$$Scores = q_3 \times K^T = \begin{pmatrix} 1 & 1 \end{pmatrix} \times \begin{pmatrix} 1 & 0 & 1 \newline 1 & 1 & 2 \end{pmatrix}$$

$$Scores = \begin{pmatrix} (1*1 + 1*1) & (1*0 + 1*1) & (1*1 + 1*2) \end{pmatrix} = \begin{pmatrix} 2 & 1 & 3 \end{pmatrix}$$

Softmax (approximate for readability):

$$e^2 \approx 7.39$$

$$e^1 \approx 2.72$$

$$e^3 \approx 20.09$$

$$\text{sum} \approx 30.2$$

$$Probs = \begin{pmatrix} 0.24 & 0.09 & 0.67 \end{pmatrix}$$

For readability we omit the standard 1/√d scaling factor from the softmax. Real attention computes softmax(q·Kᵀ/√d), where d is the head dimension; the scaling keeps the dot products from growing large enough to saturate the softmax, but it changes nothing about the caching argument.

Multiply by V:

$$Output = Probs \times V = \begin{pmatrix} 0.24 & 0.09 & 0.67 \end{pmatrix} \times \begin{pmatrix} 2 & 0 \newline 0 & 2 \newline 2 & 2 \end{pmatrix}$$

$$Output = \begin{pmatrix} (0.48 + 0 + 1.34) & (0 + 0.18 + 1.34) \end{pmatrix} = \begin{pmatrix} 1.82 & 1.52 \end{pmatrix}$$

We arrived at our output, but computing $X \times W_k$ and $X \times W_v$ scaled linearly with the length of our sequence. If our context was 10,000 tokens long, we would have done matrix multiplication for 10,000 tokens just to generate one new one. Check out Appendix A for a deeper dive into the calculations needed when there are two layers in the model.

---

Scenario 2: Generation WITH KV Cache

Now, let's assume we are caching our states. During the previous generation step (when we predicted $x_3$ ), we saved the Keys and Values for $x_1$ and $x_2$ in GPU memory.

KV Cache in Memory:

$$K_{past} = \begin{pmatrix} 1 & 1 \newline 0 & 1 \end{pmatrix}, \quad V_{past} = \begin{pmatrix} 2 & 0 \newline 0 & 2 \end{pmatrix}$$

When token $x_3$ arrives, we do not feed the whole sequence into the weight matrices. We only pass the brand-new token.

$$x_{new} = \begin{pmatrix} 1 & 1 \end{pmatrix}$$

Step 1: Calculate Q, K, and V for ONLY the new token

This is a massive compute saving. We are now doing a simple Vector-Matrix multiplication instead of a Matrix-Matrix multiplication.

$$q_{new} = x_{new} \times W_q = \begin{pmatrix} 1 & 1 \end{pmatrix} \times \begin{pmatrix} 1 & 0 \newline 0 & 1 \end{pmatrix} = \begin{pmatrix} 1 & 1 \end{pmatrix}$$

$$k_{new} = x_{new} \times W_k = \begin{pmatrix} 1 & 1 \end{pmatrix} \times \begin{pmatrix} 1 & 1 \newline 0 & 1 \end{pmatrix} = \begin{pmatrix} 1 & 2 \end{pmatrix}$$

$$v_{new} = x_{new} \times W_v = \begin{pmatrix} 1 & 1 \end{pmatrix} \times \begin{pmatrix} 2 & 0 \newline 0 & 2 \end{pmatrix} = \begin{pmatrix} 2 & 2 \end{pmatrix}$$

Step 2: Update the Cache

We append our newly calculated $k_{new}$ and $v_{new}$ to our existing cache in memory.

$$K = \text{concat}(K_{past}, k_{new}) = \begin{pmatrix} 1 & 1 \newline 0 & 1 \newline 1 & 2 \end{pmatrix}$$

$$V = \text{concat}(V_{past}, v_{new}) = \begin{pmatrix} 2 & 0 \newline 0 & 2 \newline 2 & 2 \end{pmatrix}$$

Step 3: Compute Attention

From here, the math is exactly the same as Step 2 in the previous scenario. We multiply our $q_{new}$ by the updated $K$ cache, apply softmax, and multiply by the updated $V$ cache.

$$Scores = q_{new} \times K^T = \begin{pmatrix} 1 & 1 \end{pmatrix} \times \begin{pmatrix} 1 & 0 & 1 \newline 1 & 1 & 2 \end{pmatrix} = \begin{pmatrix} 2 & 1 & 3 \end{pmatrix}$$

The final output is identical: $\begin{pmatrix} 1.82 & 1.52 \end{pmatrix}$ .

---

The Compute vs. Memory Trade-off

By using the KV cache, we completely eliminated the redundant calculations of previous tokens.

• Without Cache: Generating the N-th token takes O(N²). The model recomputes the whole sequence from scratch — and in a multi-layer model that means running attention at every position, not just the last one: N queries against N keys.
• With Cache: Generating the N-th token takes O(N). Only the new token is projected, and its single query attends to the N−1 cached keys and values.

The Price of the KV Cache
This massive speedup comes at a steep cost: GPU Memory (VRAM). As a sequence gets longer, the $K$ and $V$ matrices grow linearly.

The prompt length is known when the request arrives, and the maximum length is bounded by the context window — but the output length is unknowable until generation terminates. A system must therefore either over-reserve for the worst case (wasting memory) or allocate incrementally (which is what PagedAttention enables).

This exact problem—managing the massive, dynamic memory footprint of the KV cache—is what paved the way for memory management innovations like PagedAttention, which stores this cache in scattered, fixed-size pages. I will discuss PagedAttention in my next article.

There are also model-side innovations to reduce demand on KV cache, such as grouped-query attention (GQA) implemented in Llama models, and multi-head latent attention (MLA) implemented by Deepseek.

---

Code

Here is the PyTorch implementation of both scenarios. Notice how the three projection matmuls stay fixed at (1, 2) × (2, 2) no matter how long the sequence gets. The attention matmuls, q_new @ K_cache.T and probs @ V_cache, still grow with the cache length. This is the cost the KV cache cannot remove.

• q_new @ K_cache.T is (1, 2) x (2, N-1)
• probs @ V_cache is (1, N-1) x (T, 2)

import torch
import torch.nn.functional as F

# --------------------------------------------------------------------------
# 1. SETUP
# --------------------------------------------------------------------------
# Weight matrices (2x2)
W_q = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
W_k = torch.tensor([[1.0, 1.0], [0.0, 1.0]])
W_v = torch.tensor([[2.0, 0.0], [0.0, 2.0]])

# Past tokens and the new token
x1 = torch.tensor([[1.0, 0.0]])
x2 = torch.tensor([[0.0, 1.0]])
x3 = torch.tensor([[1.0, 1.0]]) # The new token

print("=== SCENARIO 1: WITHOUT KV CACHE ===")
# Stack all tokens into a single input matrix X (shape: 3x2)
X = torch.cat([x1, x2, x3], dim=0)
print(f"Input X shape: {X.shape}")

# Matrix multiply the ENTIRE sequence (Wasted compute!)
Q_full = X @ W_q
K_full = X @ W_k
V_full = X @ W_v
print(f"K_full shape: {K_full.shape} (Computed from scratch)")

# We only want the attention for the latest token (q3)
q3 = Q_full[-1:] # Shape 1x2

# Attention calculation
scores_no_cache = q3 @ K_full.T
probs_no_cache = F.softmax(scores_no_cache, dim=-1)
output_no_cache = probs_no_cache @ V_full

print(f"Output without cache:\n{output_no_cache}\n")


print("=== SCENARIO 2: WITH KV CACHE ===")
# Assume we have K and V from the previous step saved in memory
# In a real system this would already be in VRAM from the previous
# step
K_past = torch.cat([x1, x2], dim=0) @ W_k 
V_past = torch.cat([x1, x2], dim=0) @ W_v
print(f"K_past shape in memory: {K_past.shape}")

# We ONLY project the new token (Massive compute savings!)
q_new = x3 @ W_q
k_new = x3 @ W_k
v_new = x3 @ W_v
print(f"k_new shape: {k_new.shape} (Only computed for 1 token)")

# Update the cache
K_cache = torch.cat([K_past, k_new], dim=0)
V_cache = torch.cat([V_past, v_new], dim=0)

# Attention calculation
scores_cache = q_new @ K_cache.T
probs_cache = F.softmax(scores_cache, dim=-1)
output_cache = probs_cache @ V_cache

print(f"Output with cache:\n{output_cache}\n")

# Verify they are mathematically identical
assert torch.allclose(output_no_cache, output_cache)
print("SUCCESS: Both methods are mathematically identical!")

Output:

=== SCENARIO 1: WITHOUT KV CACHE ===
Input X shape: torch.Size([3, 2])
K_full shape: torch.Size([3, 2]) (Computed from scratch)
Output without cache:
tensor([[1.8199, 1.5105]])

=== SCENARIO 2: WITH KV CACHE ===
K_past shape in memory: torch.Size([2, 2])
k_new shape: torch.Size([1, 2]) (Only computed for 1 token)
Output with cache:
tensor([[1.8199, 1.5105]])

SUCCESS: Both methods are mathematically identical!

---

Appendix A: Worked example with 2 layers

What happens if our model has two layers instead of just one?

The Setup:
We have 2 tokens ( $x_1, x_2$ ). We just received the new token $x_3$ .
Our model has two layers: L1 and L2. Each has its own weights (e.g., $W_q^{(L1)}$ for Layer 1, $W_q^{(L2)}$ for Layer 2).

---

Step 1: The "Shortcut" Attempt at Layer 1

We want to save compute, so we decide to only calculate the Layer 1 attention output for our newest token, $x_3$ .

1. We project the raw embeddings into $q_3^{(L1)}$ .
2. Because we have no cache, we recalculate the keys and values for all tokens from scratch: $k_1^{(L1)}, v_1^{(L1)}, k_2^{(L1)}, v_2^{(L1)}, k_3^{(L1)}, v_3^{(L1)}$ .
3. We calculate the attention for token 3 and get our result: $Output_3^{(L1)}$ .

We take our single vector, $Output_3^{(L1)}$ , and pass it up to Layer 2.

---

Step 2: Hitting the Wall at Layer 2

Layer 2 receives $Output_3^{(L1)}$ and begins its work.

1. It calculates the query for the new token:
   $$q_3^{(L2)} = Output_3^{(L1)} \times W_q^{(L2)}$$
2. Now, Layer 2 needs to perform attention. To do this, it needs the Keys and Values for the past context. Let us just focus on getting the keys for now, i.e.
   $k_1^{(L2)}, k_2^{(L2)}$
   
3. It tries to calculate
   $k_1^{(L2)}$
   

The formula for a key at Layer 2 is to multiply the output of Layer 1 by Layer 2's Key weights:

$$k_1^{(L2)} = Output_1^{(L1)} \times W_k^{(L2)}$$

CRITICAL ERROR: We don't have $Output_1^{(L1)}$ .

In Step 1, we took a shortcut and only calculated the Layer 1 output for token 3. We completely ignored the Layer 1 outputs for tokens 1 and 2. Because Layer 2's Keys and Values are derived from Layer 1's outputs, Layer 2 is now paralyzed.

---

Step 3: The Cascading Compute Penalty

To fix this, we are forced to go back to Layer 1 and do the full attention calculation for every single past token just to get their $Output^{(L1)}$ states.

• Calculate for Token 1:
  

  $$q_1^{(L1)} = x_1 \times W_q^{(L1)}$$
  $$Output_1^{(L1)} = \text{softmax}(q_1^{(L1)} \times [k_1^{(L1)}]^T) \times [v_1^{(L1)}]$$

• Calculate for Token 2:
  

  $$q_2^{(L1)} = x_2 \times W_q^{(L1)}$$
  $$Output_2^{(L1)} = \text{softmax}(q_2^{(L1)} \times [k_1^{(L1)}, k_2^{(L1)}]^T) \times [v_1^{(L1)}, v_2^{(L1)}]$$

---

Step 4: Finally Executing Layer 2

Now that Layer 1 has re-processed the entire sequence, we pass the full stack of outputs ( $Output_1^{(L1)}, Output_2^{(L1)}, Output_3^{(L1)}$ ) to Layer 2.

1. Layer 2 can finally calculate the past Keys and Values:
   $$k_1^{(L2)} = Output_1^{(L1)} \times W_k^{(L2)}$$
   $$k_2^{(L2)} = Output_2^{(L1)} \times W_k^{(L2)}$$
2. Layer 2 calculates attention for the 3rd token and outputs the final result.

The Takeaway:

Without a KV cache, you cannot just pass the newest token through the network. Layer 1 must compute attention for ALL tokens, to pass ALL outputs to Layer 2, which computes attention for ALL tokens, to pass ALL outputs to Layer 3...

If your model has 32 layers, this $O(N^2)$ recomputation penalty happens 32 times per generated word, which is inefficient at long contexts.