Grouped Query Attention (GQA): Efficient Multi-Head Attention

In May 2023, Google Research published a fix that quietly reshaped how every large language model has been served ever since. The paper wasn't about a smarter attention mechanism, a bigger model, or a new loss function. It was about a memory accounting problem: when a transformer generates text one token at a time, the GPU spends most of its time not computing attention scores, but reading numbers back out of memory that it cached from every previous token. Meta's Llama 2 70B, released the same year, adopted the fix (a 34B variant was trained but withheld from public release, per the Llama 2 paper, over insufficient red-teaming time). Mistral 7B adopted it. Llama 3, Gemini, and most production LLMs since have adopted it. The fix is called Grouped Query Attention (GQA), and understanding it requires understanding exactly what breaks when you scale ordinary multi-head attention from a lecture-slide example to a system serving requests.

The bottleneck: autoregressive decoding is not attention's finest hour

You already know multi-head attention (MHA) from studying transformers: split a token's representation into h heads, project each head's query, key, and value with its own learned matrices, compute scaled dot-product attention independently per head, concatenate, and project back down. That description is correct and complete for a single forward pass over a fixed sequence — the kind you do during training, where the entire input is known and processed in one shot.

Text generation is different. A chatbot serving lakhs of concurrent conversations — think of the backend behind an Indian fintech's UPI support bot, or a Swiggy-scale customer-service assistant — generates one token, appends it to the sequence, and generates the next token conditioned on everything so far. Recomputing keys and values for the entire growing sequence at every single step would be quadratically wasteful, so every serving system caches them: for each layer, for each head, the key and value vectors of every token generated so far are stored in GPU memory. This is the KV cache, and at each decoding step the model only needs to compute Q, K, V for the one new token, then attend over the cached K/V of all prior tokens.

The KV cache solves the recomputation problem, but it creates a new one. Its size is:

KV cache bytes = 2 (K and V) × num_kv_heads × head_dim × bytes_per_element × num_layers × seq_len × batch_size

Every term in that product except num_kv_heads is fixed by the model architecture or the workload. And in ordinary MHA, num_kv_heads equals the full head count h — commonly 32, 64, even 96 in large models. As the batch size (concurrent users) and sequence length (context window, chat history) grow, the KV cache grows linearly with both, and at inference time it routinely exceeds the memory footprint of the model's own weights. Worse, at each decode step the GPU must stream every byte of that cache from memory to compute a single new token — decoding is memory-bandwidth bound, not compute bound, so a bigger cache directly means slower, more expensive generation. This is the real motivation for GQA: not FLOPs, but bytes that have to move.

Two extremes, and why neither survived contact with production

Noam Shazeer's 2019 paper proposed the aggressive fix: Multi-Query Attention (MQA). Keep h separate query heads, but give them all a single shared key head and a single shared value head. The K/V cache shrinks by a factor of h — for a 64-head model, a 64× reduction. Decoding gets dramatically cheaper. But a single K/V subspace has to simultaneously serve every query head's very different "questions" about the sequence, and empirically this degrades quality and destabilizes training, especially as models scale up. MQA works, but you pay for the memory savings in output quality.

Plain MHA sits at the other extreme: every query head gets its own private key/value head, maximum representational flexibility, maximum KV cache. GQA, introduced by Ainslie et al. (Google, 2023), is the deliberate middle: partition the h query heads into G groups. Within a group, all query heads share one key head and one value head, but different groups get different key/value heads. Set G = h and you recover ordinary MHA exactly. Set G = 1 and you recover MQA exactly. GQA is not a third mechanism bolted on top of attention — it is the same scaled dot-product attention formula, with the number of independent K/V projections turned into a tunable knob between the two extremes.

Llama-2-70B uses h = 64 query heads grouped into G = 8 groups, so each group has 8 query heads sharing one K/V head — an 8× reduction in KV cache with quality far closer to full MHA than to MQA. The paper also showed a practical trick for adoption: you don't need to train a GQA model from scratch. Take an existing MHA checkpoint, mean-pool each group's original per-head K and V projection matrices into a single shared projection, and continue training ("uptraining") for a small fraction of the original compute budget. The model recovers almost all of its original quality while permanently shrinking its serving cost.

Worked example 1 — sizing the KV cache for Llama-2-70B

Let's put real numbers through the formula above, using Llama-2-70B's published architecture: d_model = 8192, 64 query heads, so head_dim = 8192 / 64 = 128, and 80 transformer layers, running in fp16 (2 bytes per number).

d_model, num_layers, head_dim = 8192, 80, 128
bytes_per_elem = 2  # fp16

def kv_cache_bytes_per_token(num_kv_heads):
    return 2 * num_kv_heads * head_dim * bytes_per_elem * num_layers

for name, n_kv in [("MHA (64 kv heads)", 64),
                    ("GQA-8 (actual Llama-2-70B)", 8),
                    ("MQA (1 kv head)", 1)]:
    b = kv_cache_bytes_per_token(n_kv)
    print(name, "->", b, "bytes/token =", round(b/1024, 2), "KB/token")

# MHA (64 kv heads) -> 2621440 bytes/token = 2560.0 KB/token
# GQA-8 (actual Llama-2-70B) -> 327680 bytes/token = 320.0 KB/token
# MQA (1 kv head) -> 40960 bytes/token = 40.0 KB/token

Every token of context costs 2.5 MB of KV cache under plain MHA, but only 320 KB under the GQA-8 configuration Llama-2-70B actually ships with — exactly the 64/8 = 8× reduction the head-count ratio predicts, because cache size is linear in num_kv_heads and every other factor is unchanged. Now scale that to a realistic serving load: a batch of 16 simultaneous conversations, each with a 4096-token context.

seq_len, batch = 4096, 16
for name, n_kv in [("MHA", 64), ("GQA-8", 8)]:
    total = kv_cache_bytes_per_token(n_kv) * seq_len * batch
    print(name, "->", round(total/1e9, 3), "GB")

# MHA  -> 171.799 GB
# GQA-8 -> 21.475 GB

171.8 GB does not fit in any single GPU available today (an H100 has 80 GB of HBM); serving that batch under plain MHA would force splitting the KV cache across multiple GPUs just to hold it, before a single token of generation happens. At 21.5 GB, GQA-8 lets that same batch and context length fit comfortably alongside the model weights on far less hardware. This — not any change to model quality — is the entire commercial reason GQA is now standard.

It's worth being precise about what GQA does not save as much as you might guess. GQA shrinks the K and V projection weight matrices too, since each now maps down to only num_kv_heads × head_dim instead of the full d_model:

def wk_wv_params(n_kv):
    return 2 * d_model * (n_kv * head_dim)   # Wk + Wv, one layer

mha, gqa = wk_wv_params(64), wk_wv_params(8)
saved = (mha - gqa) * num_layers
print(round(saved/1e9, 3), "B params saved,",
      round(saved/70e9*100, 1), "% of a 70B model")

# 9.395 B params saved, 13.4 % of a 70B model

That 13.4% one-time weight saving is real but fixed — it doesn't change no matter how many users you serve or how long their conversations run. The 8× KV-cache saving, by contrast, multiplies with every concurrent request and every token of context. A model owner cares about both; a serving engineer building the inference stack cares almost entirely about the second one, because it is the number that determines how many users fit on how many GPUs.

Worked example 2 — tracing the arithmetic of one attention step by hand

The KV-cache accounting explains why GQA exists, but it says nothing about what changes inside the attention computation itself when heads share a K/V pair. Trace it directly with a deliberately tiny setup: h = 4 query heads, head_dim = 2, grouped into G = 2 groups of 2 heads each. Heads 1–2 form Group A and share one key/value head; heads 3–4 form Group B and share a different key/value head. There are two cached tokens, and we compute the attention output for a new query at each head.

Group A's shared key/value head holds two cached token vectors:

k1 = [1, 0]   v1 = [2, 0]
k2 = [0, 1]   v2 = [0, 2]

Group B's shared key/value head holds a different pair:

k1 = [1, 1]   v1 = [1, 1]
k2 = [1, -1]  v2 = [3, 3]

Each of the four query heads has its own query vector, projected by its own W_q, even though heads within a group will read the same keys and values:

head 1 (Group A): q = [1, 0]
head 2 (Group A): q = [0, 1]
head 3 (Group B): q = [1, 1]
head 4 (Group B): q = [1, -1]

Standard scaled dot-product attention: score_i = (q · k_i) / √head_dim, softmax over the scores, then a weighted sum of the value vectors. With head_dim = 2, the scale factor is √2 ≈ 1.4142. Working head 1 by hand:

score with k1: (1×1 + 0×0)/1.4142 = 0.7071
score with k2: (1×0 + 0×1)/1.4142 = 0.0000
softmax([0.7071, 0.0000]) = [0.6698, 0.3302]
output = 0.6698×[2,0] + 0.3302×[0,2] = [1.3395, 0.6605]

The same procedure run for all four heads (verified numerically):

head 1 (Grp A, q=[1,0]):  weights=[0.6698,0.3302]  out=[1.3395, 0.6605]
head 2 (Grp A, q=[0,1]):  weights=[0.3302,0.6698]  out=[0.6605, 1.3395]
head 3 (Grp B, q=[1,1]):  weights=[0.8044,0.1956]  out=[1.3911, 1.3911]
head 4 (Grp B, q=[1,-1]): weights=[0.1956,0.8044]  out=[2.6089, 2.6089]

concat(h1,h2,h3,h4) = [1.3395, 0.6605, 0.6605, 1.3395,
                        1.3911, 1.3911, 2.6089, 2.6089]

This 8-dimensional concatenation is then multiplied by the usual output projection W_O (a fixed learned matrix, not computed here) to produce the layer's attention output, exactly as in ordinary MHA — GQA changes nothing about how heads are combined at the end, only about how many independent K/V projections feed them.

Correcting a misconception this example makes visible

A common misreading of GQA is: "if two heads share the same keys and values, they must produce redundant, near-duplicate outputs — so why keep them as separate heads at all?" The worked trace above refutes this directly. Heads 1 and 2 read from the identical key/value pool (Group A), yet head 1's output is [1.3395, 0.6605] and head 2's is [0.6605, 1.3395] — different vectors, because the softmax weighting is driven entirely by each head's own query. Head 1's query [1,0] aligns with k1, so it pulls mostly from v1; head 2's query [0,1] aligns with k2, so it pulls mostly from v2. Sharing keys and values only constrains what content is available to attend to within a group — it does not force the heads to attend to it the same way. That distinction is exactly why GQA retains most of MHA's quality: the diversity in a multi-head layer comes substantially from having many independent query "lenses" over a shared pool of information, not from every head owning a private, non-overlapping pool.

How the three regimes map onto query and key/value heads

Choosing G in practice

G is a design-time architectural choice, fixed before training and unchanged at inference — you cannot dial it per-request. In practice it is chosen empirically by measuring the quality/cache-size trade-off on a held-out set: the GQA paper's own ablations on T5-XXL-scale models found that quality rises steeply as G moves up from 1, then flattens well before reaching full MHA, so a modest group count captures most of the achievable quality at a fraction of the cache. Llama-2-70B's G = 8 (out of 64 query heads, a 1:8 ratio) and Mistral-7B's G = 8 (out of 32 query heads, a 1:4 ratio) both sit in that flat region. The practical rule of thumb used across current open models is to pick the smallest G that keeps benchmark quality within noise of full MHA — usually landing between h/8 and h/4 — rather than deriving it from a closed-form formula, because the right value depends on model width and training data, not on the mechanism alone.

Two things fall out of this that matter for anyone reading a model's config file: the number reported as num_key_value_heads (Hugging Face configs) or n_kv_heads is exactly the G of this chapter, and dividing num_attention_heads by it recovers the group size directly — for Llama-2-70B, 64/8 = 8 query heads share every K/V head.

Active recall

Attempt each question before reading its answer.

1. A model has 32 query heads and d_model = 4096. Under GQA with G = 4, how many key/value heads does it have, and what is each K/V head's dimension?
2. Why is KV-cache size, not attention FLOPs, the resource GQA is optimizing for?
3. In the four-head worked example, heads 1 and 2 shared identical keys and values yet produced different outputs. Explain precisely why, in terms of the attention formula.
4. A 40B-parameter model with 40 layers, 40 query heads, head_dim = 128, switches from MHA to GQA with G = 8, at fp16. Compute the KV cache bytes per token before and after, and the reduction factor.
5. Would converting an MHA model to MQA (G = 1) ever hurt quality more on a wide model (large h) than a narrow one, given that a single shared K/V head has to serve every query head's view? Justify with the mechanism, not intuition alone.
6. True or false: GQA reduces the number of parameters in the feed-forward sublayers of a transformer. Justify.

Answers

1. G directly is the number of K/V heads, so 4 K/V heads. Each head's dimension is unchanged by grouping — it is still d_model / num_query_heads = 4096/32 = 128, exactly as in MHA. Grouping changes how many independent K/V projections exist, not the width of any individual head.
2. Because autoregressive decoding computes one new token at a time and must stream the entire cached K/V history from memory at every step; this makes decoding memory-bandwidth bound rather than compute bound, so the dominant cost scales with cache size, not with the number of multiply-adds in the attention score computation.
3. Attention output for a head is a softmax-weighted sum of the value vectors, where the weights come from that head's own query dotted against the (possibly shared) keys. Sharing keys/values only fixes the pool of content available; each head's distinct query vector still produces a distinct softmax distribution over that pool, so the weighted sums differ. Head 1's query aligned more with key 1, head 2's with key 2, despite both drawing from the same {key1,key2,value1,value2} pool.
4. Before (MHA, 40 kv heads): 2 × 40 × 128 × 2 bytes × 40 layers = 819,200 bytes/token ≈ 800 KB/token. After (GQA-8): 2 × 8 × 128 × 2 × 40 = 163,840 bytes/token = 160 KB/token. Reduction factor = 819200/163840 = 5.0×, matching 40/8 exactly, since every other factor in the formula is unchanged and cache size is linear in num_kv_heads.
5. Yes. A wider model has more query heads competing to extract different information through one shared K/V projection, so the bottleneck each query head faces is more severe as h grows — the single shared subspace has to be "wide enough" to support more simultaneous, potentially conflicting query views. This matches the empirical finding that MQA's quality gap versus MHA widens at larger scale, and is exactly why GQA (partial sharing, scaling G with model size) rather than full MQA became the standard for large models.
6. False. GQA only changes the key and value projection matrices inside the attention sublayer (W_k, W_v); it leaves W_q, the output projection W_O, and the entire feed-forward sublayer untouched. In a typical large transformer the feed-forward block holds the majority of parameters (for Llama-2-70B, roughly 56B of the ~70B total), so even an 8× cut to W_k/W_v only trims a double-digit percentage off total model size.

Think About It

Think about this: How would you explain grouped query attention (gqa): efficient multi-head attention 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.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind grouped query attention (gqa): efficient multi-head attention, 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.