{"slug": "kv-cache-by-hand", "title": "KV Cache by hand", "summary": "An engineer's blog post explains the Key-Value (KV) cache optimization in large language model inference, walking through the math of attention calculation to show how caching keys and values avoids redundant computation. The post demonstrates that without a KV cache, the model recalculates keys and values for all previous tokens at every generation step, wasting GPU FLOPs, and notes that queries are not cached because they are used only once due to causal masking.", "body_md": "Table of Contents\n\nThe Key-Value (KV) cache is an important optimization in Large Language Model (LLM) inference. Working through the math helps with:\n\nThis article is written with the assistance of AI.\n\nLLMs generate text **autoregressively**, meaning they predict the sequence one token at a time. To predict token `N`\n\n, the model needs to look at all previous tokens from `1`\n\nto `N-1`\n\n.\n\nIn the Transformer attention mechanism, every token is projected into a **Query (Q)**, **Key (K)**, and **Value (V)** vector.\n\nWhen 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**.\n\nIf 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.\n\nQuick question: why do we not cache `Q`\n\n? 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.\n\nWe will walk through the attention calculation for the decode step that predicts the 4th token.\n\nAssume our embedding dimension is d=2 . We have three weight matrices for our attention head: Wq , Wk , Wv .\n\nLet our weight matrices be:\n\nAssume our sequence currently consists of 2 tokens, x1 and x2 , and we just generated the newest token x3 :\n\nWe now need to perform the attention calculation for x3 to predict the next token, x4\n\nWithout 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.\n\nLet our input matrix `X`\n\nbe the stack of all three tokens:\n\nWe must multiply the entire input sequence by our weight matrices.\n\n**Calculate Q:**\n\n**Calculate K:**\n\n**Calculate V:**\n\n*Notice what just happened.* The first two rows of `K`\n\nand `V`\n\nare exactly the same as they were during the previous generation step. We wasted valuable GPU FLOPs recalculating them.\n\nTo predict token 4, we only care about the attention output of token 3. We take the 3rd row of `Q`\n\n(\nq3\n) and multiply it by the transposed `K`\n\nmatrix.\n\n*(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.)*\n\n**Unnormalized Attention Scores:**\n\n**Softmax (approximate for readability):**\n\n*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.*\n\n**Multiply by V:**\n\nWe 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.\n\nNow, 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.\n\n**KV Cache in Memory:**\n\nWhen token\nx3\narrives, we **do not** feed the whole sequence into the weight matrices. We only pass the brand-new token.\n\nThis is a massive compute saving. We are now doing a simple Vector-Matrix multiplication instead of a Matrix-Matrix multiplication.\n\nWe append our newly calculated knew and vnew to our existing cache in memory.\n\nFrom 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.\n\nThe final output is identical: (1.821.52) .\n\nBy using the KV cache, we completely eliminated the redundant calculations of previous tokens.\n\n**The Price of the KV Cache**\n\nThis massive speedup comes at a steep cost: GPU Memory (VRAM). As a sequence gets longer, the\nK\nand\nV\nmatrices grow linearly.\n\nThe 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).\n\nThis 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.\n\nThere 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.\n\nHere 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`\n\nand `probs @ V_cache`\n\n, still grow with the cache length. This is the cost the KV cache cannot remove.\n\n`q_new @ K_cache.T`\n\nis (1, 2) x (2, N-1)`probs @ V_cache`\n\nis (1, N-1) x (T, 2)\n\n``` python\nimport torch\nimport torch.nn.functional as F\n\n# --------------------------------------------------------------------------\n# 1. SETUP\n# --------------------------------------------------------------------------\n# Weight matrices (2x2)\nW_q = torch.tensor([[1.0, 0.0], [0.0, 1.0]])\nW_k = torch.tensor([[1.0, 1.0], [0.0, 1.0]])\nW_v = torch.tensor([[2.0, 0.0], [0.0, 2.0]])\n\n# Past tokens and the new token\nx1 = torch.tensor([[1.0, 0.0]])\nx2 = torch.tensor([[0.0, 1.0]])\nx3 = torch.tensor([[1.0, 1.0]]) # The new token\n\nprint(\"=== SCENARIO 1: WITHOUT KV CACHE ===\")\n# Stack all tokens into a single input matrix X (shape: 3x2)\nX = torch.cat([x1, x2, x3], dim=0)\nprint(f\"Input X shape: {X.shape}\")\n\n# Matrix multiply the ENTIRE sequence (Wasted compute!)\nQ_full = X @ W_q\nK_full = X @ W_k\nV_full = X @ W_v\nprint(f\"K_full shape: {K_full.shape} (Computed from scratch)\")\n\n# We only want the attention for the latest token (q3)\nq3 = Q_full[-1:] # Shape 1x2\n\n# Attention calculation\nscores_no_cache = q3 @ K_full.T\nprobs_no_cache = F.softmax(scores_no_cache, dim=-1)\noutput_no_cache = probs_no_cache @ V_full\n\nprint(f\"Output without cache:\\n{output_no_cache}\\n\")\n\nprint(\"=== SCENARIO 2: WITH KV CACHE ===\")\n# Assume we have K and V from the previous step saved in memory\n# In a real system this would already be in VRAM from the previous\n# step\nK_past = torch.cat([x1, x2], dim=0) @ W_k \nV_past = torch.cat([x1, x2], dim=0) @ W_v\nprint(f\"K_past shape in memory: {K_past.shape}\")\n\n# We ONLY project the new token (Massive compute savings!)\nq_new = x3 @ W_q\nk_new = x3 @ W_k\nv_new = x3 @ W_v\nprint(f\"k_new shape: {k_new.shape} (Only computed for 1 token)\")\n\n# Update the cache\nK_cache = torch.cat([K_past, k_new], dim=0)\nV_cache = torch.cat([V_past, v_new], dim=0)\n\n# Attention calculation\nscores_cache = q_new @ K_cache.T\nprobs_cache = F.softmax(scores_cache, dim=-1)\noutput_cache = probs_cache @ V_cache\n\nprint(f\"Output with cache:\\n{output_cache}\\n\")\n\n# Verify they are mathematically identical\nassert torch.allclose(output_no_cache, output_cache)\nprint(\"SUCCESS: Both methods are mathematically identical!\")\n```\n\nOutput:\n\n```\n=== SCENARIO 1: WITHOUT KV CACHE ===\nInput X shape: torch.Size([3, 2])\nK_full shape: torch.Size([3, 2]) (Computed from scratch)\nOutput without cache:\ntensor([[1.8199, 1.5105]])\n\n=== SCENARIO 2: WITH KV CACHE ===\nK_past shape in memory: torch.Size([2, 2])\nk_new shape: torch.Size([1, 2]) (Only computed for 1 token)\nOutput with cache:\ntensor([[1.8199, 1.5105]])\n\nSUCCESS: Both methods are mathematically identical!\n```\n\nWhat happens if our model has *two* layers instead of just one?\n\n**The Setup:**\n\nWe have 2 tokens (\nx1,x2\n). We just received the new token\nx3\n.\n\nOur model has two layers: **L1** and **L2**. Each has its own weights (e.g.,\nWq(L1)\nfor Layer 1,\nWq(L2)\nfor Layer 2).\n\n**Step 1: The \"Shortcut\" Attempt at Layer 1**\n\nWe want to save compute, so we decide to only calculate the Layer 1 attention output for our newest token, x3 .\n\nWe take our single vector, Output3(L1) , and pass it up to Layer 2.\n\n**Step 2: Hitting the Wall at Layer 2**\n\nLayer 2 receives Output3(L1) and begins its work.\n\nThe formula for a key at Layer 2 is to multiply the *output of Layer 1* by Layer 2's Key weights:\n\n**CRITICAL ERROR:** We don't have\nOutput1(L1)\n.\n\nIn 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.\n\n**Step 3: The Cascading Compute Penalty**\n\nTo 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\nOutput(L1)\nstates.\n\n**Calculate for Token 1:**\n\n**Calculate for Token 2:**\n\n**Step 4: Finally Executing Layer 2**\n\nNow 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.\n\n**The Takeaway:**\n\nWithout 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...**\n\nIf your model has 32 layers, this O(N2) recomputation penalty happens 32 times per generated word, which is inefficient at long contexts.", "url": "https://wpnews.pro/news/kv-cache-by-hand", "canonical_source": "https://dev.to/lewis_won/kv-cache-by-hand-26i4", "published_at": "2026-08-21 14:47:59+00:00", "updated_at": "2026-08-21 15:15:35.225628+00:00", "lang": "en", "topics": ["large-language-models", "artificial-intelligence", "machine-learning"], "entities": [], "alternates": {"html": "https://wpnews.pro/news/kv-cache-by-hand", "markdown": "https://wpnews.pro/news/kv-cache-by-hand.md", "text": "https://wpnews.pro/news/kv-cache-by-hand.txt", "jsonld": "https://wpnews.pro/news/kv-cache-by-hand.jsonld"}}