One-sentence summary: MHA, MQA, and GQA are one parameterized family—change the number of KV heads and you move along a quality–memory–speed trade-off that must be measured on the actual model and hardware.

23.1 The problem is not “how many tensors”

Chapter 22 established what happens when the model generates a new token. Every layer computes a fresh Q, K, and V. Q is consumed immediately; the new K and V are appended to that layer's KV Cache.

Here is an easy mistake to make: 32 heads do not normally mean 64 separate cache tensors per layer. A common implementation still holds two logical tensors:

K cache: [B, H_KV, S, d_h]
V cache: [B, H_KV, S, d_h]

The head count $H_{KV}$ is one axis of those tensors. Storage is determined by the length of that axis:

\text{bytes}=B\times L\times 2\times H_{KV}\times S\times d_h\times b

Here $B$ is batch size, $L$ the number of layers, $S$ cached tokens, $d_h$ head dimension, and $b$ bytes per element. The factor 2 is for K and V.

MHA, GQA, and MQA as one family defined by query-head and KV-head counts

Two numbers describe the whole family:

$H_Q$ : number of Query heads
$H_{KV}$ : number of Key heads and Value heads

We require $H_Q$ to be divisible by $H_{KV}$ . The number of Query heads sharing each K/V pair is:

R=\frac{H_Q}{H_{KV}}

Therefore:

Mechanism	Condition	Example with 8 Query heads
MHA	$H_{KV}=H_Q$	8 KV heads; 1 Q head per group
GQA	$1<H_{KV}<H_Q$	2 KV heads; 4 Q heads per group
MQA	$H_{KV}=1$	1 KV head shared by all 8 Q heads

MHA and MQA are not two unrelated inventions, and GQA is not a fourth piece of machinery wedged between them. They are the two endpoints and the interior of one parameterized family.

23.2 Intuition: eight luthiers, several pattern books

Imagine eight luthiers inspecting a half-restored nyckelharpa. Each can ask a different question: Is the key action binding? Which sympathetic string is buzzing? Has the bridge warped?

MHA gives every luthier a private pattern book. Each asks a distinct question and keeps a distinct catalogue of the instrument's details.
MQA keeps the eight questions, but everyone consults one shared pattern book.
GQA divides the luthiers into several benches; each bench shares a book, while different benches retain different catalogues.

The questions correspond to Q projections. The way history is indexed and carried forward corresponds to K and V. Sharing K/V reduces the stored history, but it does not make the Query heads identical: their Q rows and attention outputs remain distinct.

Q, K, V, and KV-cache tensor shapes in grouped-query attention

The three mechanisms share one set of shapes:

q: [B, H_Q,  T_q, d_h]
k: [B, H_KV, T_k, d_h]
v: [B, H_KV, T_k, d_h]

For $H_Q=8$ and $H_{KV}=2$ , the mapping is:

Q0 Q1 Q2 Q3  -> KV0
Q4 Q5 Q6 Q7  -> KV1

Conceptually, each KV head can be repeated $R$ times before ordinary multi-head attention. That is useful for explanation and reference code. An efficient kernel should map each Query head straight to its KV head, without inflating the cache in memory.

23.3 What exactly becomes smaller?

23.3.1 KV Cache

Reuse Chapter 22's dimensions: 32 layers, $d_h=128$ , 1,024 tokens, FP16 or BF16 at 2 bytes per element, and batch size 1.

Mechanism	$H_{KV}$	Exact bytes	Binary unit	Relative to MHA
MHA	32	536,870,912	512 MiB	100%
GQA	8	134,217,728	128 MiB	25%
MQA	1	16,777,216	16 MiB	3.125%

KV-cache and attention-projection parameter comparison for MHA, GQA, and MQA

At 4,096 tokens, those numbers become 2 GiB, 512 MiB, and 64 MiB. The ratios are exact. “Four times the concurrency” and “32 times the context” are not. A server must also hold weights, workspaces, activations, allocator slack, and fragmented blocks. A model's usable context also depends on positional encoding and what lengths it was trained to handle.

23.3.2 Projection parameters

Ignoring biases, let model width be $D=H_Qd_h$ . Q and output projections each contain $D^2$ weights. K and V each contain $D(H_{KV}d_h)$ :

P_{attn}=2D^2+2D^2\frac{H_{KV}}{H_Q}

For $H_Q=32$ :

Mechanism	Attention-projection weights
MHA, $H_{KV}=32$	$4D^2$
GQA, $H_{KV}=8$	$2.5D^2$
MQA, $H_{KV}=1$	$2.0625D^2$

Moving from 32/32 MHA to 32/8 GQA removes 37.5% of the attention-projection weights. It does not remove 37.5% of a complete 7B model, nor does it save a universal 6%. The whole-model fraction depends on its FFN, embeddings, depth, and weight sharing.

23.4 One implementation for MHA, GQA, and MQA

The code below uses [B, H, T, d_h] throughout. grouped_attention_reference explicitly repeats K and V, which makes it easy to inspect. grouped_attention uses PyTorch's native enable_gqa path.

import torch
import torch.nn.functional as F


def repeat_kv(x, num_query_heads):
    """Reference expansion: [B, H_kv, T, D_h] -> [B, H_q, T, D_h]."""
    num_kv_heads = x.size(1)
    if num_query_heads <= 0 or num_kv_heads <= 0:
        raise ValueError("head counts must be positive")
    if num_query_heads % num_kv_heads != 0:
        raise ValueError("num_query_heads must be divisible by num_kv_heads")
    return x.repeat_interleave(num_query_heads // num_kv_heads, dim=1)


def validate_gqa_shapes(q, k, v):
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError("q, k, and v must be four-dimensional")
    if q.size(0) != k.size(0) or q.size(0) != v.size(0):
        raise ValueError("q, k, and v must have the same batch size")
    if q.size(1) <= 0 or k.size(1) <= 0:
        raise ValueError("head counts must be positive")
    if q.size(1) % k.size(1) != 0 or k.size(1) != v.size(1):
        raise ValueError("invalid GQA head counts")
    if q.size(-1) != k.size(-1):
        raise ValueError("q and k must have the same head dimension")
    if k.size(-2) != v.size(-2):
        raise ValueError("k and v must have the same sequence length")


def grouped_attention_reference(q, k, v, *, is_causal):
    validate_gqa_shapes(q, k, v)
    k_full = repeat_kv(k, q.size(1))
    v_full = repeat_kv(v, q.size(1))
    return F.scaled_dot_product_attention(
        q, k_full, v_full, is_causal=is_causal, dropout_p=0.0
    )


def grouped_attention(q, k, v, *, is_causal):
    validate_gqa_shapes(q, k, v)
    return F.scaled_dot_product_attention(
        q,
        k,
        v,
        is_causal=is_causal,
        dropout_p=0.0,
        enable_gqa=q.size(1) != k.size(1),
    )

Within the same function:

8 heads in q and 8 in k/v is MHA;
8 heads in q and 2 in k/v is GQA;
8 heads in q and 1 in k/v is MQA.

is_causal is deliberately required rather than hidden behind a default. Prefill and cached Decode do not always use the same value:

for equal-length training or Prefill, pass is_causal=True;
for single-token Decode, when K/V contain only past and current tokens, pass is_causal=False because the cache contains no future token to mask;
for chunked Decode with several new tokens, build a lower-right-aligned causal mask from the actual cache positions.

PyTorch uses an upper-left-aligned bias when is_causal=True is applied to a non-square attention matrix. On a cached step shaped [T_q=1, T_k>1], blindly passing True therefore keeps only the earliest key instead of exposing the full cache. This is a mask-alignment boundary, independent of whether the attention is MHA, GQA, or MQA.

As of July 2026, PyTorch still documents enable_gqa as experimental and lists backend and tensor-type constraints. The API existing does not prove that a particular GPU, dtype, and software version will select the best kernel; inspect the selected backend and profile it. The explicit reference always expands K and V. Whether the native API avoids that expansion depends on the backend and kernel actually selected.

The executable checks compare outputs and gradients for MHA, GQA, and MQA in float32 and float64. They also cover single-token cached Decode, the non-square causal-mask trap, and a Value head dimension different from the Query/Key head dimension.

23.5 Turning an MHA checkpoint into GQA

The GQA paper's uptraining procedure is more than changing one configuration field:

Mean-pool each group of MHA K heads and V heads to initialize the smaller GQA projections. Q and output projections stay unchanged.
Continue pretraining so the model can adapt to the shared K/V representation.

Mean-pooling MHA K/V heads and continuing pretraining to produce GQA

The paper's “5%” means 5% of the original pretraining compute, not a universal recipe using 5% of the data. In its T5.1.1 experiments, the paper reported uptrained GQA close to MHA quality with speed comparable to MQA. That is a scoped experimental result, not a cross-model guarantee.

PyTorch stores nn.Linear.weight as [out_features, in_features], so the head axis is on the output side when pooling a K or V projection:

def mean_pool_kv_projection(weight, num_query_heads, num_kv_heads):
    """Pool [H_q * D_h, D_model] into [H_kv * D_h, D_model]."""
    if weight.ndim != 2:
        raise ValueError("weight must be two-dimensional")
    if num_query_heads <= 0 or num_kv_heads <= 0:
        raise ValueError("head counts must be positive")
    if num_query_heads % num_kv_heads != 0:
        raise ValueError("num_query_heads must be divisible by num_kv_heads")
    if weight.size(0) % num_query_heads != 0:
        raise ValueError("weight rows must be divisible by num_query_heads")
    head_dim = weight.size(0) // num_query_heads
    group_size = num_query_heads // num_kv_heads
    return (
        weight.reshape(num_query_heads, head_dim, weight.size(1))
        .reshape(num_kv_heads, group_size, head_dim, weight.size(1))
        .mean(dim=1)
        .reshape(num_kv_heads * head_dim, weight.size(1))
    )

If the linear layer has a bias, pool it by the same head groups. After conversion, re-evaluate loss, downstream tasks, long-context behavior, and serving performance. Mean-pooling is an initialization, not evidence that quality survived.

23.6 How real model configurations differ

The table deliberately shows a few snapshots that can be checked in papers or official configurations. Sizes and revisions within one family can differ; always read the config.json for the exact checkpoint being loaded.

Model snapshot	Q heads	KV heads	Mechanism	Q heads per group
Llama 2 7B	32	32	MHA	1
Llama 2 70B	64	8	GQA	8
Mistral-7B-v0.1	32	8	GQA	4
Qwen2-7B	28	4	GQA	7

Reading query-head and KV-head fields from a model config, then checking native kernel support

A Hugging Face-style configuration commonly exposes:

{
  "num_attention_heads": 32,
  "num_key_value_heads": 8
}

The interpretation is mechanical:

equal values mean MHA;
num_key_value_heads equal to 1 means MQA;
a value in between means GQA;
the Query-head count must also be divisible by the KV-head count.

“Eight KV heads is the sweet spot” is not a law. Eight divides neatly across some 2-, 4-, or 8-way tensor-parallel layouts, but systems may need replication or a different sharding plan when KV heads are fewer than devices. Quality curves also change with the model, data, and training budget. Choose candidate $H_{KV}$ values, then measure validation quality, TPOT/throughput, peak memory, and the target parallel layout together.

23.7 Reading the paper claims at their proper scale

The 2019 MQA paper reported much faster decoding with a small quality degradation on the translation tasks it evaluated. The 2023 GQA paper reported, in its T5.1.1 uptraining experiments, that GQA could approach MHA quality while retaining speed comparable to MQA.

Both sentences contain an implicit phrase: “in the evaluated setting.” They do not prove that:

MHA has the best quality for every model;
MQA is unacceptable at frontier scale;
GQA is lossless;
reducing KV heads from 32 to 8 makes decoding exactly four times faster;
eight KV heads is right for every architecture.

The architecture changes the available trade-off. Training, kernels, hardware, batch size, context length, and serving software decide where a real system lands.

23.8 Chapter summary

KV Cache is commonly two logical tensors per layer. $H_{KV}$ is an axis, not “one tensor per head.”
MHA, GQA, and MQA form an $H_Q/H_{KV}$ family, with endpoints $H_{KV}=H_Q$ and $H_{KV}=1$ .
Cache bytes scale linearly with $H_{KV}$ , but real concurrency and speed depend on more than that ratio.
A native GQA kernel maps Query heads directly to KV heads. Explicit repeat_kv is primarily a reference implementation.
Mean-pooling an MHA checkpoint only initializes GQA; the paper then spent 5% of the original pretraining compute on uptraining.
There is no timeless eight-head default. Read the exact config and measure quality and performance on the target system.

References

Attention Is All You Need (Vaswani et al., 2017)
Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron et al., 2023)
Official Mistral-7B-v0.1 configuration
Official Qwen2-7B configuration
PyTorch scaled_dot_product_attention documentation

Next chapter

GQA reduces how much K/V data each token contributes, but it does not change which historical tokens a Query attends to. Chapter 24 asks the next question: what if attention stops looking at the complete history and visits only selected positions? That is the problem Sparse Attention sets out to solve.