Grouped Query Attention (GQA)

Grouped Query Attention (GQA) changes how the key and value projections in self-attention are shared across heads, so that a model can serve requests faster and cheaper without retraining from scratch or meaningfully hurting output quality. Standard multi-head attention gives every query head its own key and value head; GQA instead splits the query heads into a smaller number of groups, and every query head in a group reads from one shared key/value head. Introduced in 2023 and adopted almost immediately by Llama 2’s larger checkpoints, GQA is now the default attention shape for transformer-based LLMs shipping in production, from Llama and Mistral to Qwen, Gemma, and GPT-OSS.

The Inference Bottleneck: The KV Cache

During autoregressive decoding, a transformer caches the key and value vectors for every previous token so it never has to recompute them, one cache slot per layer, per head. That KV cache has to live in GPU memory for the entire length of a generation, and it grows linearly with both sequence length and batch size. For a model with many attention heads and a large context window, the KV cache can become larger than the model’s own weights, and every token generated has to stream all of it out of high-bandwidth memory before the next token can be produced.

This is a memory-bandwidth problem, not a compute problem. Decoding one token barely uses the GPU’s FLOPs, since matrix multiplications are small when the batch is one token wide, but it still has to read the full KV cache off HBM. Flash Attention solves the analogous bandwidth problem for training and prefill by tiling the computation; GQA solves it for decoding by shrinking the thing that has to be read in the first place. Cut the number of distinct key/value heads by 8x and the KV cache traffic per token drops by roughly 8x too.

The Mechanism: Query Heads Share a Key/Value Head

Let h be the number of query heads and g the number of key/value head groups, where 1 ≤ g ≤ h. GQA partitions the h query heads into g groups of h/g heads each, and every query head in a group attends against the same, single key and value head. Three settings of g recover three named mechanisms:

• g = h (Multi-Head Attention, MHA): every query head has its own key/value head. Maximum expressiveness, maximum KV cache.
• 1 < g < h (Grouped Query Attention, GQA): query heads share key/value heads within a group. A tunable middle ground.
• g = 1 (Multi-Query Attention, MQA): all query heads share a single key/value head. Minimum KV cache, but the sharpest quality cost, since MQA was proposed by Noam Shazeer in 2019, well before GQA generalized it into a dial rather than an all-or-nothing switch.

graph TB
    classDef default fill:#ffffff,stroke:#4338CA,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef data    fill:#EEF0F7,stroke:#0D9488,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef process fill:#F7F8FC,stroke:#6366F1,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef output  fill:#4338CA,stroke:#4338CA,stroke-width:2px,color:#ffffff,rx:8px,ry:8px;

    subgraph MHA["Multi-Head Attention (g = h = 4)"]
        Q1([Q1]):::data --> K1[(K1, V1)]:::process
        Q2([Q2]):::data --> K2[(K2, V2)]:::process
        Q3([Q3]):::data --> K3[(K3, V3)]:::process
        Q4([Q4]):::data --> K4[(K4, V4)]:::process
    end
    subgraph GQA["Grouped Query Attention (g = 2)"]
        G1([Q1]):::data --> GK1[(K, V · group A)]:::process
        G2([Q2]):::data --> GK1
        G3([Q3]):::data --> GK2[(K, V · group B)]:::process
        G4([Q4]):::data --> GK2
    end
    subgraph MQA["Multi-Query Attention (g = 1)"]
        M1([Q1]):::data --> MK[(shared K, V)]:::output
        M2([Q2]):::data --> MK
        M3([Q3]):::data --> MK
        M4([Q4]):::data --> MK
    end

In code, GQA is usually implemented by projecting queries with the full h heads but projecting keys and values with only g heads, then repeating each key/value head h/g times before the batched attention matmul so the shapes line up:

# scenario: a serving-side attention kernel for a Llama-3-style
# model with 64 query heads and 8 KV heads
import torch

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    # x: (batch, num_kv_heads, seq_len, head_dim)
    b, g, s, d = x.shape
    if n_rep == 1:
        return x
    x = x[:, :, None, :, :].expand(b, g, n_rep, s, d)
    return x.reshape(b, g * n_rep, s, d)

# q: (batch, 64, seq_len, head_dim), k/v: (batch, 8, seq_len, head_dim)
k_full = repeat_kv(k, n_rep=q.shape[1] // k.shape[1])  # -> (batch, 64, ...)
v_full = repeat_kv(v, n_rep=q.shape[1] // v.shape[1])
attn_out = torch.nn.functional.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

Only 8 distinct key/value tensors are ever stored or read from the KV cache; the repeat_kv expansion happens on the fly, entirely in fast on-chip memory, so it doesn’t add HBM traffic.

Where g Sits Between Two Extremes

The GQA paper’s key empirical finding is that quality degrades gracefully as g shrinks from h, but not linearly: dropping from MHA to a moderate group count (Llama 3 70B uses g=8 against h=64) costs almost nothing on held-out benchmarks, while collapsing all the way to MQA’s g=1 gives up noticeably more. That’s why virtually no current frontier open-weight model ships as pure MQA, but almost none ship as pure MHA either; GQA’s tunable middle sits in the sweet spot models actually use.

Try It: KV Cache vs. Group Count

The formula behind the calculator below is exact, not illustrative: KV cache size in bytes is 2 (K and V) × layers × g × head_dim × sequence_length × batch_size × bytes_per_element. Only g (the number of KV heads) and the sequence length change here; everything else is fixed to Llama-3-70B’s real shape (80 layers, 64 query heads, head_dim 128, fp16). Drag g down from 64 and watch both the grid of query heads regroup by color and the cache size fall in lockstep.

Drag g from 64 (MHA) down to 1 (MQA) and slide the context length out to 128K to see why nobody serves full MHA at long context.

Training GQA Models: Uptraining, Not Just From-Scratch

The original GQA paper’s other contribution was showing that an existing MHA checkpoint doesn’t have to be retrained from scratch to gain GQA’s efficiency. Its “uptraining” recipe mean-pools each group’s existing key/value heads into a single shared head, then continues pretraining for a small fraction, roughly 5%, of the original compute budget. That’s what let Llama 2’s 34B and 70B checkpoints ship with GQA without a from-scratch retrain, and it’s the same pattern later re-used for converting other architectures’ attention layers post-hoc.

# scenario: converting an already-trained MHA checkpoint to GQA
# by mean-pooling key/value heads within each new group, before
# a short continued-pretraining run (the "uptraining" recipe)
import torch

def uptrain_to_gqa(k_heads: torch.Tensor, num_groups: int) -> torch.Tensor:
    # k_heads: (num_original_heads, head_dim) for one layer
    h = k_heads.shape[0]
    heads_per_group = h // num_groups
    k_grouped = k_heads.view(num_groups, heads_per_group, -1)
    return k_grouped.mean(dim=1)  # (num_groups, head_dim), then continue pretraining

Where GQA Is Used Today

Beyond GQA: Multi-Head Latent Attention

GQA reduces the KV cache by sharing whole key/value heads across groups. DeepSeek-V2 pushed the same goal further with Multi-Head Latent Attention (MLA), which instead compresses each token’s keys and values into a single low-rank latent vector and reconstructs per-head keys and values from it on the fly. MLA reports a smaller KV cache than an equivalent GQA setup at comparable or better quality, and later work (TransMLA) shows that any GQA model can be re-expressed as an MLA model, framing MLA as a strictly more expressive generalization of the same underlying idea: don’t store more key/value information per token than attention actually needs to read back. Where GQA quantizes that trade-off into a small number of discrete groups, MLA makes it continuous through the latent dimension’s rank, at the cost of a more involved reconstruction step at attention time.

What’s New (2025–2026)

• MLA displacing GQA at the frontier. DeepSeek’s V2/V3-family models and their MLA variant have pushed several other labs to evaluate low-rank KV compression as a GQA alternative rather than a niche technique, particularly for very long context serving.
• GQA-μP. Recent work on maximal-update parameterization for GQA studies how learning rates and initialization should scale specifically when query and key/value head counts diverge, closing a gap in standard μP theory that assumed MHA’s one-to-one head correspondence.
• GQA combined with sliding window attention as the default recipe, not an either/or choice: most serving-optimized 2025–2026 open-weight models (Mistral, Gemma, GPT-OSS) now ship both simultaneously, since GQA shrinks the width of the KV cache per token and sliding window attention shrinks the number of tokens that need caching per layer, two independent, stackable savings.
• Key-driven grouping. Rather than grouping query heads arbitrarily (e.g. by index), some newer approaches choose group membership based on learned key similarity, aiming to keep the heads sharing a KV head closer in what they actually attend to.

from transformers import LlamaConfig, LlamaForCausalLM

# Llama-3-70B's actual GQA-8 setup: 64 query heads grouped into
# 8 KV heads, so every KV head is shared by 64/8 = 8 query heads.
config = LlamaConfig(
    hidden_size=8192,
    num_attention_heads=64,
    num_key_value_heads=8,        # <- this field is what makes it GQA
    num_hidden_layers=80,
    max_position_embeddings=8192,
)
model = LlamaForCausalLM(config)

# num_key_value_heads == num_attention_heads    -> standard MHA
# num_key_value_heads == 1                      -> MQA
# 1 < num_key_value_heads < num_attention_heads -> GQA