A telecom-support assistant built on a 70-billion-parameter model goes live for a carrier with 400 million subscribers. On the day the postpaid bills generate, complaint volume spikes and the assistant is fielding 30,000 simultaneous chats. The model itself has not changed since the demo that impressed everyone last month. What changes is that every one of those 30,000 conversations forces the GPU fleet to keep a private, growing scratchpad of every token it has already read, for the entire length of that conversation, for every layer, for every attention head. The engineering team's pager does not go off because attention cannot be computed fast enough — it goes off because the GPUs run out of memory bandwidth moving numbers around before they can even start computing. This chapter is about that scratchpad: what it costs, why it is the actual bottleneck in serving a transformer at scale, and the two families of fixes — grouped-query attention and FlashAttention — that let the same architecture serve orders of magnitude more users on the same silicon.
What this chapter assumes, and what it adds
You already know how a single attention head computes softmax(QKᵀ/√d)V, why the scaling factor exists, how multiple heads are stacked, and how positional information and residual connections hold the whole encoder or decoder together. None of that is repeated here. What follows is the part of the transformer story that research papers on architecture skip and production systems papers live in: what happens to attention the moment you stop training on fixed batches and start generating text one token at a time, forever, for paying users, on a GPU that costs real money per hour.
Generation is not training: the KV cache
During training, a transformer sees a full sequence at once and computes attention over it in one pass — this is why training is compute-bound and GPUs (which are extremely good at dense matrix multiplication) are the right tool. Autoregressive generation is structurally different. To produce token t+1, the model needs the key and value vectors of every token from 1 to t, for every layer and every head, so that the new query at position t+1 can attend over all of them. Recomputing those keys and values from scratch at every new token would mean redoing the entire forward pass for the whole prefix on every single step — quadratic waste for no reason, since the key and value vectors of past tokens never change once computed.
The fix, used in every production LLM serving stack, is the KV cache: after each forward pass, the key and value tensors for the newly processed token are stored, per layer, per head. Generating the next token only requires computing Q, K, V for the one new token, then attending that single query against the cached keys and values of everything before it. The cost of one decode step becomes roughly linear in the sequence length already generated, not quadratic in the whole thing recomputed. The price for this saving is memory: the cache is a live tensor sitting in GPU high-bandwidth memory (HBM) for as long as that conversation is open, and it grows by one token's worth of K and V at every step.
The size of that cache, in bytes, is exactly:
KV_cache_bytes = 2 × n_layers × n_kv_heads × head_dim × seq_len × batch × bytes_per_valueThe leading 2 is for storing both K and V. n_kv_heads is the number of attention heads that maintain their own key/value projections — in standard multi-head attention (MHA) this equals the number of query heads, but as the next section shows, it does not have to.
Worked Example 1 — sizing the cache for a real 70-billion-parameter model
Llama 2's 70B configuration (Touvron et al., 2023, Meta AI) has n_layers = 80, n_heads = 64 query heads, head_dim = 128, and — this is the detail that matters for this chapter — only n_kv_heads = 8, because the 34B and 70B variants use grouped-query attention rather than standard MHA (the mechanism is explained below). Take a single 4,096-token conversation, batch size 1, weights and activations stored in fp16 (2 bytes per value).
def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, batch=1, bytes_per_value=2):
return 2 * n_layers * n_kv_heads * head_dim * seq_len * batch * bytes_per_value
mha_equivalent = kv_cache_bytes(80, 64, 128, 4096) # if it used standard MHA
gqa_actual = kv_cache_bytes(80, 8, 128, 4096) # Llama-2-70B's real config
mqa_hypothetical = kv_cache_bytes(80, 1, 128, 4096) # single shared KV head
for name, b in [("MHA-equivalent", mha_equivalent),
("GQA (actual)", gqa_actual),
("MQA", mqa_hypothetical)]:
print(f"{name}: {b/1e9:.2f} GB")Tracing the arithmetic by hand for the GQA row: 2 × 80 × 8 × 128 × 4096 × 1 × 2 = 1,342,177,280 bytes, i.e. 1.34 GB. Running the same formula for the other two configurations, the code above prints:
MHA-equivalent: 10.74 GB
GQA (actual): 1.34 GB
MQA: 0.17 GBRead that middle row again: a single 4,096-token conversation costs 1.34 GB of GPU memory just to keep the conversation's own state alive — before a single weight is loaded, before a second user connects. If the model had used standard MHA with 64 independent KV heads instead of grouping them into 8, the same conversation would cost 10.74 GB, eight times as much, for identical model quality on the query side. On a GPU with roughly 40 GB free for KV cache after the model weights are resident, that is the difference between serving about 3 concurrent 4,096-token conversations and serving about 29 of them — 40 / 1.34 ≈ 29.8 versus 40 / 10.74 ≈ 3.7. That roughly 8× gap in concurrent-user capacity, for the same GPU fleet, is the entire commercial argument for grouped-query attention.
Why decoding starves the GPU: the roofline argument
It is tempting to think the bottleneck above is just "attention is a big matrix multiply, and big matrix multiplies are slow." That is the wrong mental model for decoding, and naming it precisely matters. During single-token decoding, the query is one vector, not a matrix — there is very little arithmetic to do (a handful of dot products against the cached keys, then a weighted sum of the cached values). What dominates the time is moving the entire KV cache from HBM into the GPU's on-chip compute units, once per generated token, forever.
The standard way engineers reason about this is arithmetic intensity: FLOPs performed per byte moved from memory. Per token, per layer, decode-step attention performs roughly 4 × d_model floating-point operations (two multiply-adds for the query-key scores, two more for the weighted sum over values), while it must read 2 × n_kv_heads × head_dim × bytes_per_value bytes of cached K and V for that token. Dividing:
arithmetic_intensity = (4 × d_model) / (2 × n_kv_heads × head_dim × bytes_per_value)With bytes_per_value = 2 (fp16), the 4 in the numerator and the 2×2=4 in the denominator cancel, leaving the clean form arithmetic_intensity = d_model / (n_kv_heads × head_dim).
For Llama-2-70B, d_model = 8192, head_dim = 128. With standard MHA (n_kv_heads = 64): intensity = 1.0 FLOP per byte. With the actual GQA config (n_kv_heads = 8): intensity = 8.0 FLOPs per byte. An A100 GPU delivers roughly 312 TFLOP/s of fp16 tensor-core compute and roughly 2.0 TB/s of HBM bandwidth, giving a ridge point — the intensity at which compute and memory bandwidth are equally limiting — of 312e12 / 2.0e12 = 156 FLOPs per byte. Both 1.0 and 8.0 sit far below 156, which is the formal statement of "decoding is memory-bound": the GPU's compute units spend almost all their time idle, waiting for bytes to arrive. Grouped-query attention does not change that qualitative fact — decoding stays memory-bound either way — but it moves the achieved intensity from 0.6% of the ridge point to 5.1% of it, meaning eight times fewer bytes are needed per token generated, which is exactly the 8× cache-size and concurrency gap measured above. The FLOPs were never the problem; the bytes were.
Shrinking the cache without shrinking the model: MQA and GQA
Multi-query attention, introduced by Noam Shazeer in Fast Transformer Decoding: One Write-Head Is All You Need (2019), keeps all query heads exactly as they are — same number, same dimension, same learned projections — but has every query head read from a single shared key/value projection instead of its own. Grouped-query attention (Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón & Sanghai, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023) is the middle ground: query heads are split into g groups, and each group shares one KV head, interpolating between full MHA (g = number of heads, no sharing) and MQA (g = 1, total sharing). Llama-2-70B's 64 query heads in groups of 8, sharing 8 KV heads, is a GQA model with a group size of 8.
The mechanism is a change to which tensors get cached, not to how attention scores are computed. Each query head still produces its own attention distribution over the sequence and its own output — the QKᵀ and softmax steps are unchanged per head. What is shared is the K and V projection weight matrices themselves, so heads within a group are scoring against and reading from an identical set of key/value vectors rather than each maintaining a private copy. Fewer distinct K/V tensors means a proportionally smaller cache and proportionally fewer bytes to move per decode step, exactly matching the arithmetic in the previous two sections.
Correcting a misconception
The instinct, on first meeting GQA, is that sharing key/value projections across heads must throw away information and make the model weaker — heads are supposed to specialize, and specialization implies they need different views of the sequence. This conflates two things that GQA deliberately keeps separate: what a head looks for (its query projection, fully independent per head, unchanged) versus what a head reads from (its key/value projection, shared within a group). A head's specialization lives primarily in how it forms its query and interprets the resulting attention weights, not in maintaining a private key/value dictionary. Ainslie et al. (2023) uptrained existing multi-head checkpoints (including T5-XXL) into grouped-query models with a small fraction of the original pre-training compute and found quality within roughly a percentage point of full MHA on summarization and translation benchmarks, while cutting inference latency to close to that of MQA. The cache-size arithmetic in the previous sections is exact and unconditional; the quality cost, empirically, is close to free. The one real trade Meta made explicit in the Llama 2 paper: they applied GQA only to the 34B and 70B models, where the KV-cache-memory problem is largest, and kept standard MHA on the smaller 7B and 13B models where it was not yet the bottleneck — a sign that this is a deployment-scale decision, not a universal upgrade.
FlashAttention: attention that respects the memory hierarchy
GQA attacks the size of the cache. FlashAttention (Dao, Fu, Ermon, Rudra & Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022) attacks a different inefficiency: the way the standard attention computation is implemented moves far more data through HBM than the arithmetic strictly requires, even before the KV cache enters the picture. A naive implementation computes the full N × N score matrix QKᵀ, writes it to HBM, reads it back to apply the softmax, writes the result back, reads it again to multiply by V. On an A100, HBM bandwidth is roughly 1.5–2.0 TB/s, while the on-chip SRAM each streaming multiprocessor uses for its working set moves data at roughly 19 TB/s — nearly an order of magnitude faster, but with capacity measured in a few hundred kilobytes rather than tens of gigabytes. Every trip of that N × N matrix out to HBM and back is a trip at the slow, high-capacity tier when it could have stayed at the fast, low-capacity tier if the computation were restructured to fit.
FlashAttention restructures it by tiling: it loads a block of queries and a block of keys/values into SRAM, computes the partial attention scores and output for just that block, and never materializes the full N × N matrix in HBM at all. The obstacle to doing this naively is that softmax needs a normalization constant (the sum of exponentials) computed over the entire row before any single weight can be finalized — which seems to force seeing all the keys before finishing any block. The trick that removes this obstacle is the online softmax (its numerically-stable running form traces back to Milakov & Gimelshein, 2018): process one block, get a provisional output and a provisional normalizer, then when the next block arrives, rescale everything computed so far by a correction factor and fold in the new block's contribution. The result, after all blocks are seen, is mathematically identical to computing the softmax over the full row at once — not an approximation, an algebraic rearrangement.
Worked Example 2 — tracing the online-softmax update by hand
Take one query q = [1, 1] attending over four keys, split into two blocks of two (a tiny stand-in for a much larger SRAM tile), with scores computed as q·k / √d, d = 2:
k1=[1,0], v1=[10,0] k2=[0,1], v2=[0,10] ← block 1
k3=[2,2], v3=[5,5] k4=[-1,-1], v4=[-5,0] ← block 2
s1 = (1)/√2 = 0.7071 s2 = (1)/√2 = 0.7071
s3 = (4)/√2 = 2.8284 s4 = (-2)/√2 = -1.4142Computing the ordinary full softmax over all four scores at once, as a reference to check against: weights [0.0956, 0.0956, 0.7974, 0.0115], giving output w·v = [4.8854, 4.9427].
Now trace the block-wise online update, which never sees blocks 1 and 2 at the same time. Maintain a running max m, running normalizer l, and running unnormalized output acc, initialized to (-∞, 0, [0,0]):
import math
def online_softmax_block(state, scores, values):
m, l, acc = state
m_new = max(m, max(scores))
correction = math.exp(m - m_new) if m != -math.inf else 0.0
l_new = l * correction + sum(math.exp(s - m_new) for s in scores)
dim = len(values[0])
acc_new = [acc[d]*correction +
sum(math.exp(s - m_new) * v[d] for s, v in zip(scores, values))
for d in range(dim)]
return (m_new, l_new, acc_new)Block 1 (scores = [0.7071, 0.7071]): the running max was -∞, so it becomes m = 0.7071, correction is defined as 0 (nothing to rescale — no prior block exists). l = e⁰ + e⁰ = 2.0000. acc = 1×[10,0] + 1×[0,10] = [10, 10]. If generation stopped here, the (wrong, partial) output would be acc/l = [5, 5] — a plausible-looking number that is silently missing block 2 entirely, which is exactly why the running state, not just the ratio, has to be carried forward.
Block 2 (scores = [2.8284, -1.4142]): the block's own max, 2.8284, exceeds the running max 0.7071, so m_new = 2.8284 and every quantity accumulated in block 1 is now stale relative to this new max — it must be rescaled by correction = exp(0.7071 - 2.8284) = exp(-2.1213) = 0.1199 before block 2's contribution is added. l = 2.0000 × 0.1199 + (e⁰ + e^-4.2426) = 0.2398 + 1.0143 = 1.2541. acc = [10,10] × 0.1199 + [5,5] + [-5,0] × e^-4.2426 ≈ [1.199, 1.199] + [5,5] + [-0.071, 0] = [6.1269, 6.1987].
Final output: acc / l = [6.1269/1.2541, 6.1987/1.2541] = [4.8854, 4.9427] — matching the full-softmax reference to four decimal places, computed without ever holding all four scores or writing an N × N matrix anywhere. The correction factor is what makes this exact rather than approximate: every time the running max updates, everything accumulated so far is retroactively rescaled to be consistent with the new max, so the final ratio is identical to what a single global softmax would have produced. This is the computation GPU kernels perform block-by-block inside fast SRAM, writing only the small running state back to HBM between blocks instead of the full score matrix.
How the two mechanisms compose
GQA and FlashAttention solve different layers of the same problem and are used together, not as alternatives. GQA reduces how many bytes exist to move — fewer distinct KV heads means a smaller cache in absolute terms. FlashAttention reduces how wastefully those bytes are moved during the score-and-softmax computation, whether the cache is large or small. A production serving stack — vLLM, TensorRT-LLM, and similar systems all do this — runs GQA-shaped models through FlashAttention-style fused kernels, so the cache is already smaller thanks to grouping, and what remains is read with close to the minimum possible HBM traffic. Neither technique changes the number of parameters in the model or what the attention mechanism conceptually computes; both are proofs that the same mathematical object can be implemented at wildly different hardware efficiencies, and that the gap between those implementations is where most of the cost of running an LLM at scale actually lives.
Diagram: from query-head grouping to cache size
Active recall
Q1. Why is single-token decoding memory-bandwidth-bound rather than compute-bound, even though attention is fundamentally a matrix multiplication?
Q2. A classmate argues: "GQA must produce worse answers, because it throws away information by forcing heads to share." What is wrong with this reasoning?
Q3. A different decoder has 40 layers, 32 query heads, head_dim 128, and uses GQA with 4 KV heads. For a 2,048-token sequence, batch size 1, fp16, what is the KV cache size in GB?
Q4. Starting from the Llama-2-70B GQA configuration in Worked Example 1 (8 KV heads, 4,096-token sessions, 1.34 GB each, ~29 concurrent sessions in a 40 GB budget): the serving team doubles the context window to 8,192 tokens and switches to a speculative-decoding setup that keeps 2 sequences (draft + verify) per user. What happens to the cache size per user, and how many concurrent users now fit in the same 40 GB budget?
Q5. In the online-softmax trace, block 2's correction factor was exp(m_old − m_new) = exp(0.7071 − 2.8284) = 0.1199. Why is the exponent m_old − m_new and not the other way around, and what would go wrong numerically if it were computed as exp(m_new − m_old) instead?
Q6. FlashAttention is described as computing exact attention, not an approximation. Using the worked trace as evidence, explain why tiling plus the running-correction rescale gives a result identical to full softmax rather than merely close to it.
Answers
A1. Decoding one token performs very little arithmetic per layer (roughly 4 × d_model FLOPs against the cached keys and values for that step) but must read the entire per-layer KV cache — potentially thousands of tokens' worth of stored keys and values — from HBM to do it. The arithmetic intensity (FLOPs per byte moved) works out to about 1.0 for standard MHA on a Llama-2-70B-sized model, far below the ~156 FLOPs/byte ridge point of an A100 where compute and bandwidth are balanced. Below the ridge point, the GPU's compute units sit idle waiting for memory traffic, which is the definition of memory-bound.
A2. The argument conflates two separate things GQA keeps distinct: each query head's projection (what it looks for) stays fully independent and undiminished; only the key/value projections (what gets read) are shared within a group. A head's specialization is expressed mainly through its query, not through owning a private KV dictionary. Ainslie et al. (2023) measured quality within about a percentage point of full MHA after uptraining, which is consistent with the mechanism, not a coincidence.
A3. 2 × 40 × 4 × 128 × 2048 × 1 × 2 = 167,772,160 bytes ≈ 0.168 GB (168 MB). Note that the number of query heads (32) never enters the formula — only n_kv_heads, head_dim, seq_len, and n_layers determine cache size.
A4. Both changes multiply cache size linearly and independently, so they compound: doubling seq_len from 4,096 to 8,192 doubles the cache to 2.68 GB per single sequence; doubling the number of live sequences per user (draft + verify) doubles it again, to 5.37 GB per user — four times the original 1.34 GB, not two. In the 40 GB budget, concurrent-user capacity falls from about 29 to 40 / 5.37 ≈ 7 — a roughly 4× drop, driven equally by the context-length doubling and the per-user sequence-count doubling; missing either factor when estimating capacity would overstate it by 2×.
A5. The running max only ever increases as new blocks arrive (m_new = max(m_old, block_max) ≥ m_old), so m_old − m_new ≤ 0 and exp(m_old − m_new) is always a value in (0, 1] — a safe, bounded shrink factor applied to the stale accumulator so it is expressed relative to the new, larger max, exactly mirroring how the numerically-stable softmax subtracts the max before exponentiating. Computing it backwards as exp(m_new − m_old) would give a factor ≥ 1, growing the already-stale accumulated sum and output by the wrong direction — in this trace, exp(2.8284 − 0.7071) ≈ 8.34 instead of 0.1199, corrupting l and acc and, for a longer sequence, being exactly the kind of unbounded growth the max-subtraction trick exists to prevent (overflow for large score ranges).
A6. The trace showed the two-block online computation land on [4.8854, 4.9427], identical to the single-pass full-softmax reference to four decimal places. That equality is not a numerical coincidence of this particular example — the correction step at every block boundary exactly rescales every previously accumulated quantity to be consistent with the new global maximum seen so far, so by the final block, acc and l hold precisely the same sums a single global softmax would have produced, just built up incrementally. Tiling changes the order and locality of the computation (and therefore its speed and memory traffic); it does not change which numbers get summed or how they get weighted, which is why "IO-aware" in the paper's title is paired with "exact," not "approximate."
Think About It
Think about this: How would you explain transformers and attention mechanisms: the foundation of modern llms 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 transformers and attention mechanisms: the foundation of modern llms 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 transformers and attention mechanisms: the foundation of modern llms to at least 3 other topics you have studied.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind transformers and attention mechanisms: the foundation of modern llms, 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.