The problem no research paper mentions: serving fifty million users
Picture an Indian startup on the scale of Sarvam AI or Krutrim shipping a Hindi-English customer-support model built on a 7-billion-parameter transformer decoder, and a Swiggy-scale platform adopting it for tens of millions of chat sessions a day. The 2017 paper that introduced the architecture, Vaswani et al.'s "Attention Is All You Need," proved that stacked self-attention could learn to translate and generate language better than anything that came before it. It said almost nothing about what happens when that same architecture has to generate one token at a time, for millions of concurrent users, on GPUs that cost real money per hour to rent. That gap, between an architecture that works on a benchmark and an architecture that ships in a product, is where nearly every transformer actually deployed today (GPT-4, Llama, Gemini, Sarvam's own models) diverges from the block diagram in the original paper.
You have already seen, in earlier chapters, how scaled dot-product attention computes softmax(QKᵀ/√d)V, and how encoder and decoder layers stack with residual connections and normalization. This chapter assumes that machinery and goes somewhere the diagram-level view never has to: what changes when the same architecture has to run millions of times a second, one token at a time, on hardware with a fixed, finite amount of memory. Four specific engineering decisions, each traceable to a named paper, separate a transformer that exists in a proof of concept from one that exists in production: the KV cache, grouped-query attention, rotary position embeddings, and a pair of quiet substitutions inside every block, RMSNorm and SwiGLU. Every one of them shows up, by name, in the config file of Llama, Mistral, or any other open-weights model you can inspect today.
Why generating text one token at a time is worse than it looks
A decoder-only transformer generates text autoregressively: predict the next token, append it to the sequence, feed the whole thing back in, predict the next one. The naive way to implement this is to literally do that: at generation step t, run the entire sequence of t tokens through every layer of the network from scratch, and read off the prediction for position t. This is correct, and it is enormously wasteful, because token 1's key and value vectors at layer 5 are exactly the same at step 50 as they were at step 2. Nothing about token 1's representation changes once it has been computed; only the query for the newest token is new.
Quantify the waste. At step t, computing attention for a sequence of length t costs roughly O(t²·d) for the attention score matrix and weighted sum, plus O(t·d²) for the linear projections that produce Q, K, and V, where d is the model width. If you naively recompute the full forward pass at every one of n generation steps, the attention term alone sums to roughly O(n³·d) across the whole generation (the sum of t² for t = 1 to n grows like n³/3), and the projection term sums to O(n²·d²). Generating a 500-token response this way does not cost 500 times as much as generating one token; it costs on the order of 500³ relative to a single step's marginal attention work. That cubic blowup, not the attention formula itself, is the first thing any production system has to eliminate.
The KV cache: compute once, reuse for the rest of the generation
The fix is to notice that a token's key and value vectors, once computed in a given layer, never change. So instead of recomputing them, store them. For every layer and every attention head, the KV cache holds the K and V vectors for every token generated so far. At step t, the model computes Q, K, and V for only the new token (O(d²) work, independent of how long the sequence has grown), appends the new K and V to the cache, and computes attention between the new query and every cached key (O(t·d) work). Summed across n steps, this brings total cost down to O(n·d²) for projections plus O(n²·d) for attention, cutting the naive cubic baseline down to quadratic by removing every bit of redundant recomputation of the past.
The cache has a concrete memory cost, and it is worth computing exactly, because it is what actually constrains how many users a GPU can serve at once. For a decoder with L layers, H key/value heads, head dimension d_head, and a sequence of length T, stored in b bytes per number (2 for fp16), the cache needs 2 (one copy for K, one for V) × L × H × d_head × T × b bytes, per sequence in the batch.
def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, bytes_per_val=2, batch=1):
return 2 * n_layers * n_kv_heads * head_dim * seq_len * batch * bytes_per_val
mha = kv_cache_bytes(n_layers=32, n_kv_heads=32, head_dim=128, seq_len=4096)
gqa = kv_cache_bytes(n_layers=32, n_kv_heads=8, head_dim=128, seq_len=4096)
print("MHA (32 KV heads):", round(mha / 1024**3, 2), "GiB")
print("GQA (8 KV heads):", round(gqa / 1024**2, 1), "MiB")Trace it by hand: 2 × 32 layers × 32 heads × 128 dims × 4096 tokens × 2 bytes = 2,147,483,648 bytes, exactly 2.0 GiB, for one 4096-token conversation, in one layer's worth of heads multiplied across all 32 layers. That is the output of the first print statement: MHA (32 KV heads): 2.0 GiB. These are the real published dimensions of Llama 2 7B (32 layers, 32 heads, d_model = 4096, so head_dim = 128; Touvron et al., 2023). The second call previews a technique explained in the next section: sharing key/value heads across groups of query heads cuts the head count from 32 to 8, and the same arithmetic gives 536,870,912 bytes, exactly 512.0 MiB, a quarter of the plain-MHA figure. That is GQA (8 KV heads): 512.0 MiB.
The misconception: "the KV cache makes generation linear"
It is tempting to conclude, from the fact that the KV cache removes redundant recomputation, that generation cost per token becomes constant, i.e. that a 10,000-token response costs roughly ten times a 1,000-token one and nothing worse. That is wrong, and it is wrong in a way that matters for anyone estimating serving cost. The cache eliminates the wasted re-derivation of old tokens' K and V, which is real and valuable, but it does not eliminate the attention computation itself: at step t, the new query still has to be compared against all t cached keys and the result still has to be weighted-summed against all t cached values. That per-step cost grows linearly with the current sequence length, and summed across a full n-token generation, total attention cost is still O(n²), not O(n). The KV cache converts a naive cubic-in-n baseline into a quadratic one; it does not convert a quadratic mechanism into a linear one. This is precisely why context length, not just token count, drives serving cost, and why systems research on inference (Pope et al., 2022, "Efficiently Scaling Transformer Inference") treats the growing cache, not the fixed model weights, as the dominant cost driver for long conversations.
Grouped-Query Attention: giving heads fewer things to remember
If the cache is the bottleneck, the next question is whether every one of 32 query heads genuinely needs its own private set of keys and values, or whether heads can share. Two prior answers exist. Multi-Query Attention (Shazeer, 2019, "Fast Transformer Decoding: One Write-Head is All You Need") takes the extreme position: keep 32 separate query heads, but give them all one single shared key/value head. This shrinks the cache dramatically but also measurably hurts output quality, because a single shared K/V representation is a real bottleneck on what different heads can attend to. Grouped-Query Attention (Ainslie et al., 2023, "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints") is the compromise that production models actually converged on: split the query heads into groups (say, four groups of eight, or eight groups of four), and give each group its own shared key/value head, rather than one shared head for everyone or one head per query.
The diagram below shows the three patterns side by side, using a schematic of 8 query heads per panel for legibility (the worked example above uses the real 32-head configuration). In each panel, the blue boxes are query heads and the red boxes are the key/value heads they read from; the number of red boxes is the number of separate K/V representations the cache actually has to store.
Tie this back to the hook. A 7B-class model's weights, stored in fp16, take about 13 GiB (7 billion parameters × 2 bytes ≈ 14 billion bytes ≈ 13 GiB). On a GPU with 80 GB of HBM, after loading weights and reserving roughly 7 GiB for activations and overhead, about 60 GiB remains for KV caches. With plain multi-head attention at 2.0 GiB per 4096-token conversation, that GPU can hold about 30 conversations at once. With GQA at 512 MiB per conversation (the same 4× reduction shown in the diagram: 8 KV heads instead of 32), the same 60 GiB holds about 120 conversations, a 4× increase in how many users a single GPU can actually serve simultaneously, for the same model weights and no change to the attention formula's shape. Meta's own Llama 2 70B uses this technique in production, with 64 query heads sharing just 8 KV heads (Touvron et al., 2023); its 7B and 13B siblings do not use GQA, which is exactly why the smaller checkpoints are proportionally more expensive to serve at long context, parameter-for-parameter, than the 70B model is.
Rotary position embeddings: encoding position as a rotation, not an addition
The original transformer adds a fixed sinusoidal vector to each token's embedding once, at the input layer, so that position information has to survive being mixed through every subsequent layer alongside content information. Rotary Position Embedding, RoPE (Su et al., 2021, "RoFormer: Enhanced Transformer with Rotary Position Embedding"), does something structurally different: instead of adding a position signal at the input, it rotates each query and key vector by an angle proportional to that token's position, immediately before the attention dot product is computed, inside every layer, every time.
Take the simplest case, a 2-dimensional vector (x₁, x₂) at position m, rotated by angle m·θ using the standard 2D rotation matrix: the rotated vector is (x₁cos(mθ) − x₂sin(mθ), x₁sin(mθ) + x₂cos(mθ)). The payoff shows up when you take the dot product of a rotated query at position m with a rotated key at position n: because rotating two vectors by angles α and β and then taking their dot product is equivalent to rotating one of them by the difference (α − β), the resulting attention score depends only on (m − n), the relative distance between the two tokens, never on their absolute positions. Two tokens three positions apart produce the identical raw attention-score contribution from this term whether they are the 3rd and 6th tokens of the sequence or the 3,003rd and 3,006th. This relative-position property is exactly what the original chapters' additive sinusoidal encoding does not guarantee, and it is the reason RoPE extrapolates far better to sequence lengths not seen during training.
import math
def rope_rotate(x1, x2, pos, theta):
angle = pos * theta
cos_a, sin_a = math.cos(angle), math.sin(angle)
return x1 * cos_a - x2 * sin_a, x1 * sin_a + x2 * cos_a
theta = 0.3
q = (1.0, 0.0)
k = (1.0, 0.0)
for m, n in [(2, 5), (10, 13)]:
qm = rope_rotate(*q, m, theta)
kn = rope_rotate(*k, n, theta)
dot = qm[0] * kn[0] + qm[1] * kn[1]
print(f"m={m}, n={n}, m-n={m-n}, dot={dot:.4f}")Trace the first pair by hand: at m = 2, angle = 0.6 rad, so qm = (cos 0.6, sin 0.6) ≈ (0.8253, 0.5646). At n = 5, angle = 1.5 rad, so kn = (cos 1.5, sin 1.5) ≈ (0.0707, 0.9975). Their dot product is 0.8253 × 0.0707 + 0.5646 × 0.9975 ≈ 0.6216. Now the second pair, with the same relative gap of −3 but at completely different absolute positions: at m = 10, angle = 3.0 rad, qm ≈ (−0.9900, 0.1411); at n = 13, angle = 3.9 rad, kn ≈ (−0.7259, −0.6878); dot product ≈ (−0.9900)(−0.7259) + (0.1411)(−0.6878) ≈ 0.6216. Both print as dot=0.6216, confirming the score depends only on m − n = −3, not on the absolute positions 2, 5 or 10, 13.
RMSNorm and SwiGLU: the quiet substitutions inside every block
Two smaller changes appear in essentially every modern open-weights model alongside GQA and RoPE. The first replaces LayerNorm with RMSNorm (Zhang and Sennrich, 2019, "Root Mean Square Layer Normalization"). LayerNorm normalizes a vector by subtracting its mean and dividing by its standard deviation before applying a learned scale and shift. Zhang and Sennrich's ablations found that the mean-subtraction (re-centering) step contributes little to LayerNorm's training-stability benefit; the re-scaling step, dividing by a magnitude statistic, is what actually matters. RMSNorm keeps only that part: RMSNorm(x) = (x / RMS(x)) · g, where RMS(x) = √((1/d)Σxᵢ² + ε) and g is a learned per-dimension gain. Dropping the mean computation removes one full reduction pass over the vector at every normalization call, in every layer, for every token, with no measured quality cost in their experiments, which is a meaningful, compounding saving at the scale of a model with dozens of layers processing billions of tokens.
The second swaps the feed-forward network's ReLU activation for a gated variant, SwiGLU (Shazeer, 2020, "GLU Variants Improve Transformer"). Where a standard transformer FFN computes ReLU(xW₁)W₂ using two weight matrices, SwiGLU computes (SiLU(xW₁) ⊙ (xV))W₂ using three: one gate matrix W₁ passed through SiLU (x·σ(x)), one value matrix V, elementwise-multiplied together, then projected back down by W₂. Because this uses three matrices of the intermediate width instead of two, matching it to a ReLU-FFN's total parameter count requires shrinking the intermediate dimension by a factor of about 2/3. This is exactly why Llama 7B's published intermediate_size is 11,008 rather than the "obvious" 4 × 4096 = 16,384 you would expect from the original paper's 4× expansion rule: (2/3) × 4 × 4096 ≈ 10,923, rounded up to the nearest multiple of 256, giving 11,008.
Assembling the modern block
Put together, a Llama-style decoder block looks like this: x + Attention(RoPE(RMSNorm(x)), using GQA-shared KV heads), followed by x + SwiGLU-FFN(RMSNorm(x)), with normalization applied before each sublayer (pre-norm) rather than after it, which itself improves training stability at depth. Every piece of that sentence is a named, citable substitution for a piece of the 2017 block: RMSNorm for LayerNorm, RoPE for additive sinusoidal encoding, GQA for full multi-head attention, and SwiGLU for the ReLU feed-forward network. None of these changes the shape of what a transformer computes at a diagram level; all of them change, by a large constant factor, what it costs to actually run one.
Active recall
Attempt each question before reading its answer.
Q1. In one sentence, what does the KV cache store, and what does it not eliminate?
Q2. A model has 24 layers, 16 KV heads, head_dim = 64 (so d_model = 1024), stored in fp16. Compute the KV cache size for one 2048-token sequence.
Q3. Using the 7B-class example above (32 layers, 8 KV heads via GQA, head_dim = 128, 60 GiB of headroom on an 80 GB GPU), suppose the deployment doubles maximum context length from 4096 to 8192 tokens. Trace the full effect: (a) new cache size per sequence, (b) new number of concurrent sequences the GPU can hold, (c) how per-token attention cost and total attention cost across a full generation each change, (d) does RoPE require retraining to support the longer context?
Q4. Why does removing mean-centering (RMSNorm vs. LayerNorm) still keep training stable, given that LayerNorm's normalization was long assumed to need mean-subtraction first?
Q5. SwiGLU's FFN uses three weight matrices instead of ReLU-FFN's two. What do production models typically do to the intermediate dimension to keep total FFN parameters roughly matched, and why?
Q6. True or false: KV caching reduces the complexity of generating n tokens from O(n³) to O(n²). Justify your answer precisely.
A1. The KV cache stores the already-computed key and value vectors for every previously generated token, in every layer and head, so they never have to be recomputed. It does not eliminate the attention computation itself: each new token's query still has to be compared against every cached key and weighted-summed against every cached value, so per-step cost still grows with sequence length.
A2. Using 2 × L × H × d_head × T × bytes: 2 × 24 × 16 × 64 × 2048 × 2 = 201,326,592 bytes = 192.0 MiB.
A3. (a) Cache size scales linearly in sequence length: 512 MiB × 2 = 1.0 GiB per 8192-token sequence. (b) 60 GiB of headroom ÷ 1.0 GiB per sequence ≈ 60 concurrent sequences, down from about 120 at 4096 tokens: doubling context length halves batch capacity, all else equal. (c) Per-token attention cost at the end of an 8192-token generation is about double what it was at the end of a 4096-token one (proportional to current cache length), but total attention compute integrated across the whole generation scales as O(n²), so doubling the full context length roughly quadruples total attention compute for filling the sequence, not merely doubles it; the "only 2×" answer is the one most students give and it is incomplete. (d) No retraining is strictly required, since RoPE's rotation angles are a deterministic function of position rather than learned parameters, but naive extrapolation past the training-time maximum length typically degrades quality; production systems instead apply a scaling correction such as positional interpolation (Chen et al., Meta AI, 2023, "Extending Context Window of Large Language Models via Positional Interpolation") to remap the rotation angles before extending the usable context.
A4. Zhang and Sennrich's ablations showed empirically that LayerNorm's re-centering (mean-subtraction) invariance contributes little to its stabilizing effect; the re-scaling invariance, normalizing by a magnitude statistic so gradients do not explode or vanish across layers, is what actually matters. RMSNorm keeps only the re-scaling step and skips computing the mean, which removes one reduction pass per normalization call without a measurable quality cost in their experiments.
A5. They shrink the FFN's intermediate dimension to roughly 2/3 of the original 4× expansion, since three matrices at the original width would otherwise cost about 1.5× the parameters and compute of the two-matrix ReLU baseline; scaling the width down by 2/3 keeps the total roughly matched. Llama 7B's intermediate_size of 11,008 is (2/3) × 4 × 4096 ≈ 10,923, rounded up to the nearest multiple of 256.
A6. True, but only relative to the specific naive baseline this chapter defined: recomputing the full forward pass over the entire sequence-so-far at every generation step, which costs roughly O(n³) summed across n steps because of the redundant re-derivation of unchanged past tokens. The KV cache removes that redundant recomputation, leaving the inherently quadratic O(n²) cost of attending over a growing cache, which caching does not and cannot remove. It is not correct to say caching makes generation "linear"; it is correct to say it removes one whole power of n from a needlessly wasteful implementation.
Think About It
Think about this: How would you explain building the transformer: the architecture that changed ai to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.
Practice Exercises
Now it is time to practice! Complete these challenges to solidify your understanding:
- Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
- Exercise 2: Find a real-world example where building the transformer: the architecture that changed ai is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
- Exercise 3: Create a mind-map connecting building the transformer: the architecture that changed ai to at least 3 other topics you have studied.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind building the transformer: the architecture that changed ai, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.