[{"content":"When we judge how capable a model is, the first number we look at is its parameter count. Scaling laws back that instinct up: a model\u0026rsquo;s loss falls predictably as you scale up model size, training data and compute together. The catch is that a bigger model also costs more compute, in training and in inference. Mixture of Experts, or MoE, breaks that link. It lets a model\u0026rsquo;s parameter count grow tens of times over while the compute spent on each token barely moves. The Switch Transformers paper puts it this way: \u0026ldquo;a sparsely-activated model \u0026ndash; with outrageous numbers of parameters \u0026ndash; but a constant computational cost.\u0026rdquo; Today almost every flagship open model, DeepSeek, Qwen and Kimi among them, is built on it. This post walks through how it works, one step at a time.\nDense models: every parameter in the FFN does work for every token I\u0026rsquo;ve covered the Transformer decoder in earlier posts, so here\u0026rsquo;s a quick recap. A decoder layer has two main blocks. The input first goes through multi-head attention, whose job is to move information between tokens. It then goes through the FFN, the feed-forward network. Once each token has picked up its context, the FFN mixes and transforms that information to produce the layer\u0026rsquo;s output. Attention is the only block in a Transformer that looks across tokens. The FFN treats every token independently, so all tokens can go through it in parallel.\nInside the FFN, each token\u0026rsquo;s \\(d\\)-dimensional vector is first projected up to a wider intermediate size, \\(d_{\\mathit{ff}}\\). A nonlinear activation function is applied in that wider space, and the result is projected back down to \\(d\\) dimensions.\nIn matrix form, say the prompt has \\(n\\) tokens and each token is a \\(d\\)-dimensional vector. The input \\(X\\) is an \\(n \\times d\\) matrix. Multiplying it by the up-projection \\(W_1\\) gives an \\(n \\times d_{\\mathit{ff}}\\) intermediate result. After the activation, multiplying by the down-projection \\(W_2\\) brings the output \\(Y\\) back to \\(n \\times d\\). The shape that goes in is the shape that comes out.\n(The sizes in these diagrams are illustrative. In real models \\(d\\) is in the thousands.)\nAn FFN built this way is what makes a dense model: every token goes through the same complete network, and every neuron and every parameter in the FFN takes part in the computation. So the bigger the model, the more compute each token needs. In a dense model, adding parameters means adding compute.\nSparse models: four experts, and each token uses one A sparse model does it differently. It makes many copies of the FFN, and each copy is called an expert. When the model runs, each token goes through only one of them.\nTake the prompt \u0026ldquo;What is 1 + 1 ?\u0026rdquo; and four experts that handle punctuation, verbs, conjunctions and numbers. When the sparse model sees that the current token is most closely related to numbers, it sends the token to the numbers expert. The other three experts sit idle and do no computation at all, and the output of the whole sparse layer is the numbers expert\u0026rsquo;s output.\n(The expert labels are for illustration. Real experts don\u0026rsquo;t divide the work this neatly: the Mixtral paper found no clear specialization by topic.)\nAt its core, then, a sparse model picks a compute path based on its input. Each token activates only part of the model\u0026rsquo;s parameters, which is how the total parameter count can grow without adding compute per token.\nCompare that with the multi-head attention block. Whatever the input looks like, every head computes its attention. The results of all the heads are concatenated into one large matrix, and a linear projection turns that into the block\u0026rsquo;s output. So of the two blocks in a decoder layer, attention puts every one of its neurons to work, while the sparse layer uses only some of them.\nRouting: a router scores the experts, and the top one runs That raises the key question: who decides which expert a token goes to? MoE adds a component for this called the router. It is a small layer followed by a softmax. Its input is the token\u0026rsquo;s vector, and its output is \\(E\\) numbers, one probability per expert. The higher the probability, the better suited the router thinks that expert is for this token.\nSay a token comes in and the router\u0026rsquo;s softmax gives the four experts probabilities of 0.45, 0.19, 0.05 and 0.31. The first expert scores highest, so the token goes to expert 1, and that expert\u0026rsquo;s output is multiplied by its probability, 0.45, to give the layer\u0026rsquo;s output.\nThat multiplication matters. Picking the expert means taking the maximum, and taking a maximum can\u0026rsquo;t be differentiated, so no gradient flows back through the choice itself. Multiplying by 0.45 puts the probability into the formula for the output. The gradient can then flow back through that multiplication into the router\u0026rsquo;s parameters and tell it whether this choice was a good one.\nIn matrix form, if the input is an \\(n \\times d\\) matrix, the router produces an \\(n \\times E\\) matrix, and softmax is applied to each row. Each row is one token\u0026rsquo;s probability distribution over the \\(E\\) experts.\nThe router box in the figure is labeled FFN, but in Mixtral and Switch Transformers the router is in fact a single linear layer: one \\(d \\times E\\) matrix, \\(W_g\\) above, followed by the softmax.\nThe forward pass: dispatch, then combine Now let\u0026rsquo;s follow one complete forward pass through an MoE layer. The input is six tokens, a 6×3 matrix, and each token gets a color tag so we can track where it goes.\nNext the router scores them, producing a 6×4 matrix. Each row holds one token\u0026rsquo;s probabilities for the four experts, and the highest cell in each row is highlighted in that expert\u0026rsquo;s color. t1\u0026rsquo;s highest score is 0.62, for expert 2. t2\u0026rsquo;s is 0.70, for expert 1.\nThen comes the important step, dispatch. Following the experts just chosen, the six tokens are regrouped into four piles. t2 and t5 both went to expert 1, so they are pulled out into a new matrix of their own. In the same way, t1, t4 and t6 are rearranged into a separate matrix that becomes expert 2\u0026rsquo;s input, and t3 goes to expert 4. Notice that expert 3 got no tokens at all, so its input is empty.\nEach expert then runs its FFN on its own input matrix and multiplies every row of the result by the probability that token was routed with. The last step, combine, puts all the weighted outputs back in the tokens\u0026rsquo; original order, so each position again holds the result for the token that started there.\nThe dispatch and combine steps look a lot like MapReduce from distributed computing. Much as attention\u0026rsquo;s query, key and value borrow from information retrieval, ideas from classic algorithms turn up all over the LLM stack.\nOne thing to keep in mind from this example: how many tokens each expert receives is decided entirely by the data. The counts are naturally unequal and can even be zero. This is one of the central problems in MoE models, load balancing, and we\u0026rsquo;ll come back to it shortly.\nTop-K: several experts, added up by weight Routing each token to a single expert has a weakness. A model that can send its input through only one expert is limited in what it can express. When several experts score about the same, a tiny difference in routing scores can send a token down a completely different path. That easily makes training unstable and feeds the load-balancing problem we just saw. So real models pick more than one expert, usually the \\(K\\) with the highest scores. DeepSeek-V4-Pro has 384 experts and activates 6 of them for each token. Kimi-K3 picks 16 out of 896. Mixtral, by comparison, picks 2 out of 8. The number of experts has grown much faster than \\(K\\), so each token touches a smaller and smaller share of the model.\nTop-K routing changes little in the computation. The one new step comes after the top \\(K\\) experts are chosen: their weights are renormalized. In the left example, the two highest probabilities are 0.45 and 0.31. Before the outputs are combined, 0.45 is divided by their sum, 0.76, to get 0.59, and 0.31 divided by 0.76 gives 0.41. These renormalized weights are what the two experts\u0026rsquo; outputs are multiplied by and added up with. The right column runs the same steps for the token \u0026ldquo;is\u0026rdquo;, which lands on two different experts.\nLoad balancing: how collapse happens, and how it\u0026rsquo;s kept in check We saw in the forward pass that different experts get different amounts of work. An expert that isn\u0026rsquo;t chosen gets no gradient at all, and that sets up a feedback loop. At the start of training the router is randomly initialized. Some expert happens to receive a few more tokens, so it gets a little more gradient and its parameters improve a little more. At the next step the router is more inclined to pick it, so it receives a few more tokens still. The router gradually funnels tokens to a handful of experts, and the rest, rarely chosen, barely get trained. Shazeer et al. described this back in 2017: the gating network \u0026ldquo;tends to converge to a state where it always produces large weights for the same few experts.\u0026rdquo; We want every expert to receive roughly the same number of tokens, so that all of them get trained. This is load balancing in MoE: keeping the network from leaning on a few experts.\nThe fix is an auxiliary loss built from two quantities, \\(P_i\\) and \\(f_i\\). \\(P_i\\) is the average probability the router gives expert \\(i\\) across the tokens in a batch.\n$$ P_i=\\frac{1}{n}\\sum_{x\\in B}p_i(x) $$ On the left of the table below is an unbalanced router. \\(P_1 = 0.58\\) is the average of the first probability across the four tokens\u0026rsquo; distributions, so \\(P_i\\) shows that the router strongly favors expert 1 and gives the others little. On the right, with balanced routing, each expert\u0026rsquo;s average probability hovers around a quarter, 0.25.\n\\(f_i\\) is the fraction of tokens in the batch that the router actually sends to expert \\(i\\), which is how much traffic that expert gets in practice.\n$$ f_i=\\frac{1}{n}\\sum_{x\\in B}\\mathbf{1}\\{\\arg\\max p(x)=i\\} $$ On the left, expert 1\u0026rsquo;s average probability is only 0.58, but it is always the largest one, so the router picks it every time. Its traffic is 100%, \\(f_1 = 1\\): one expert takes everything. On the right, \\(f\\) is four 0.25s, since each expert is picked once.\nThe auxiliary loss is \\(\\alpha\\) times the number of experts \\(E\\) times the sum, over all experts, of \\(f_i\\) times \\(P_i\\).\n$$ \\text{Auxiliary Loss}=\\alpha\\cdot E\\sum_{i=1}^{E} f_i\\,P_i $$ On the left, only the first expert contributes: \\(4 \\times 1 \\times 0.58 \\approx 2.33\\) (\\(P_1\\) is 0.5825 before rounding). The right is easy to see too. Every \\(f_i\\) is 0.25, so \\(E \\cdot f_i = 1\\), and the \\(P_i\\) add up to 1 because every token\u0026rsquo;s probabilities do. When the load is perfectly balanced, the loss reaches its minimum of 1. The table shows the part after \\(\\alpha\\); Switch Transformers sets \\(\\alpha\\) to \\(10^{-2}\\), so the terms added to the total loss are about 0.0233 and 0.01. The product matters because \\(f_i\\) is a count and can\u0026rsquo;t be differentiated, while \\(P_i\\) comes straight from the softmax and can. Adding this term to the total loss gradually evens out how many tokens each expert gets during training.\nSummary: more experts, without more of them per token Back to the question we started with. A dense model can only add parameters by making its FFN wider, and every bit of extra width is extra compute for every token. MoE instead makes more copies of the FFN. The number of experts can grow very large, while the number each token activates doesn\u0026rsquo;t have to grow with it. That is how parameter count and compute come apart: the parameters grow tens of times over while the compute per token barely moves.\nReferences Noam Shazeer et al. (@NoamShazeer). “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” ICLR 2017. William Fedus et al. (@LiamFedus). “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” JMLR 2022. Maarten Grootendorst (@MaartenGr). “A Visual Guide to Mixture of Experts (MoE).” 2024. ","permalink":"https://zhangchitc.github.io/posts/mixture-of-experts/","summary":"Mixture of Experts lets a model\u0026rsquo;s parameter count grow tens of times over while the compute per token barely moves. We build it up step by step: copying the FFN into experts, the router, dispatch and combine, top-K routing, and the auxiliary loss that keeps experts balanced.","title":"How Mixture of Experts Scales Model Parameters"},{"content":"The FlashAttention paper has a striking benchmark. Training GPT-2 small on eight A100s took 9.5 days with the Hugging Face implementation and 4.7 days with NVIDIA\u0026rsquo;s Megatron-LM. With FlashAttention it took 2.7 days, and the trained model was exactly as good: all three reached the same perplexity of 18.2. In another benchmark, reads and writes to GPU memory dropped from 40.3 GB to 4.4 GB, about 9× less. (These are the paper\u0026rsquo;s 2022 comparisons. Both baselines have since adopted fused attention kernels, so the gap would be smaller today.)\nFlashAttention-style kernels are now a standard part of the modern LLM stack. PyTorch 2, Hugging Face and vLLM all ship them as attention backends. So how does it pull this off? This post builds it up one step at a time.\nGPU memory comes in tiers Start with the fact that GPU memory is layered. At the top is SRAM, which sits on each streaming multiprocessor (SM), the GPU\u0026rsquo;s compute units. On an A100 its bandwidth is about 19 TB/s, but all the SRAM across the whole chip adds up to only about 20 MB. In the middle is HBM, high-bandwidth memory, the kind SK hynix and Micron make. An A100 has 40 GB of it, at 1.5 TB/s. At the bottom is the CPU\u0026rsquo;s DRAM, which can reach terabytes but delivers only 12.8 GB/s, in a different league entirely from the two above.\nThink of HBM as the warehouse. Model weights, the KV cache and every layer\u0026rsquo;s activations all live there. SRAM is the workbench: it holds only the small piece of data being worked on right now and is freed as soon as that work is done. The compute units can\u0026rsquo;t operate on numbers in HBM directly. Any number that takes part in a computation has to be copied from HBM into on-chip SRAM first. So the way to make a GPU fast is to finish as much work as possible in SRAM and avoid reading and writing HBM in the middle of a computation.\nAttention: the matrices don\u0026rsquo;t fit, the vectors do With that in mind, look at how much data attention handles. I covered the formula in my post on the KV cache, so here\u0026rsquo;s the short version. Say the prompt has n tokens, and each token\u0026rsquo;s Q, K and V vectors have d dimensions. Then Q, K and V are each an \\(n \\times d\\) matrix. Attention multiplies Q by K transposed, divides by the square root of d, applies softmax to each row to get an \\(n \\times n\\) table, and finally multiplies by V to get the output O, another \\(n \\times d\\) matrix.\nThe key point is that n is the context length, which can reach a hundred thousand tokens or even a million, while d is the dimension of one head, a constant like 32, 64 or 128 in most models. Both the \\(n \\times d\\) matrices and the \\(n \\times n\\) table are far bigger than SRAM, so they can only live in HBM. (At n = 128K and d = 128 in FP16, a single \\(n \\times d\\) matrix is 32 MB and the \\(n \\times n\\) table is 32 GB.) Only data that doesn\u0026rsquo;t grow with n, like a single token\u0026rsquo;s q, k, v and o vectors, fits into SRAM in one go. This is the most fundamental constraint on running a Transformer on a GPU.\nThree steps, and every row is independent So where exactly does attention lose time? Let\u0026rsquo;s break the computation down as far as it goes and follow the first row through each step.\nStep one is S = QK transposed. The first row is q1 dotted with k1, k2, all the way to kN: n numbers that are the first token\u0026rsquo;s attention scores against every token. At this point they\u0026rsquo;re raw scores, not yet normalized. (The diagrams leave out the division by the square root of d to keep them clean. It doesn\u0026rsquo;t change anything here.)\nStep two is P = softmax(S), applied row by row. Softmax turns a row of raw scores into a probability distribution that sums to 1: take the exponential of each score and divide by the sum of the exponentials across the row. So p11, p12, through p1N add up to exactly 1.\nStep three is O = PV. For the first token, p1 times V expands to v1 times p11, plus v2 times p12, all the way to vN times p1N. In other words, every token\u0026rsquo;s value vector is weighted by those attention probabilities and summed into one output vector, o1.\nRemember that nothing whose size depends on n fits in SRAM. So in step one, we can only bring in one q and one k from HBM at a time, compute their dot product, and write the result straight back to HBM, because SRAM has no room for a whole row of n results. Steps two and three work the same way: each one reads the previous step\u0026rsquo;s results from HBM, computes, and writes its own results back to HBM.\nKernel fusion: three round trips into one The left half of this diagram shows the standard attention computation on a GPU. Each step reads part of the data from HBM, computes, and writes the result back to HBM, and this happens three times. Since HBM is so much slower than SRAM, a large share of the time goes into writing intermediate results to HBM and reading them back.\nThat suggests an obvious idea: what if we never stored those intermediate results in HBM at all, and instead chained every step together and wrote to HBM only once at the end? In terms of the matrix math, we\u0026rsquo;d read in q1, k1 and v1, compute the dot product q1·k1, and instead of writing it back, go straight to the next step: somehow turn it into the softmax probability p11, multiply p11 by v1, and add the result into o1. Then loop over k2 and v2, and so on up to N, and write only the final o vector back to HBM.\nThat\u0026rsquo;s the intuition behind kernel fusion, shown in the right half of the diagram.\nA fused kernel, and an overflow problem There\u0026rsquo;s a fairly direct way to write this. In the code below, q is one row\u0026rsquo;s query vector, and K and V are the full matrices. The variable d is the softmax denominator, starting at 0 (here d means the denominator, not the head dimension), and o is the output vector, starting at all zeros. The loop over k and v stands for reading one k and one v at a time from HBM into SRAM. x is q dotted with k, times scale, which is the \\(1/\\sqrt{d}\\) scaling factor from the attention formula. Then exp(x) times v is added straight into o.\nWhy does that work? Softmax is exp(x) divided by the sum of all the exponentials. When we reach a given x, we don\u0026rsquo;t know that sum yet, so we set it aside and use only the numerator as the weight on v, adding it into o. Once every x has been processed, d holds the full denominator, and we return o divided by d in one step. Since the denominator is the same number for the whole row, dividing before or after the sum gives the same result. We\u0026rsquo;ve only tracked the numerator\u0026rsquo;s and the denominator\u0026rsquo;s contributions to o separately, without changing o.\nThis is the fused kernel we were after. K and V are read from HBM only once, the intermediate values x, o and d don\u0026rsquo;t grow with n so they all stay in SRAM, and only the final o is written to HBM. But it has a fatal flaw: exp(x) overflows. x is a dot product of two vectors, so its size is unbounded, while floating-point numbers have a limited range. In FP16, whose largest value is 65,504, \\(e^{12}\\) is already out of range. So real softmax implementations subtract a baseline m before taking the exponential, a trick called safe softmax. The two are clearly equivalent: it amounts to multiplying the numerator and the denominator by \\(e^{-m}\\). The usual choice is m = the maximum of all the x values. Then every \\(x_i - m\\) is at most 0, every \\(e^{x_i-m}\\) falls between 0 and 1, and the exponential can never overflow.\n$$ \\begin{aligned} \\mathrm{softmax}(x_i) \u0026= \\dfrac{e^{x_i}}{\\sum_{j} e^{x_j}}\\\\[10pt] \\text{safe softmax}(x_i) \u0026= \\dfrac{\\textcolor{#0a7d61}{e^{-m}}\\cdot e^{x_i}}{\\textcolor{#0a7d61}{e^{-m}}\\cdot\\sum_{j} e^{x_j}} = \\dfrac{e^{x_i\\textcolor{#0a7d61}{-m}}}{\\sum_{j} e^{x_j\\textcolor{#0a7d61}{-m}}} \\end{aligned} $$ Online softmax: three passes become two Here\u0026rsquo;s a straightforward implementation of safe softmax. It takes three passes. The first finds the row\u0026rsquo;s maximum m. The second subtracts m from each value, takes the exponential and adds it into the denominator d. Only the third divides each term by d to produce the softmax. Here \\(d_i\\) is the sum of the first i terms, and every term subtracts the same baseline, \\(m_N\\), the maximum over the entire row. Since \\(m_N\\) is known only after scanning the whole row, the second pass can\u0026rsquo;t start until the first one finishes.\n$$ d_i = \\sum_{j\\le i} e^{\\,x_j-m_N} $$ Is there a way to avoid all these passes? That\u0026rsquo;s what online softmax does. If the global maximum \\(m_N\\) is what forces the extra pass, get rid of it. Redefine the denominator: let \\(d'_i\\) also be the sum of the first i terms, but have each term subtract \\(m_i\\), the maximum of the first i terms, instead of the global maximum. With that change, \\(d'_i\\) depends only on the first i terms, so it can be computed while scanning. And the \\(d_N\\) we ultimately need is the same thing as \\(d'_N\\), because at i = N the two definitions coincide.\n$$ d'_i = \\sum_{j\\le i} e^{\\,x_j-\\textcolor{#0a7d61}{m_i}} $$ The key is that with the new definition, \\(d'_i\\) and \\(d'_{i-1}\\) are linked by a recurrence: \\(d'_i\\) equals \\(d'_{i-1}\\) times a correction factor \\(e^{m_{i-1}-m_i}\\), plus \\(e^{x_i-m_i}\\). Every term in \\(d'_{i-1}\\) was computed against the baseline \\(m_{i-1}\\). Now the baseline changes to \\(m_i\\), so we multiply the whole old sum by the correction factor. When the exponents add, the \\(m_{i-1}\\) in the middle cancels, leaving \\(e^{x_j-m_i}\\), so every term now uses the new baseline. Then we add the i-th term against the new baseline, and that\u0026rsquo;s the new \\(d'_i\\).\n$$ \\begin{aligned} d'_i \u0026= d'_{i-1}\\,\\textcolor{#0a7d61}{e^{\\,m_{i-1}-m_i}} + e^{\\,x_i-m_i}\\\\[6pt] \u0026= \\Big(\\sum\\nolimits_{j\\le i-1} e^{\\,x_j-m_{i-1}}\\Big)\\textcolor{#0a7d61}{e^{\\,m_{i-1}-m_i}} + e^{\\,x_i-m_i}\\\\[4pt] \u0026= \\sum\\nolimits_{j\\le i-1} e^{\\,x_j-m_i} + e^{\\,x_i-m_i} \\;=\\; \\sum\\nolimits_{j\\le i} e^{\\,x_j-m_i} \\end{aligned} $$ So the online softmax code needs only a small change. It introduces m_old for the previous maximum, updates m, and then follows the recurrence: d becomes d times exp(m_old − m), plus exp(xi − m). The maximum m and the denominator d are now computed together in a single pass. Three passes have become two.\nFlashAttention: the same trick, applied once more Now let\u0026rsquo;s apply online softmax to attention. That gives the code below, which needs two passes. The first loops over every k vector and uses the recurrence to compute the maximum m and the denominator d. The second goes over K and V again, computes exp(q·k·scale − m) for each term as safe softmax defines it, weights v by it, and adds the result into o.\nThis version can\u0026rsquo;t overflow, it\u0026rsquo;s code that actually runs, and it needs only two passes compared with three for the naive approach. But can we go further? Look at how \\(o_i\\) is defined under safe softmax: every term uses the global maximum \\(m_N\\) and the global denominator \\(d_N\\). Doesn\u0026rsquo;t that look a lot like \\(d_i\\) in the previous section? Can we use the same trick, redefining \\(o_i\\) so its terms no longer depend on the global maximum and the global denominator?\n$$ \\boldsymbol{o}_i = \\sum_{j\\le i}\\frac{e^{\\,x_j-m_N}}{d_N}\\,\\boldsymbol{v}_j $$ We can. Define \\(o'_i\\) by replacing \\(m_N\\) and \\(d_N\\) with \\(m_i\\), the maximum of the first i terms, and \\(d'_i\\), the denominator of the first i terms. Now \\(o'_i\\) depends only on the first i terms, and the \\(o_N\\) we ultimately need is the same thing as \\(o'_N\\).\n$$ \\boldsymbol{o}'_i = \\sum_{j\\le i}\\frac{e^{\\,x_j-\\textcolor{#0a7d61}{m_i}}}{\\textcolor{#0a7d61}{d'_i}}\\,\\boldsymbol{v}_j $$ Just as before, \\(o'_i\\) and \\(o'_{i-1}\\) are linked by a recurrence: \\(o'_i\\) equals \\(o'_{i-1}\\) times a correction factor, plus the i-th term\u0026rsquo;s new contribution. The correction factor is \\(d'_{i-1}\\) times \\(e^{m_{i-1}-m_i}\\), divided by \\(d'_i\\). You can read it in two parts. The exponential in the numerator does the same job as before, switching the old baseline to the new one. The ratio of denominators swaps the old denominator for the new one. Multiplied together, they make this single factor.\n$$ \\begin{aligned} \\boldsymbol{o}'_i \u0026= \\boldsymbol{o}'_{i-1}\\,\\textcolor{#0a7d61}{\\frac{d'_{i-1}\\,e^{\\,m_{i-1}-m_i}}{d'_i}} + \\frac{e^{\\,x_i-m_i}}{d'_i}\\,\\boldsymbol{v}_i\\\\[8pt] \u0026= \\Big(\\sum\\nolimits_{j\\le i-1}\\frac{e^{\\,x_j-m_{i-1}}}{d'_{i-1}}\\boldsymbol{v}_j\\Big)\\textcolor{#0a7d61}{\\frac{d'_{i-1}\\,e^{\\,m_{i-1}-m_i}}{d'_i}} + \\frac{e^{\\,x_i-m_i}}{d'_i}\\boldsymbol{v}_i\\\\[6pt] \u0026= \\sum\\nolimits_{j\\le i}\\frac{e^{\\,x_j-m_i}}{d'_i}\\boldsymbol{v}_j \\end{aligned} $$ From this recurrence we can write the code below. Compared with the two-pass version, it adds one line, d_old = d, to keep the denominator from before the update. The rest follows the recurrence directly: compute the new m and the new d, assemble the correction factor rescale, then set o to o times rescale, plus exp(x − m) divided by d, times v. Once q, k and v are read in from HBM, all the intermediate values x, m, d and o stay in SRAM. Nothing intermediate is read from or written to HBM until every k and v has been scanned and the final o vector is written back. And that\u0026rsquo;s FlashAttention.\nTiling: moving blocks, not vectors To keep the explanation simple, we\u0026rsquo;ve been reading one q vector and one k, v pair at a time, computing one dot product at a time, and producing one output vector. A real GPU doesn\u0026rsquo;t work that way, because it would leave both SRAM and the Tensor Cores mostly idle. In practice, FlashAttention splits Q, K and V into blocks of many vectors, called tiles, and sends a whole tile to the compute units at once. Where before each step produced one output vector o, now each step produces a batch of output vectors, a block of O. And since different rows of Q are completely independent, all of those outputs can be computed in parallel.\nThe code changes very little. The maximum m and the denominator d were scalars; now they\u0026rsquo;re vectors, holding one maximum and one denominator per row. The output o was a vector; now it\u0026rsquo;s a block of rows, a matrix. The rest is the same algorithm as before, just running on tile_size vectors in parallel instead of one. If the previous version made sense, this one needs no new ideas.\n(The loop order here follows FlashAttention-2, with Q in the outer loop and K and V in the inner loop. The original FlashAttention paper put K and V on the outside.)\nWhat FlashAttention actually saves To sum up, FlashAttention comes down to three techniques. The first is kernel fusion, which collapses the three round trips for computing scores, applying softmax and multiplying by V into one, so intermediate results never touch HBM. The second is online softmax, which uses a correction factor so that the denominator and the output can be computed while scanning, without waiting for the whole row. That\u0026rsquo;s what makes the fusion possible in the first place. The third is tiling, which replaces \u0026ldquo;read one vector at a time\u0026rdquo; with \u0026ldquo;move one block at a time,\u0026rdquo; so the SRAM and the Tensor Cores are actually kept busy.\nBack to the question we started with. FlashAttention doesn\u0026rsquo;t do less arithmetic. Every score in that \\(n \\times n\\) table still gets computed, so attention is still \\(O(n^2)\\) work; it isn\u0026rsquo;t true that FlashAttention turns attention into \\(O(n)\\). What it saves is all the extra data movement between HBM and SRAM. A lot of optimization on modern GPUs isn\u0026rsquo;t about making them compute less. It\u0026rsquo;s about making them wait less.\nReferences Zihao Ye. “From Online Softmax to FlashAttention.” UW CSE 599M lecture notes, 2023. Dao et al. (@tri_dao). “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022. Tri Dao (@tri_dao). “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” ICLR 2024. Horace He (@cHHillee). “Making Deep Learning Go Brrrr From First Principles.” 2022. Biao He (@hebiao064). “What is Flash Attention?” 2025. ","permalink":"https://zhangchitc.github.io/posts/flash-attention/","summary":"FlashAttention does every multiplication standard attention does, yet runs several times faster. We build it up step by step from the GPU memory hierarchy: kernel fusion, safe softmax, online softmax, and tiling.","title":"Why FlashAttention is So Fast"},{"content":"When I was digging into prompt caching in OpenAI\u0026rsquo;s Codex, I noticed something odd: the number of cached tokens always went up in steps of 512. The reason is that before a prompt reaches the model, it gets cut into blocks of a fixed number of tokens, and both prefix matching and the KV cache work one block at a time.\nThat block-based design has a name: PagedAttention. I only sketched the idea back then. This post walks through how PagedAttention works in vLLM, the open-source inference engine it was built for. (I\u0026rsquo;m using vLLM because its implementation is public. I\u0026rsquo;m not claiming OpenAI\u0026rsquo;s servers run the same code.)\nTwo kinds of memory fragmentation To understand PagedAttention, start with the problem it solves. Serving LLMs runs into two problems with how GPU memory gets used.\nThe first is internal fragmentation. When a request starts, you usually don\u0026rsquo;t know how long its output will be, so the simplest approach is to reserve memory for the maximum context length up front. Most requests never use all of it, and whatever they don\u0026rsquo;t use is wasted.\nThe second is external fragmentation. Some systems allocate memory on demand, in regions of different sizes. But as requests finish and free their memory, the free space breaks up into gaps of different sizes scattered between allocations. You might have 1 GB of free memory in total and still be unable to fit a new request, because no single contiguous 1 GB region exists.\nPagedAttention was designed to get rid of both kinds of fragmentation and use GPU memory far more efficiently.\nBorrowing paging from the operating system Operating systems solved this problem decades ago with paging. A program doesn\u0026rsquo;t need one contiguous chunk of physical memory. Its address space is cut into fixed-size pages, each page goes wherever physical memory has room, and a page table maps the contiguous virtual addresses the program sees onto physical pages scattered anywhere.\nPagedAttention brings the same idea to GPU memory, with a block playing the role of a page. In this example the block size is 4 tokens, so prompt 1 and prompt 2 each split into two blocks. The block table in the middle is the page table: it records which physical block each logical block lives in. ABCD sits in physical block 5 and EFGH in physical block 0. The numbers have nothing to do with each other, and blocks that are adjacent logically can sit anywhere physically. Both prompts start with ABCD, so their first blocks point at the same physical block.\nThat\u0026rsquo;s the payoff of paging. Every block is the same size, so any free block can go to any request, and there\u0026rsquo;s never a gap that\u0026rsquo;s too small to use. External fragmentation disappears. The only waste left is inside a request\u0026rsquo;s last block, which may not be full. In vLLM a block is typically 16 tokens, so each request wastes fewer than 16 token slots, which is negligible.\nWhat a physical block looks like Now let\u0026rsquo;s look at how vLLM manages the physical memory. When vLLM starts up, it cuts the GPU memory set aside for the KV cache into physical blocks all at once. The free blocks are chained into a doubly linked list called free_block_queue. When a block is needed, it\u0026rsquo;s taken off this list, and when it\u0026rsquo;s no longer needed, it goes back. Each node in the list is a KVCacheBlock object, the metadata for one physical block. It holds four kinds of information:\nblock_id, the number of the physical block it describes. ref_cnt, how many requests are currently using the block. _block_hash, the hash of the prefix stored in the block, which we\u0026rsquo;ll come back to shortly. prev_free_block and next_free_block, the two pointers that chain free blocks into the list. This free list lives in CPU memory and only does the bookkeeping. The actual blocks live in GPU memory. At the top level they\u0026rsquo;re split into L layers: a Transformer has L decoder layers, and each layer needs its own K and V. So a physical block is really a column that runs through every layer. In the diagram, the dashed yellow outline marks block_id 2, the same slot in every layer. Across a single layer there are num_blocks physical blocks.\nZoom in on one block and it\u0026rsquo;s a tensor of shape [2, block_size, num_kv_heads, head_dim], where the 2 holds K and V.\nSo the total size of the KV cache is all of these numbers multiplied together: L layers, num_blocks blocks per layer, block_size tokens per block, 2 for K and V, the number of KV heads, the dimension of each head, and finally the bytes per number.\n$$ \\begin{aligned} \\text{KV cache} = {}\u0026\\textcolor{#1d4ed8}{L} \\times \\textcolor{#b45309}{\\mathit{num\\_blocks}} \\times \\mathit{block\\_size} \\times \\textcolor{#b45309}{2}\\\\ \u0026\\times \\mathit{num\\_kv\\_heads} \\times \\mathit{head\\_dim} \\times \\textcolor{#6b7280}{\\text{bytes per element}} \\end{aligned} $$ This is the same formula from my previous post on KV cache size. The only difference is that the total context length is now written as num_blocks × block_size.\nBlock hash: telling whether two prefixes match How is the _block_hash in each KVCacheBlock computed? With a chained hash. The formula is short: the hash of block n is the hash of block n−1\u0026rsquo;s hash together with all the tokens in block n.\n$$ \\textcolor{#6d28d9}{\\mathrm{hash}(\\text{Block } n)} = \\mathrm{hash}\\big(\\textcolor{#6d28d9}{\\mathrm{hash}(\\text{Block } n-1)},\\ \\textcolor{#b45309}{\\mathrm{tokens}(n)}\\big) $$\nBecause the hash is chained, block n\u0026rsquo;s hash already folds in every token from the first block through block n. So when two blocks have the same hash, we know that block and everything before it are identical. A matching prefix is exactly what prompt caching needs before it can reuse anything.\nTwo lookup structures PagedAttention in vLLM relies on two more data structures. The first is BlockHashToBlockMap, a map from block hash to KVCacheBlock. There\u0026rsquo;s a single global copy shared by all requests. When a new request arrives, vLLM looks up each of its block hashes in this map. A hit means some earlier request already computed that prefix, and the map says where in GPU memory the result is stored.\nThe second is the block table, called req_to_blocks in the code. For each request it keeps the list of KVCacheBlocks that hold the request\u0026rsquo;s KV cache, in order. This list is exactly the logical-to-physical mapping from the paging picture earlier.\nThat\u0026rsquo;s every concept PagedAttention needs. Now let\u0026rsquo;s follow a complete example.\nWalking through two requests At time 0, the first request arrives. Its prompt is 12 tokens, ABCD through IJKL, which splits into 3 logical blocks at a block size of 4. The three block hashes come out as 0xABCD, 0xABGH and 0xABKL. (Since each hash covers the whole prefix, I\u0026rsquo;m labeling it with the first two and last two tokens of that prefix to keep it readable. Real hashes look nothing like this.) BlockHashToBlockMap starts out empty, so all three lookups miss. Every logical block needs a new physical block, so vLLM takes blocks 0, 1 and 2 off the front of free_block_queue and records them in req_to_blocks[0]. On the GPU side, each of the three blocks gets ref_cnt 1, meaning one request is using it, and records its own hash. Since nothing hit, all three blocks have to go through prefill, so they\u0026rsquo;re all marked COMPUTE.\nAt time 1, the second request arrives. Its first eight tokens are identical to the first request\u0026rsquo;s, followed by XYZ. Its first two block hashes therefore come out as 0xABCD and 0xABGH, matching the first request exactly. The third block holds only three tokens, and there\u0026rsquo;s a rule here: only full blocks get a hash. This block is the tail of the prompt and will soon be filled by tokens from decode, so its hash is computed once it\u0026rsquo;s full.\nLooking up BlockHashToBlockMap, the first two hashes are already there. Those two blocks don\u0026rsquo;t need new physical blocks or copies, because their KV cache is already sitting in physical blocks 0 and 1. Only the third block needs a new physical block, so vLLM takes block 3 off the queue, and req_to_blocks[1] becomes blocks 0, 1 and 3. The ref_cnt of blocks 0 and 1 goes from 1 to 2, meaning two requests now share that memory. Block 3 isn\u0026rsquo;t full, so its hash is empty and its ref_cnt is 1. It\u0026rsquo;s also the only block that needs computing. The results for the first two blocks are already in memory, so they\u0026rsquo;re skipped.\nThat\u0026rsquo;s the core rule for new requests. When a block\u0026rsquo;s prefix hash hits, the block is neither copied nor recomputed. vLLM finds the physical block, adds 1 to its ref_cnt and reads the result directly. Only blocks freshly allocated from the free queue need actual computation.\nWhen a request finishes What happens when a request is done? At time 2, the first request finishes. Its three blocks, 0, 1 and 2, each have their ref_cnt reduced by 1. Blocks 0 and 1 are still in use by the second request, so their counts drop to 1 and they stay put. Block 2 drops to 0, meaning no one is using it, so it goes back into free_block_queue.\nNotice that BlockHashToBlockMap hasn\u0026rsquo;t changed at all. The entry mapping 0xABKL to block 2 is still there. Block 2 can now be handed to a future request, but its computed KV data is still sitting in GPU memory. If some future request produces a hash that matches what\u0026rsquo;s in block 2, that result shouldn\u0026rsquo;t go to waste. It can be reused instead of recomputed.\nAt time 3, the second request finishes too. The ref_cnt of blocks 0, 1 and 3 all drop to zero, and all three go back into the free queue. The hash map still has the same three entries, none deleted.\n(The diagrams simplify the queue order. In vLLM, freed blocks are appended to the end of the queue and new blocks are taken from the front, so the blocks that have sat unused the longest are the first to be reused.)\nSo when a request ends, returning a block to the free queue doesn\u0026rsquo;t mean throwing it away. The KV data in those four physical blocks is still there, and so are the hash map entries. The blocks are only marked as unused and available. Until one of them is actually handed out again, any request with the same prefix can still hit it.\nPicking freed blocks back up At time 4, a third request arrives, and it\u0026rsquo;s an interesting one. Its first eight tokens are again A through H, the same as the first two requests, followed by four Ps and four Qs, for four blocks in total. The first two hashes are again 0xABCD and 0xABGH, and the last two are new: 0xABPP and 0xABQQ. The lookup gives two hits and two misses. The hit blocks, 0 and 1, are sitting in the free queue with ref_cnt 0. But as we just saw, their results are still valid and still in GPU memory, so vLLM pulls them straight out of the queue and sets their ref_cnt from 0 back to 1. None of their KV cache needs recomputing, so prefill skips the first two blocks.\nThe last two blocks missed, so they need new blocks from the free queue: physical blocks 3 and 2, taken from the front. And here\u0026rsquo;s the interesting part. Block 2 still holds the data for hash 0xABKL. Now that block is about to store the four Qs, its old data will be overwritten, so the 0xABKL entry in the hash map is no longer valid. Only at this moment is it evicted, replaced by 0xABQQ. Eviction isn\u0026rsquo;t a periodic cleanup, and it isn\u0026rsquo;t a cache expiring. A hash entry becomes invalid only when its physical block is handed out again and the old contents are about to be overwritten. The rest follows naturally: req_to_blocks[2] is blocks 0, 1, 3 and 2, and only the last two need computing.\nWrapping up To sum up, PagedAttention cuts the KV cache into blocks that can be paged, mapped and shared. Physical blocks in GPU memory are small fixed-size slots for KV tensors. The hash map lets different requests find the same block by its content. ref_cnt keeps a shared block from being freed while someone still uses it. And the free queue lets blocks from finished requests be picked back up by the next request. Together, these make far better use of GPU memory when serving LLMs.\n(The example uses a block size of 4, a tiny pool of blocks and short labels like 0xABCD so everything fits in a diagram. A real scheduler also handles preemption, swapping and more, which this post leaves out.)\nReferences Kwon et al. (@woosuk_k). “Efficient Memory Management for Large Language Model Serving with PagedAttention.” SOSP 2023. vLLM (@vllm_project). “Automatic Prefix Caching.” vLLM documentation. vLLM (@vllm_project). “Paged Attention.” vLLM documentation. Aleksa Gordić (@gordic_aleksa). “Inside vLLM: Anatomy of a High-Throughput LLM Inference System.” 2025. Sankalp (@dejavucoder). “How prompt caching works: Paged Attention and Automatic Prefix Caching plus practical tips.” 2025. ","permalink":"https://zhangchitc.github.io/posts/paged-attention/","summary":"Why serving LLMs fragments GPU memory, and how vLLM\u0026rsquo;s PagedAttention fixes it with OS-style paging: physical blocks, a block table, chained block hashes, and a walkthrough of how blocks are allocated, shared, freed and evicted.","title":"How PagedAttention Works: Paging the KV Cache Like an Operating System"},{"content":"In my last post I explained that a prompt cache stores the KV cache for your prompt. Several readers followed up with the obvious next question: how much physical GPU memory does that KV cache take up? And with memory chips as expensive as they are right now, how do you fit more context into the memory you have?\nThis post works out how big the KV cache is on a real inference server. We\u0026rsquo;ll start from the simplest possible attention and build the formula up one factor at a time.\nA quick recap: what the KV cache stores Inference in a Transformer revolves around attention, the mechanism that lets the model search the context for the tokens most relevant to the one it\u0026rsquo;s about to generate. The token embeddings are projected into three matrices, Q, K and V. Multiplying Q by K transposed gives a table of scores, softmax turns those scores into weights, and the weights are used to take a weighted sum of the rows of V. That sum is the output.\nDuring decode, each new token brings a fresh Query, which is compared against every earlier Key and then used to weight every earlier Value. The Query is new at every step and thrown away afterwards. K and V, on the other hand, are needed again for every token the model generates from then on, so they\u0026rsquo;re kept. That\u0026rsquo;s the KV cache.\nNow some numbers. Say the prompt is \\(n\\) tokens long and each token\u0026rsquo;s vector has \\(d\\) dimensions. Then K and V are each an \\(n \\times d\\) matrix, and the memory the KV cache needs is:\n$$ \\text{KV cache} = \\textcolor{#d97706}{2} \\times n \\times d \\times \\textcolor{#6b7280}{\\text{bytes per element}} $$ The 2 is there because we store two matrices, K and V. Bytes per element is how much space one number takes. With FP16, the most common format, that\u0026rsquo;s 2 bytes.\nMulti-head attention: every head stores its own K and V A single set of Q, K and V has a limitation: it can only judge how tokens relate from one angle. Language is full of different kinds of relationships. Which earlier word does \u0026ldquo;it\u0026rdquo; refer to? Who did something to whom? Is \u0026ldquo;Apple\u0026rdquo; the company or the fruit? One attention pattern struggles to capture all of these at once. What you tend to get is an average of them, which captures none of them precisely.\nMulti-head attention fixes this. Instead of computing attention once, it splits attention into several heads, each with its own Q, K and V, so each head can look at the same input from a different angle in parallel.\nHere\u0026rsquo;s how the computation goes. Each head projects the tokens into its own Q, K and V and computes attention on its own. In this example there are 4 heads, so we get 4 output matrices, each \\(n \\times d\\). We only want one output, though, so the 4 matrices are concatenated side by side into a single \\(n \\times 4d\\) matrix, which holds everything the 4 heads found. A \\(4d \\times d\\) projection matrix, \\(W^O\\), then turns it into the final output of multi-head attention.\n(To keep the picture simple, each head here works in the full d dimensions. In real models each head usually works in a smaller slice, so the heads\u0026rsquo; dimensions add up to the model\u0026rsquo;s width. Qwen3-8B, which we\u0026rsquo;ll meet at the end, has 32 heads of 128 dimensions each, for a width of 4,096. From here on, \\(d\\) in the formula means the dimension of one head.)\nThe important part for memory is that every head computes its own K and V, and all of them have to be kept. So the KV cache grows by a factor of \\(h\\), the number of attention heads:\n$$ \\text{KV cache} = \\textcolor{#d97706}{2} \\times n \\times d \\times \\textcolor{#d97706}{h} \\times \\textcolor{#6b7280}{\\text{bytes per element}} $$ Stacked decoder layers: different K and V at every layer Everything so far happens inside a single decoder layer. A real Transformer stacks many of these layers on top of each other.\nWhy stack them? In one attention layer, each token gets only one chance to mix in information from the context, and that\u0026rsquo;s not enough to build much understanding. Language has to be understood in layers of abstraction. It\u0026rsquo;s a bit like a convolutional network recognizing an image: the bottom layers see only pixels, the next layers assemble pixels into edges and textures, the layers above that combine those into eyes and ears, and only the top layers recognize the animal in the picture. Stacking many identical decoder layers lets a Transformer refine each token\u0026rsquo;s representation the same way, step by step, into something more abstract and complete.\nEvery layer has the same structure. First comes the multi-head attention we just covered, then a feed-forward network, and each of the two is followed by an Add \u0026amp; Norm step: a residual connection plus normalization. Those steps mainly keep a network dozens of layers deep stable during training. They don\u0026rsquo;t take part in attention and don\u0026rsquo;t produce any K or V, so they don\u0026rsquo;t affect the KV cache. The token embeddings enter the first layer, each layer\u0026rsquo;s output becomes the next layer\u0026rsquo;s input, and after the top layer the model uses the final output to predict the next token.\nBecause every layer has its own \\(h\\) sets of K and V, the formula picks up another factor, \\(L\\), the number of decoder layers:\n$$ \\text{KV cache} = \\textcolor{#d97706}{2} \\times n \\times d \\times \\textcolor{#d97706}{h} \\times \\textcolor{#d97706}{L} \\times \\textcolor{#6b7280}{\\text{bytes per element}} $$ With this many factors, it\u0026rsquo;s fair to ask whether every layer\u0026rsquo;s K and V really need to be kept. They do. Every new token passes through the whole stack starting from the first layer, computing attention at each one. If any layer\u0026rsquo;s K and V weren\u0026rsquo;t cached, that layer\u0026rsquo;s cost per decode step would jump from the order of \\(n\\) to the order of \\(n^2\\). However fast the other layers are, that one layer would hold up the whole decode, so the K and V of every layer have to stay.\nMQA and GQA: replacing h with h_kv Multi-head attention is the most expressive option, since every head can look at the context from its own angle. The downside is just as clear: with many heads, each one storing its own K and V puts enormous pressure on memory. Engineers have come up with several ways around this.\nThe most aggressive is multi-query attention (MQA). Every head keeps its own Q, so the model still looks from many angles, but all the heads share a single K and V. That cuts the K and V you need to store by a factor of \\(h\\). The memory savings are real, but so is the cost: with every head attending through the same K and V, much of the benefit of having multiple heads is lost, and output quality drops noticeably.\nThe compromise is grouped-query attention (GQA), which is what most new models use today. Instead of giving every head its own K and V like MHA, or making every head share one like MQA, it splits the heads into groups, and each group shares one K and V. Pair up 8 heads, for example, and you only need 4 sets of K and V. Seen this way, MHA and MQA are both special cases of GQA: MQA is GQA with one group, and MHA is GQA with as many groups as heads. GQA gets the best of both, striking a good balance between expressiveness and memory.\nThat gives us the final formula, where \\(h_{\\text{kv}}\\) is the number of K and V heads actually stored:\n$$ \\text{KV cache} = \\textcolor{#d97706}{2} \\times n \\times d \\times \\textcolor{#059669}{h_{\\text{kv}}} \\times \\textcolor{#d97706}{L} \\times \\textcolor{#6b7280}{\\text{bytes per element}} $$ Checking the formula against a real calculator We can check this against a real tool. The LMCache project has a KV cache calculator on its documentation site that works out the KV cache size for open models. With its default model, Qwen3-8B, the data type set to BF16 and the context set to 100,000 tokens, it reports that the KV cache for this one request takes 13.73 GB of GPU memory.\nQwen3-8B has 36 layers and 32 attention heads, but only 8 KV heads, and each head has 128 dimensions. BF16 takes 2 bytes per element. Plug those into the formula:\n$$ \\begin{aligned} \\text{KV cache} \u0026= \\textcolor{#d97706}{2} \\times \\underbrace{100{,}000}_{n} \\times \\underbrace{128}_{d} \\times \\underbrace{\\textcolor{#059669}{8}}_{h_{\\text{kv}}} \\times \\underbrace{\\textcolor{#d97706}{36}}_{L} \\times \\underbrace{\\textcolor{#6b7280}{2}}_{\\text{bytes}}\\\\ \u0026= 14{,}745{,}600{,}000\\ \\text{bytes} \\approx 13.73\\ \\text{GiB} \\end{aligned} $$ That\u0026rsquo;s about 7.37 billion numbers, or 14.7 billion bytes. The calculator labels the result GB, but it divides by 1024³, so strictly speaking it\u0026rsquo;s 13.73 GiB, or about 14.7 GB in decimal units. Either way, it matches the formula exactly.\n(This counts only the K and V tensors themselves. The total memory a GPU process uses also includes the model weights, activations and runtime workspace, so this formula can\u0026rsquo;t predict that total on its own.)\nReferences Vaswani et al. (@ashVaswani). “Attention Is All You Need.” NeurIPS 2017. Ainslie et al. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” EMNLP 2023. LMCache (@lmcache). “KV Cache Size Calculator.” LMCache documentation. ","permalink":"https://zhangchitc.github.io/posts/kv-cache-memory/","summary":"We build the KV cache size formula one factor at a time, from single-head attention to multi-head, stacked layers and MQA/GQA, then check it against a real calculator: 13.73 GiB for 100,000 tokens on Qwen3-8B.","title":"How Much GPU Memory Does the KV Cache Really Take?"},{"content":"I\u0026rsquo;ve covered prompt caching a few times now, and the same questions keep coming back. Does the server save the model\u0026rsquo;s previous answer? When my prompt prefix hits the cache, does the model skip running altogether and just hand back a stored response, like a database lookup?\nBoth come from common misconceptions about how LLMs work, and they boil down to the same question: what exactly is sitting in a prompt cache? Let\u0026rsquo;s answer that properly. Along the way we\u0026rsquo;ll meet its close cousin, the KV cache, and by the end the difference between the two will be clear.\nTwo phases: prefill and decode To get there, we need to start with how a Transformer generates a response. It happens in two phases.\nThe first is prefill. The model reads your entire prompt in one go, building up its understanding of the context as it reads, so it\u0026rsquo;s ready to start answering.\nThe second is decode, where the model writes the answer. Here\u0026rsquo;s the catch: even if your prompt is a thousand tokens long, prefill reads it all at once, but on the way out the model can only produce one token at a time. And before each new token, it has to look back over everything that came before: your whole prompt, plus every token it has generated so far. Generating text this way is called autoregressive.\nPrompt caching mainly speeds up the first phase, prefill.\nA quick refresher on attention To see why, we need to revisit attention, the core of the Transformer and the key to understanding the KV cache.\nAttention exists because, when the model predicts the next token, the earlier words aren\u0026rsquo;t equally important. It needs a way to tell which ones matter most right now.\nHere\u0026rsquo;s an example. The prompt is \u0026ldquo;The quick brown fox jumps\u0026rdquo;, and the model has to predict what comes next. In the diagram, thicker arrows mean more attention. The thickest ones land on \u0026ldquo;quick\u0026rdquo;, \u0026ldquo;brown\u0026rdquo; and \u0026ldquo;fox\u0026rdquo;, while \u0026ldquo;The\u0026rdquo; and \u0026ldquo;jumps\u0026rdquo; barely register. The model outputs \u0026ldquo;over\u0026rdquo;.\n(The weights are illustrative, not measured from a real model.)\nWhy \u0026ldquo;over\u0026rdquo;? Because \u0026ldquo;The quick brown fox jumps over the lazy dog\u0026rdquo;, the pangram that packs every letter of the alphabet into one sentence and is a staple of typing drills and font previews, appears in its training data more times than anyone can count. That co-occurrence is baked into its weights. Once it spots \u0026ldquo;quick brown fox\u0026rdquo; as the words that matter, \u0026ldquo;over\u0026rdquo; is the obvious next pick. That\u0026rsquo;s attention in a nutshell: instead of treating every word equally, it weighs them by how much they matter at this moment.\nSo how is it computed? Every token gets projected into three vectors: a Query, a Key and a Value.\nThe Query is what this token is looking for. The Key is what this token can offer. The Value is the information this token actually carries. The names come from information retrieval and databases. The Query is like what you type into a search box, the Key is like an index key in a database, and the Value is the data you get back for that key.\nThe current token takes its Query and scores it against the Key of every earlier token. A higher score means that token is more relevant. Those scores then become weights for a weighted sum of all the tokens\u0026rsquo; Values. That\u0026rsquo;s the whole attention computation, and the formula says exactly the same thing:\n$$ \\text{Attention}(Q,K,V)=\\text{softmax}\\!\\left(\\frac{QK^{\\top}}{\\sqrt{d_k}}\\right)V $$ \\(QK^{\\top}\\) is \u0026ldquo;score the current Query against every Key\u0026rdquo;. Dividing by \\(\\sqrt{d_k}\\) is a scaling trick from the original paper that keeps the scores from getting too large. The softmax normalizes the scores into weights that sum to 1. Multiplying by \\(V\\) is \u0026ldquo;take the weighted sum of every token\u0026rsquo;s Value\u0026rdquo;. What prefill and decode compute Back to the computation. Say your prompt has \\(n\\) tokens, and each token\u0026rsquo;s Q, K and V vectors have \\(d\\) dimensions. Stack the \\(n\\) Query vectors row by row and you get an \\(n \\times d\\) matrix. The same goes for K and V.\nHere\u0026rsquo;s what prefill does. It computes all three matrices for the whole prompt. Then it multiplies Q by K transposed in a single step, producing an \\(n \\times n\\) table of scores between every pair of tokens in the prompt. That table is used to weight the entire V matrix. The attention for the whole prompt is done in one shot.\nBecause that table is \\(n \\times n\\), it grows quickly as the prompt gets longer. The amount of computation is huge, which is why prefill is limited by the GPU\u0026rsquo;s raw compute.\nNow decode. Generating one new token adds just one row for that token: its own Q, K and V, each a \\(1 \\times d\\) vector. So the K and V matrices each grow by one row, to \\(n+1\\). The new token\u0026rsquo;s Query is compared against every row of K, and the result is used to take a weighted sum over V.\nThe second token, the third, and every one after that work the same way: K and V keep growing by one row at a time. And all the earlier rows get used again for every single new token.\nWhat the KV cache stores, and how much it saves Decoding one token at a time like this means recomputing K and V for every earlier token at every step. But once those K and V are computed, they never change. So engineers had the obvious idea: keep them around during decode. That\u0026rsquo;s the KV cache.\nWhy store K and V but not Q? A token\u0026rsquo;s Query is only used once, at the step that produces that token. At the next step there\u0026rsquo;s a new token with a new Query, and the old one is useless, so storing it would be a waste. K and V, on the other hand, never change once computed and get pulled out again at every decode step, so they\u0026rsquo;re worth keeping.\nThe KV cache exists to speed up every decode step. Let \\(N\\) be the current context length. Without a KV cache, every new token means recomputing everything for all \\(N\\) earlier tokens from scratch, which costs on the order of \\(N^2\\) per step. With a KV cache, the model only computes K and V for the one new token and reads the rest from the cache, bringing the cost per step down to the order of \\(N\\).\nThat\u0026rsquo;s a factor of \\(N\\) less work per step. With a few thousand tokens of context, each step does thousands of times less work, and with tens of thousands of tokens, tens of thousands of times less.\nPrefill fills the cache, decode appends to it With a KV cache in place, here\u0026rsquo;s what the full inference process looks like.\nDuring prefill, the model computes K and V for every token in the prompt at once and writes them into an empty KV cache. Prefill is where the KV cache gets initialized.\nDuring decode, each new token only needs its own K and V, which are appended to the end of the cache. Decode keeps adding one row at a time.\nIn the diagram below, the orange arrows are writes: prefill fills the cache from above, and each decode step appends one row from below (the highlighted last column of \\(K^{\\top}\\) and last row of V). The gray arrows on the right are reads: both the prefill and the decode attention take K and V straight from the cache.\nSo, back to the question we started with: what exactly is stored in a prompt cache? It\u0026rsquo;s that KV cache, the one prefill builds after your prompt has gone through the Transformer. It\u0026rsquo;s an intermediate result from halfway through the model\u0026rsquo;s computation. It\u0026rsquo;s not your last answer, and it\u0026rsquo;s certainly not a string that can be handed straight back to you.\nWhat changes with a prompt cache So far, the KV cache only removes repeated work within a single response. Inside one request, decode doesn\u0026rsquo;t recompute the earlier K, V or attention. A prompt cache extends that reuse across requests.\nWhat does that mean? Without a prompt cache, prefill starts from scratch for every new request. Your system prompt, your tool definitions, the earlier turns of the conversation: K and V get recomputed for every one of those tokens.\nWith a prompt cache, things change. Prefill notices that this request\u0026rsquo;s prefix is identical to last time, so it pulls out the K and V it stored back then and uses them directly, skipping that whole stretch of attention computation. The model only has to process the small part you added this time. Once that\u0026rsquo;s done, it joins up with the cached part and decode carries on as usual. (The blue marker in the diagram shows where the cached prefix ends, which Anthropic\u0026rsquo;s API sets with cache_control.)\nNote that a cache hit doesn\u0026rsquo;t mean you get a ready-made answer. What it saves is the prefill computation. The model still has to generate every token of its reply. That\u0026rsquo;s also why cached tokens are so much cheaper: you\u0026rsquo;re not making the GPU recompute them.\nKV cache vs prompt cache To wrap up, here\u0026rsquo;s how the two compare.\nIn one sentence: the KV cache saves repeated decode work within a request, and the prompt cache saves repeated prefill work across requests. Used together, one covers decode and the other covers prefill.\nReferences Vaswani et al. (@ashVaswani). “Attention Is All You Need.” NeurIPS 2017. Sankalp (@dejavucoder). “How prompt caching works: Paged Attention and Automatic Prefix Caching plus practical tips.” 2025. vLLM. “Automatic Prefix Caching.” vLLM documentation. Gim et al. “Prompt Cache: Modular Attention Reuse for Low-Latency Inference.” MLSys 2024. ","permalink":"https://zhangchitc.github.io/posts/kv-cache-vs-prompt-cache/","summary":"A prompt cache doesn\u0026rsquo;t store your last answer. We walk through prefill, decode and attention to see what the KV cache keeps, and how a prompt cache reuses it across requests.","title":"KV Cache vs. Prompt Cache: What's the Difference?"},{"content":"\u0026ldquo;Dog bites man\u0026rdquo; and \u0026ldquo;man bites dog\u0026rdquo; contain exactly the same words. Swap two of them and the meaning flips.\nLeft to itself, a Transformer can\u0026rsquo;t tell these two sentences apart. Attention, the mechanism at the heart of the model, never looks at where a token sits. As far as attention is concerned, a sentence is a bag of tokens, not a sequence.\nThis post walks through why Transformers need positional encoding, and how the sine-and-cosine scheme from the original paper captures position so neatly.\nWhat attention sees without positions Let\u0026rsquo;s start with plain attention and see what goes wrong when it has no position information. I won\u0026rsquo;t re-derive attention here. All we need is the formula:\n$$ \\text{Attention}(Q,K,V)=\\text{softmax}\\!\\left(\\frac{QK^{\\top}}{\\sqrt{d_k}}\\right)V $$ Multiply \\(Q\\) by \\(K\\) transposed, scale by \\(\\sqrt{d_k}\\), apply softmax, and use the resulting weights to mix the value vectors in \\(V\\).\nNow take two sentences:\n\u0026ldquo;Alice hits Bob. Who\u0026rsquo;s hurt?\u0026rdquo; \u0026ldquo;Bob hits Alice. Who\u0026rsquo;s hurt?\u0026rdquo; Suppose the model has just read the last token, \u0026ldquo;hurt\u0026rdquo;, and is about to predict the next one. Here\u0026rsquo;s what happens in that decoding step.\nIn the first sentence, the query vector for \u0026ldquo;hurt\u0026rdquo; is dotted with the key vector of every earlier token. That produces a row of scores: \\(q_{\\text{hurt}} \\cdot k_{\\text{Alice}}\\), \\(q_{\\text{hurt}} \\cdot k_{\\text{hits}}\\), \\(q_{\\text{hurt}} \\cdot k_{\\text{Bob}}\\), and so on. Softmax turns those scores into weights, and the output is a weighted sum of the tokens\u0026rsquo; value vectors.\nIn the second sentence, Alice and Bob trade places, so the scores come out as \\(q_{\\text{hurt}} \\cdot k_{\\text{Bob}}\\), \\(q_{\\text{hurt}} \\cdot k_{\\text{hits}}\\), \\(q_{\\text{hurt}} \\cdot k_{\\text{Alice}}\\). It\u0026rsquo;s the same set of terms in a different order, and the value vectors get reshuffled in exactly the same way.\nThat\u0026rsquo;s the whole problem. A value vector\u0026rsquo;s weight depends only on its own score, not on where it shows up in the list. So every value vector gets the same weight in both sentences, and since addition doesn\u0026rsquo;t care about order, both sentences produce exactly the same output vector.\nGeneration is autoregressive, so the same output means the model continues both sentences the same way. It has no way to tell who\u0026rsquo;s hurt. On its own, attention is blind to word order.\nTwo obvious fixes: integers and fractions The most direct fix is to add each token\u0026rsquo;s position number to its embedding.\nSay the sequence is 64 tokens long and \u0026ldquo;Bob\u0026rdquo; is token 24. Every token is a \\(d\\)-dimensional embedding vector.\nThe naive version adds 24 to every dimension of Bob\u0026rsquo;s embedding. It\u0026rsquo;s simple, but the problem is obvious. Embedding values are small numbers like the ones below, and position numbers are one or two orders of magnitude larger. Add the two together and the position drowns out whatever the embedding was saying about the word.\nSo let\u0026rsquo;s scale it down. Divide the position by the sequence length \\(L\\), and every value lands between 0 and 1.\nNow the position no longer drowns out the embedding, but we\u0026rsquo;ve traded one problem for another: the encoding depends on how long the sequence is. Token 24 gets 0.375 in a 64-token sequence but 0.024 in a 1,000-token one.\n$$ \\begin{aligned} L = 64: \u0026\\quad \\frac{24}{64} = {\\color{#0a7d61}0.375} \\\\[4pt] L = 1000: \u0026\\quad \\frac{24}{1000} = {\\color{#d64545}0.024} \\end{aligned} $$ The same position gets a different code every time the length changes, so the model can never learn a stable meaning for any particular value.\nBinary: bounded and length-independent, but discontinuous Let\u0026rsquo;s try again. Another way to write down a position is in binary.\n24 in binary is 011000. Every digit is either 0 or 1, so the values stay in a fixed range. And position 24 is always 011000, however long the sequence is. Add those six bits to six dimensions of the embedding, one bit each, and you have a binary positional encoding.\nIt helps to picture each bit as a square wave. The horizontal axis is position, from 0 to 63, and each row tracks one bit:\nThe top row is the lowest bit. It flips at every position. The next row flips every 2 positions. The rows below flip every 4, 8, 16 and 32 positions. To read the code for a position, read down its column. For 24, going from the lowest bit to the highest, you get 0, 0, 0, 1, 1, 0. Position 0 is all zeros and position 63 is all ones.\nBinary gets a lot right, but it has one big flaw: it changes in jumps. Look at what happens between 23 and 24:\n$$ \\begin{aligned} 23 \u0026= 0\\,1\\,0\\,1\\,1\\,1_{(2)} \\\\[2pt] 24 \u0026= 0\\,1\\,{\\color{#d64545}\\mathbf{1\\,0\\,0\\,0}}_{(2)} \\end{aligned} $$ One step forward flips four bits at once. That\u0026rsquo;s bad news for a neural network. Networks learn by gradient descent, and gradient descent works best on smooth, continuous functions: if the position changes a little, the encoding should change a little too. That\u0026rsquo;s what lets the gradient tell the model which way things are moving.\nSinusoidal encoding: smoothing out the steps This is where \u0026ldquo;Attention Is All You Need\u0026rdquo;, the paper that introduced the Transformer, comes in. Its answer is sinusoidal positional encoding.\nThe idea is to split the encoding vector\u0026rsquo;s \\(d\\) dimensions into pairs. In each pair, the even dimension gets a sine and the odd dimension gets a cosine:\nBoth lines share the same argument. One takes its sine and the other its cosine. Here \\(pos\\) is the position, \\(i\\) is the index of the pair, \\(d\\) is the size of the vector, and 10000 is a fixed constant.\nIf the formula looks arbitrary, don\u0026rsquo;t worry. We\u0026rsquo;ll come at it from a few angles until the design makes sense.\nFirst, let\u0026rsquo;s plot it. Notice how much it resembles the square waves from before. The first pair changes fastest: nudge the position and its values already move. Each pair after that changes more slowly.\nThe difference is that the square waves have become smooth curves. Where binary jumped, these change continuously.\nIt ticks the other boxes too:\nBounded: sine and cosine always stay between \\(-1\\) and \\(1\\). Length-independent: the formula only involves \\(pos\\), \\(i\\) and \\(d\\). The sequence length \\(L\\) appears nowhere, so a position\u0026rsquo;s encoding depends only on the position itself. Every requirement we\u0026rsquo;ve collected so far is met.\nWavelengths: a geometric sequence from 2π to 2π·10000 Now for a different angle: the wavelength of each sine wave.\nA quick refresher. The basic sine function, \\(y = \\sin x\\) (blue), has a wavelength of \\(2\\pi\\), meaning the curve repeats every \\(2\\pi\\).\nWhat about \\(y = \\sin(x/2)\\) (orange)? Its wavelength is \\(4\\pi\\). Because \\(x\\) is divided by 2, it has to travel twice as far to complete one cycle, so the wave is stretched to twice the length.\nThe general rule: dividing \\(x\\) by \\(a\\) stretches the wave by a factor of \\(a\\).\n$$ y = \\sin\\frac{x}{a} \\quad\\Longrightarrow\\quad \\text{wavelength} = 2\\pi \\cdot a $$ In the positional encoding, the divisor is \\(10000^{2i/d}\\), so the wavelength of pair \\(i\\) is:\n$$ y = \\sin\\frac{x}{10000^{2i/d}} \\quad\\Longrightarrow\\quad \\text{wavelength} = 2\\pi \\cdot 10000^{2i/d} $$ Let\u0026rsquo;s plug in \\(d = 8\\), which gives four pairs.\nFor \\(i = 0\\), the exponent \\(2i/d\\) is 0, and \\(10000^0 = 1\\). The wavelength is \\(2\\pi\\), so this pair repeats roughly every 6 positions. For \\(i = 1\\), the exponent is 0.25, and \\(10000^{0.25}\\) is exactly 10. The wavelength is \\(20\\pi\\). The pattern continues: each pair\u0026rsquo;s wavelength is 10 times the previous one. (The curves on the right show the pattern, not the true scale. Each row is drawn twice as wide as the one above. At the real 10× ratio, the bottom rows would look flat.)\nIn general, the wavelengths form a geometric sequence that starts at \\(2\\pi\\) and climbs toward \\(2\\pi \\cdot 10000\\) for the last pair. You can read a wavelength as the number of positions before that pair\u0026rsquo;s values repeat.\nThis is the same structure as binary. In binary, the bits flip every 1, 2, 4, 8 positions, doubling each time. In the sinusoidal encoding, the ratio between neighbors is set by 10000 and \\(d\\) instead of being 2. Either way, low dimensions change fast and high dimensions change slowly.\nWavelengths also explain the choice of 10000. The longest wavelength is close to\n$$ 2\\pi \\cdot 10000 \\approx 62{,}832 \\ \\text{tokens} $$ So the slowest pair doesn\u0026rsquo;t repeat for roughly 60,000 positions, which means any sequence shorter than that gets a distinct code at every position. That\u0026rsquo;s plenty for typical context lengths. You could use a bigger base, like 100,000, but there\u0026rsquo;s little reason to: when processing text, a token rarely needs context from 100,000 tokens back.\nEach sin/cos pair is a clock hand One question is left: why does the encoding come in sine/cosine pairs?\nRecall from high-school trig that the cosine and sine of the same angle are the coordinates of a point on the unit circle.\nSo let\u0026rsquo;s rewrite the formula. Define \\(\\omega_i\\) as the reciprocal of the divisor, which physicists call the angular frequency:\n$$ \\omega_i = \\frac{1}{10000^{2i/d}} $$ Pair \\(i\\) then becomes \\(\\sin(\\omega_i \\cdot pos)\\) and \\(\\cos(\\omega_i \\cdot pos)\\), and \\(\\omega_i \\cdot pos\\) is simply an angle.\nSeen this way, the encoding becomes intuitive: each sine/cosine pair is a hand on a clock. Every time \\(pos\\) goes up by 1, the hand rotates counterclockwise by \\(\\omega_i\\).\nWe saw earlier that \\(10000^{2i/d}\\) grows with \\(i\\). Since \\(\\omega_i\\) is its reciprocal, it shrinks as \\(i\\) grows, so higher pairs turn more slowly.\nThe four circles below show the four hands at the same \\(pos\\). The low pairs spin quickly, and each pair up turns more slowly. At normal speed you can barely see the bottom two move, but speed things up and they start turning:\nSound familiar? It\u0026rsquo;s a clock face with a second hand, a minute hand and an hour hand. No single hand tells you the exact time, but together they pin down one unique moment.\nPositional encoding works the same way. Each pair of dimensions is a hand, and the position \\(pos\\) is the time. The wavelengths from the previous section get a new meaning too: a pair\u0026rsquo;s wavelength is how many positions its hand needs to go around once.\nReferences [1] Christopher Fleetwood. “You could have designed state of the art positional encoding.” Hugging Face Blog, 2024.\n[2] Naoki Shibuya. “Transformer’s Positional Encoding.” 2021.\n[3] Vaswani et al. “Attention Is All You Need.” NeurIPS 2017.\n","permalink":"https://zhangchitc.github.io/posts/positional-encoding/","summary":"Attention can\u0026rsquo;t tell \u0026ldquo;dog bites man\u0026rdquo; from \u0026ldquo;man bites dog\u0026rdquo;. We try integer, fractional and binary encodings, see where each one breaks, and arrive at sines and cosines: a set of clock hands spinning at different speeds.","title":"Why Transformers Need Positional Encoding"},{"content":"Hi, I\u0026rsquo;m Chi Zhang. I studied computer science at Peking University and have spent more than a decade building consumer products at Facebook and Airbnb. I started out on Facebook\u0026rsquo;s growth team in Kuala Lumpur, went on to work at both companies\u0026rsquo; headquarters in the Bay Area, and later joined the founding team of Airbnb\u0026rsquo;s Beijing office, where I helped build the engineering team from the ground up.\nThese days most of my curiosity goes into AI. Working with coding agents like Claude Code and Codex made me curious about what happens underneath, both in the agents themselves and in the LLMs that power them. So I started taking them apart. I capture the real requests, read the papers, redo the math step by step, and turn what I find into visual explainers.\nMost of this work first appears as short videos in Chinese, under the name 张司机在路上 (\u0026ldquo;Driver Zhang on the road\u0026rdquo;). This blog is where the English versions live, with the reasoning written out in full.\nOff the keyboard Canadian CASI Level 2 snowboard instructor Licensed sport pilot (fixed-wing aircraft) in China Hiker who has summited Kilimanjaro, Mount Fuji and Mount Kinabalu, Malaysia\u0026rsquo;s highest peak Scuba diver and freediver Foodie who once documented a week of breakfasts at Airbnb\u0026rsquo;s San Francisco HQ Get in touch The easiest way to reach me is on X or LinkedIn. Everything else is linked from the homepage.\n","permalink":"https://zhangchitc.github.io/about/","summary":"\u003cp\u003eHi, I\u0026rsquo;m Chi Zhang. I studied computer science at Peking University and have spent more than a decade building consumer products at Facebook and Airbnb. I started out on Facebook\u0026rsquo;s growth team in Kuala Lumpur, went on to work at both companies\u0026rsquo; headquarters in the Bay Area, and later joined the founding team of Airbnb\u0026rsquo;s Beijing office, where I helped build the engineering team from the ground up.\u003c/p\u003e\n\u003cp\u003eThese days most of my curiosity goes into AI. Working with coding agents like Claude Code and Codex made me curious about what happens underneath, both in the agents themselves and in the LLMs that power them. So I started taking them apart. I capture the real requests, read the papers, redo the math step by step, and turn what I find into visual explainers.\u003c/p\u003e","title":"About"}]