Attention, Worked Out by Hand (Part 2): MHA, GQA, MQA, FlashAttention & MLA

Part 2 of the attention series · Part 1: Attention, Step by Step

Most comparisons of these variants are a table of asymptotics. This one is arithmetic: the same tokens, run through each variant, with the intermediate numbers on the page. Every table below was computed rather than illustrated.

The five-step mechanism itself is worked out slowly, on two tokens, in Attention, Step by Step. This article assumes it and goes after the bill: what each variant has to keep in memory per token while it generates. That per-token number is what decides how many users fit on one GPU at once, which is to say what a served token costs, and by the end a config line like num_kv_heads: 8 should read as a decision rather than an incantation.

The recap, and no more than it

Every word rewrites itself as a weighted blend of the words it can see.

Part 1 works that out digit by digit; here is the sentence-length version. Each word turns its vector into a query (what it is looking for), a key (what it advertises) and a value (what it hands over if attended to). A score is one dot product, query against key, and five steps (score, scale, mask, softmax, mix) turn a row of scores into weights that add up to 1, which each word spends on the values. The arrows only ever go backwards, or curl round to the word itself: that is causal masking, a word may look at what came before it and never at what comes after.

Every variant below keeps that mechanism intact, and none of them changes the five steps. What they change is how many keys and values there are to pay out from, and where those numbers live while it happens.

Why there is more than one kind of attention

Part 1's example was two tokens; here it widens to four, the cat sat on, because the costs below only show up once there is more than one thing in the cache. The arithmetic is the five steps above; the numbers are new, made the same way: each token's vector through W_Q, W_K and W_V. You are not expected to read these tables; they are here so every number below has somewhere to come from. The row to watch is on's.

Run the same five steps on the on row and you get the baseline that every variant below is measured against:

Keep that row in mind. It is about to be computed three more ways.

The hatched cells are the mask. Rows 3 and 4 show the softmax worked out: exponentiate each live score, divide by the row's total. Row 4 is the on row above; this is where it comes from.

Everything from here on computes something very close to what you just watched. None of it is a cleverer idea about what attention should mean; it is all about what attention costs you, and the costs are specific and annoying:

Three fixes for five names is the right count: MHA is the thing being billed, and MQA turns up inside GQA's section on its own.

The first one needs a bill before it makes sense, and the bill needs a head count. So: heads first, then the bill.

First, there is more than one head

Everything computed on this page is a single head. Real models run a crowd of them in parallel: each head has its own W_Q, W_K and W_V, reads the whole token vector, and writes a head_dim-wide slice of the output. A head is the same five steps run on a thinner slice: four heads on quarter-width vectors do the same total arithmetic as one head at full width, so the split is free in FLOPs, and it buys several independent questions per layer, because one head can only ask one kind of question.

That matters here for one reason: the cache scales with the head count. The Q, K and V tables above are one head of head_dim 4. A four-head layer of the same shape would be four such blocks side by side, d_model 16. The cache figures below count that four-head layer, because a one-head cache is not a problem anyone has.

The bill

Generating text is stubbornly sequential. You produce a token, feed it back in, produce the next one. Each new token arrives with a fresh query, and that query has to score against every key so far and collect from every value so far; that is all attention is. So the old tokens' queries are spent the moment they're used, but their keys and values are needed again at every single step from now on. Recomputing them for the whole prompt every time would be quadratic and absurd, so every serving stack keeps them in a cache. Per token, per layer, that cache costs 2 × n_heads × head_dim numbers.

One head of head_dim 4 caches 4 numbers of K plus 4 of V: 8 numbers. A four-head layer of that shape caches 2 × 4 × 4 = 32. The middle 4 is the head count, not the width: each head keeps its own K and V.

Llama 3 70B has 64 heads of head_dim 128. Full multi-head attention would be 2 × 64 × 128 = 16,384 numbers per token per layer, times 80 layers, times every token in every sequence you're serving concurrently. At 8k context that's gigabytes per user. The weights are a fixed cost you pay once; the KV cache scales with users × context, and it is what decides your batch size, and with it your cost per token.

That number, not the weights, is what every variant below is trying to shrink. (Quantising it rather than shrinking it is the other lever. That is TurboQuant, and it composes with everything here.)

1. Grouped-query attention: share the keys

GQA gives every head its own queries but makes several heads share one K/V pair. In a four-head layer with two KV heads, heads 0 and 1 read group 0's keys and heads 2 and 3 read group 1's. The cache term goes from 2 × n_heads × head_dim to 2 × n_kv × head_dim.

Which raises the obvious worry: if two heads are reading the same keys, don't they just collapse into the same head?

No. They still ask different questions. Here is our head, and a second head with its own W_Q reading the identical K and V, both on the on row:

Same keys, same values, opposite answer. Head 0 asks what on attaches to and finds the verb. Head 1 asks what frames the phrase and finds the determiner. Sharing K and V constrains what can be looked up, not what gets asked, and that difference is the whole reason GQA works.

Notice what that table already is. Two heads, each with its own W_Q, reading one identical K/V set. That is multi-query attention: GQA taken to its limit, with a single KV head for every query head in the layer. The rows above are literally an MQA layer.

The bill

In the wild: Llama 3 70B runs 64 query heads over 8 KV heads, which is 2,048 numbers per token per layer instead of 16,384. An 8× cut for a quality loss that benchmarks put somewhere around nothing. MQA at the same dimensions caches 2 × 1 × 128 = 256, a 64× cut, and is measurably worse.

Two things decide where you land between them. MQA is unstable to train, which is the reason GQA was proposed at all. And big models are served sliced across several GPUs (tensor parallelism), with the KV cache divided up by head, so each GPU holds its own heads' keys and values. Drop to a single KV head and every GPU needs a full copy of it: the saving stops once n_kv reaches the GPU count. That is why production configs set the two equal rather than driving to one: Llama 3 70B runs 8 KV heads on 8-way tensor parallelism.

GQA changes the model, though not necessarily your training budget. You can average an existing MHA checkpoint's K/V heads into their groups (mean-pool them, in the paper's terms), then keep training briefly so the model recovers from the surgery, at roughly 5% of the original compute. That recipe, which the paper calls uptraining, is its actual contribution. The next one needs no training at all.

2. FlashAttention: same numbers, different schedule

This one is different in kind: it computes exactly the result from the five steps above. It is a memory-access fix, not an architecture change.

Here's the problem it fixes. The naive kernel builds the whole T×T score matrix in HBM, which is the GPU's big slow memory. It writes the thing out, reads it back for softmax, writes the result, then reads that again for @V. At T = 8,192 that matrix holds 67 million entries per head. The arithmetic is trivial. The round trips are not. Attention is memory-bound, not compute-bound.

Don't build it. Stream K and V through SRAM in small tiles (SRAM being tiny, and roughly 10× the bandwidth), and keep a running softmax as you go. Whenever a new tile turns up with a bigger max, rescale everything you've accumulated so far and carry on.

Three running numbers do the work, and every exponent is taken relative to the running max: p = exp(score - m), which is what keeps e^x from overflowing. m is the biggest score seen so far, l is the running sum of those exponentials, and acc is the running weighted sum of value vectors. Watch it on our on row, tile size 2. Its masked scores were [-0.238, -0.321, 1.200, -0.478], and softmaxing those is what produced the baseline row above.

Tile 0, keys the and cat. Scores [-0.238, -0.321]. (Scores are printed to three places throughout; the arithmetic below runs on the unrounded values, so the last digit of a product will not always match what you get from the printed operands.)

# this tile: keys "the", "cat"
scores = [-0.238, -0.321]
# running max; -0.238 is this tile's biggest
m = max(-inf, -0.238)  = -0.238
# e^0 and e^(-0.321 + 0.238)
p = exp(scores - m)    = [1.0000, 0.9197]
# running sum of p
l = 1.0000 + 0.9197    =  1.9197
acc = 1.0000·V[the] + 0.9197·V[cat]
    = [0.619, -0.728, -0.456, 0.565]

Tile 1, keys sat and on. Scores [1.200, -0.478]. The max jumps from −0.238 to 1.200, so everything accumulated so far is stale by e^(-0.238 - 1.200) ≈ e^-1.4372 = 0.2376.

# this tile: keys "sat", "on"
scores  = [1.200, -0.478]
# the max moved, so the old work is stale
m       = max(-0.238, 1.200)  =  1.200
# = e^(-0.238 - 1.200)
rescale = exp(old m - new m)  =  0.2376
# e^0 and e^(-0.478 - 1.200)
p       = exp(scores - m)     = [1.0000, 0.1869]
l = 1.9197 × 0.2376 + 1.0000 + 0.1869
  =  1.6430
acc = [0.619, -0.728, -0.456, 0.565] × 0.2376
    + 1.0000·V[sat] + 0.1869·V[on]
    = [-0.106, 1.036, -0.268, -0.457]

Finish. Divide by l:

acc / l  =  [-0.064, 0.631, -0.163, -0.278]
# ← from the table further up
out[on]  =  [-0.064, 0.631, -0.163, -0.278]

The same numbers, to 1e-12. And the T×T matrix never existed.

Here T is the sequence length, d the head width, and M the SRAM budget the tiles live in. On real hardware d²/M comes out comfortably below 1, so the second row is the smaller one, a fact the notation goes out of its way not to show.

The compute is the same, and arguably slightly more, since the backward pass recomputes the scores rather than storing them. Redoing the maths is cheaper than the trip to memory and back. That tells you everything about how lopsided modern hardware has become.

• The running-max trick isn't new. It's the standard numerically-stable softmax, applied incrementally.
• Causal masking makes it faster, not slower: whole tiles above the diagonal are simply skipped.
• It composes with everything else here. GQA + Flash and MLA + Flash are both normal.

3. Multi-head latent attention: cache one small vector

GQA shrank the cache by making heads share keys. MLA, which is what DeepSeek-V2 and V3 run, asks a cheekier question: why cache the keys at all?

Squash each token down into one small latent vector c. Cache only that. Then rebuild each head's K and V from it on demand, or better, never rebuild them at all, which is where the trick really pays. Whether you get that second option depends on RoPE, the way most models encode position; that fight gets its own section below. Either way the cache stops scaling with the number of heads.

Our four tokens, latent width 2:

Two numbers per token, down from the eight this one head's K and V would need. To use them, a learned up-projection matrix W_UK (U for up, K for keys) inflates the latent back out. c @ W_UK rebuilds this head's keys:

Close, but plainly not equal, and in one spot not even close. Look at on's third column: 1.560 came back as 0.552. The largest single error across the whole matrix is 1.008, which on numbers this size is enormous. At this point you would be forgiven for assuming the technique is broken.

Now run attention from the latent anyway, keys and values both rebuilt from c.

The largest weight shift anywhere in the matrix is 0.025. The keys are visibly wrong and the routing barely notices.

But the routing is not what the next layer reads. The output is. Take the the row alongside it. It is the first token, so the mask leaves it attending 100% to itself, [1, 0, 0, 0], in both runs. Its routing cannot be wrong at all:

Identical routing on the, and the output still moves by 0.617. None of that is where the word looked; all of it is the rebuilt value. Across the whole matrix, the largest output shift (0.617) is 24× the largest weight shift (0.025). Even on the on row, where the weights move most, the output error (0.070) is 2.7× the weight error.

So the honest statement is narrower than "attention barely notices". Where each word looks survives the squeeze almost intact; what it collects does not. Value reconstruction is the loss channel MLA has to buy back during training, and this toy overstates it: a rank-2-of-4 latent is a 50% squeeze, against DeepSeek-V3's 512-of-32,768 with sixty more layers to absorb the damage.

Before you believe the table. Three things it proves less than it seems to.

First, our down-projection comes from the SVD of the stacked token vectors, the very matrix being compressed. The SVD is the textbook answer to "what two directions lose the least when you throw the rest away": the provably optimal rank-2 squeeze, chosen with hindsight. So "what survives the squeeze carried the most variance" is a definition of the SVD here, not a discovery about MLA. The real wager is that a learned low-rank KV subspace still spans what the queries actually read, and the evidence for that is DeepSeek-V2 and V3 shipping on it, not this four-token table.

Second, the latent above is 2 numbers reconstructing the 8 that one head's K and V need, a 4× squeeze, and that is the whole of what we computed. The 2 against 32 in the tables below is a bigger claim: it assumes the same 2-number latent serves all four heads, reconstructing 32 numbers rather than 8. That is exactly what real MLA does, since one latent carries K and V for every head, which is why its cache stops scaling with head count at all. But it is an extrapolation from this page, not a result on it.

Third, at real dimensions the squeeze is far harsher again: DeepSeek-V3 compresses to 512 against 2 × 128 × 128 = 32,768, everything confined to a 512-dimensional subspace. That is a rank constraint on what the layer's keys can express, not a per-key error that averages out: every key must now be mixed from the same few ingredient directions, and whatever a key needed outside them is simply gone. Whether that constraint costs you anything is an empirical question, and DeepSeek's answer is that it doesn't.

RoPE does not fit. Here is the complication promised earlier. RoPE, rotary position embedding, is how most models encode position: before the dot product, each token's q and k are rotated by an angle proportional to that token's position, so the score qᵀk sees position only through the gap between the two tokens.

The cache saving above is only half the trick. Without RoPE the score is qᵀk = qᵀ W_UKᵀ c, and matrix products associate: fold W_UK into W_Q once, offline, and at inference a per-head key is never materialised at all; the query scores straight against the latent.

RoPE breaks the fold. The rotation depends on each token's position, so it has to happen between W_Q and W_UK, exactly where the fold was, and no position-independent product survives. You can still get the right answer by up-projecting every cached key and rotating it, but you are now doing that at every decoding step. The cache is still only the latent; the compute is no longer free. DeepSeek's fix is to split the head: a content half rebuilt from the latent, plus a small decoupled key computed straight from the token and rotated normally. That decoupled key is cached alongside the latent: the +64 below.

The bill

Two numbers per token in the toy, against MHA's 32. At DeepSeek-V3's dimensions it is 576 (512 latent + 64 rotary) against 2 · 128 · 128 = 32,768, the MHA baseline the papers quote. That baseline actually undersells the win: a DeepSeek built with literal MHA would cache more than that, because its keys and values are not the same width: 192 dims of key per head (128 content + 64 rotary) against 128 of value, so its cache is n_h · (d_k + d_v) = 128 × 320 = 40,960. Full comparison in the cheat sheet below.

Break-even. MLA's cache is a fixed width; everything else grows with the head count, so there is a crossover, and it sits low. Writing d_c for the latent width, d_rope for the decoupled rotary key width and d_h for head_dim, as the papers do: against MHA the crossover sits at (d_c + d_rope) / (2 · d_h) heads, or 2.25 at DeepSeek-V3's dimensions on the papers' own baseline. (Against the literal DeepSeek-shaped MHA above, where K and V are different widths, the divisor is d_k + d_v and the crossover is 576 / 320 = 1.8.)

Against GQA it is the same test on the KV heads: MLA is smaller whenever 2 · n_kv · d_h > d_c + d_rope, so above about 2.25 KV heads. Every GQA config at the scale where the cache is the bottleneck clears that bar: Llama 3 70B's eight KV heads cache 2,048 against MLA's 576. Small models sit below it: Qwen2-1.5B's two 128-dim KV heads cache 512, less than the latent itself, which is the crossover doing exactly what it says.

What makes MLA a large-model technique isn't the crossover, though; it's what sits either side of it. The cache is bought with compute: once W_UK folds into W_Q, the query runs at latent width, so the score dot product is 576 wide instead of 192. MLA trades FLOPs for cache, and that only pays when the cache is the thing hurting you.

Cheat sheet

† MHA, GQA and MQA at Llama 3 70B's shape (64 query heads, head_dim 128, 8 KV heads for GQA); MLA at DeepSeek-V3's (512 latent + 64 rotary). 2 · n_h · d_h assumes K and V are the same width, which holds for Llama-shaped attention but not for DeepSeek's.

So which one do you actually want?

• Training or serving anything at length → FlashAttention. Exact, and it costs a few recomputed FLOPs to save the round trips. There is no argument against it.
• Serving a dense model, long context, batch-limited by KV cache → GQA. Well understood, supported everywhere, 4–8× for almost nothing. Set n_kv to your tensor-parallel degree; under head-sharded KV, going below it buys nothing.
• Cache-bound and willing to pay in quality → MQA. Another 8× under GQA, at a loss you will measure.
• Very many heads and the KV cache still dominates → MLA. Best compression on offer, at the cost of a more complex block and the RoPE split.
• Small model → plain MHA. The rest add machinery to compress a cache that wasn't hurting you in the first place.

Check it yourself

Every table here was generated by a script that self-asserts and refuses to run if a number drifts, the FlashAttention tile trace and the MLA latent included, and part 1's numbers come from its companion. The one number you cannot redo on paper is the latent c itself, which comes out of an SVD.

Everything else is a dot product, a division and a weighted average. Take the on row of Q, dot it with each row of K, halve, mask, exponentiate, divide by the total: 0.145, 0.133, 0.609 and 0.114.

Three of the five variants here are that same arithmetic wearing different clothes. Only the bill changed, and FlashAttention did not even change that.

Frequently Asked Questions

What is the difference between MHA, GQA and MQA?

Only how many key/value sets the layer keeps. MHA gives every head its own K and V. MQA gives the whole layer one K/V set that every query head reads. GQA sits in between: heads are split into groups, and each group shares one K/V set. Every head still has its own queries in all three, so they still ask different questions. Sharing constrains what can be looked up, not what gets asked.

Does FlashAttention change the model's output?

No. It is a memory-access schedule, not an architecture change, and it computes the same result up to summation order: the tiles add the same numbers in a different sequence, which floating point notices only in the last digits. That is the 1e-12 in the tile trace above. Nothing is approximated or dropped, which is why it needs no retraining.

Why does the KV cache matter more than the model weights?

The weights are a fixed cost you pay once. The KV cache scales with users × context length, so it is the term that decides how many requests fit on a GPU at the same time: your batch size, and with it your cost per token. At Llama 3 70B's shape, full multi-head attention would cache 16,384 numbers per token per layer across 80 layers; at 8k context that is gigabytes per user.

When is MLA worth the complexity?

When the head count is high enough and the cache is genuinely the binding constraint. MLA's cache is a fixed width, so it beats GQA above roughly 2.25 KV heads. But it buys that with compute, because the query runs at latent width and the score dot product gets wider. It also forces the RoPE split, which makes the block more complicated. On a small model, whose cache was never the problem, it is machinery for nothing.

Can GQA, FlashAttention and MLA be combined?

FlashAttention composes with both. GQA + Flash and MLA + Flash are ordinary production configurations, because Flash only changes how the same arithmetic is scheduled. GQA and MLA are alternatives rather than a stack: both are answers to the same question of what the layer caches, so a model picks one.

Further reading

• Vaswani et al., Attention Is All You Need (2017)
• Shazeer, Fast Transformer Decoding: One Write-Head Is All You Need (2019), which introduced MQA
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023)
• Dao et al., FlashAttention (2022); Dao, FlashAttention-2 (2023); Shah et al., FlashAttention-3 (2024)
• DeepSeek-AI, DeepSeek-V2 (2024) and DeepSeek-V3 (2024), which introduced MLA
• Attention, Step by Step, the five steps worked out slowly on two tokens

About J33.AI

At J33.AI we size and serve models for teams who have to pay the inference bill, and that work is mostly this article: knowing which number grows with your users, and which architecture makes it smaller. Our AI foundations practice runs from model selection through to deployment, across every industry we serve.