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 trade-offs: 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 . We have three weight matrices for our attention head: , , . Let our weight matrices be: Assume our sequence currently consists of 2 tokens, and , and we just generated the newest token : - (The new token) We now need to perform the attention calculation for to predict the next token, 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: Step 1: Calculate Q, K, and V for the whole sequence 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. 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 ( ) 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 and 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 ), we saved the Keys and Values for and in GPU memory. KV Cache in Memory: When token arrives, we do not feed the whole sequence into the weight matrices. We only pass the brand-new token. 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. Step 2: Update the Cache We append our newly calculated and to our existing cache in memory. Step 3: Compute Attention From here, the math is exactly the same as Step 2 in the previous scenario. We multiply our by the updated cache, apply softmax, and multiply by the updated cache. The final output is identical: . 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 and 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: