# KV Cache by hand > Source: > Published: 2026-08-21 14:47:59+00:00 Table of Contents The Key-Value (KV) cache is an important optimization in Large Language Model (LLM) inference. Working through the math helps with: This article is written with the assistance of AI. 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. 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. 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: Wq , Wk , Wv . Let our weight matrices be: Assume our sequence currently consists of 2 tokens, x1 and x2 , and we just generated the newest token x3 : We now need to perform the attention calculation for x3 to predict the next token, x4 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: We must multiply the entire input sequence by our weight matrices. **Calculate Q:** **Calculate K:** **Calculate V:** *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. To predict token 4, we only care about the attention output of token 3. We take the 3rd row of `Q` ( q3 ) 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.)* **Unnormalized Attention Scores:** **Softmax (approximate for readability):** *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:** We arrived at our output, but computing X×Wk and X×Wv 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. Now, let's assume we are caching our states. During the previous generation step (when we predicted x3 ), we saved the Keys and Values for x1 and x2 in GPU memory. **KV Cache in Memory:** When token x3 arrives, we **do not** feed the whole sequence into the weight matrices. We only pass the brand-new token. This is a massive compute saving. We are now doing a simple Vector-Matrix multiplication instead of a Matrix-Matrix multiplication. We append our newly calculated knew and vnew to our existing cache in memory. From here, the math is exactly the same as Step 2 in the previous scenario. We multiply our qnew by the updated K cache, apply softmax, and multiply by the updated V cache. The final output is identical: (1.821.52) . By using the KV cache, we completely eliminated the redundant calculations of previous tokens. **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. 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) ``` python 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! ``` What happens if our model has *two* layers instead of just one? **The Setup:** We have 2 tokens ( x1,x2 ). We just received the new token x3 . Our model has two layers: **L1** and **L2**. Each has its own weights (e.g., Wq(L1) for Layer 1, Wq(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, x3 . We take our single vector, Output3(L1) , and pass it up to Layer 2. **Step 2: Hitting the Wall at Layer 2** Layer 2 receives Output3(L1) and begins its work. The formula for a key at Layer 2 is to multiply the *output of Layer 1* by Layer 2's Key weights: **CRITICAL ERROR:** We don't have Output1(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:** **Calculate for Token 2:** **Step 4: Finally Executing Layer 2** Now that Layer 1 has re-processed the entire sequence, we pass the full stack of outputs ( Output1(L1),Output2(L1),Output3(L1) ) to Layer 2. **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(N2) recomputation penalty happens 32 times per generated word, which is inefficient at long contexts.