Grouped Query Attention (GQA)

Grouped query attention (Ainslie et al., 2023) partitions the H query heads into G groups, each sharing one key head and one value head. The total number of KV heads equals G, with H / G query heads per group. Two degenerate cases make the design space clear: when G = 1, every query head shares one KV head — this is MQA. When G = H, every query head has its own KV head — this is standard MHA. Any value of G between 1 and H is a point on the tradeoff curve between memory and representational capacity.

Head grouping mechanics

For a model with H = 32 query heads and G = 8 KV groups, each group contains $\frac{32}{8} = 4$ query heads sharing one key and one value projection. The attention computation within each group is identical to standard attention — the only change is that multiple query heads address the same K and V.

The projection dimensions:

Q: d_model × d_model (same as MHA). 32 heads, each (d_model, d_head).
K: d_model × (G * d_head). 8 heads instead of 32. Shrinks by $\frac{H}{G} = 4x$ .
V: d_model × (G * d_head). Same reduction as K.

For $d_{\text{model}} = 4096$ , H = 32, G = 8, $d_{\text{head}} = 128$ :

MHA K/V params: $4096 \cdot 4096 = 16.8M$ each
GQA-8 K/V params: $4096 \cdot 1024 = 4.2M$ each
MQA K/V params: $4096 \cdot 128 = 0.52M$ each

Production configurations

Llama 3 70B: 64 query heads, 8 KV heads. Each KV head serves $\frac{64}{8} = 8$ query heads. KV cache is $\frac{8}{64} = 12.5%$ of MHA.

Llama 3 8B: 32 query heads, 8 KV heads. Each KV head serves $\frac{32}{8} = 4$ query heads. KV cache is $\frac{8}{32} = 25%$ of MHA.

Mistral 7B: 32 query heads, 8 KV heads. Same ratio as Llama 3 8B.

Gemma 7B (Google DeepMind, 2024): 16 query heads, 16 KV heads. G = H — this is actually full MHA, not GQA.

Gemma 2 27B: 32 query heads, 16 KV heads. G = 16, so each KV head serves 2 query heads. KV cache is 50% of MHA.

Qwen 2 72B: 64 query heads, 8 KV heads. Same configuration as Llama 3 70B.

The pattern is consistent: G = 8 has become the default for models with 32 or 64 query heads. This is not a coincidence — Ainslie et al. (2023) found that GQA-8 matches MHA quality on most benchmarks while cutting KV cache to 12.5–25% of the MHA baseline.

KV cache comparison

For a model with $d_{\text{model}} = 8192$ , H = 64, $d_{\text{head}} = 128$ , at sequence length 4096 in float16:

MHA (G = 64): $4096 \cdot 2 \cdot 64 \cdot 128 \cdot 2 = 128 \text{MB}$ per sequence

GQA-8 (G = 8): $4096 \cdot 2 \cdot 8 \cdot 128 \cdot 2 = 16 \text{MB}$ per sequence

GQA-4 (G = 4): $4096 \cdot 2 \cdot 4 \cdot 128 \cdot 2 = 8 \text{MB}$ per sequence

MQA (G = 1): $4096 \cdot 2 \cdot 1 \cdot 128 \cdot 2 = 2 \text{MB}$ per sequence

At batch size 128 (a realistic serving scenario for a 70B model on 8x A100):

MHA: $128 \cdot 128 \text{MB} = 16.4 \text{GB}$
GQA-8: $128 \cdot 16 \text{MB} = 2.0 \text{GB}$
GQA-4: $128 \cdot 8 \text{MB} = 1.0 \text{GB}$
MQA: $128 \cdot 2 \text{MB} = 0.25 \text{GB}$

GQA-8 recovers 87.5% of MQA's memory savings relative to MHA while retaining 8 independent KV representations. Those 8 independent KV heads are the difference between MQA's compression bottleneck and GQA's near-MHA quality.

PyTorch implementation

python

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_groups = n_heads // n_kv_heads  # queries per KV head
        self.d_head = d_model // n_heads

        self.W_Q = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.W_K = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.W_V = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.W_O = nn.Linear(n_heads * self.d_head, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, S, _ = x.shape

        q = self.W_Q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        k = self.W_K(x).view(B, S, self.n_kv_heads, self.d_head).transpose(1, 2)
        v = self.W_V(x).view(B, S, self.n_kv_heads, self.d_head).transpose(1, 2)

        # Expand KV heads to match query heads:
        # (B, n_kv_heads, S, d_head) -> (B, n_heads, S, d_head)
        k = k.unsqueeze(2).expand(-1, -1, self.n_groups, -1, -1)
        k = k.reshape(B, self.n_heads, S, self.d_head)
        v = v.unsqueeze(2).expand(-1, -1, self.n_groups, -1, -1)
        v = v.reshape(B, self.n_heads, S, self.d_head)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.W_O(out)

The critical step is the expand-and-reshape of the KV heads. Starting from (B, n_kv_heads, S, d_head), the unsqueeze(2) inserts a group dimension, expand repeats each KV head n_groups times without copying memory (it creates a view), and reshape merges the KV head and group dimensions back into n_heads. After this, the attention computation is identical to MHA.

Verification: GQA degenerates correctly

python

d_model = 4096
n_heads = 32

gqa_as_mqa = GroupedQueryAttention(d_model, n_heads, n_kv_heads=1)
gqa_as_mha = GroupedQueryAttention(d_model, n_heads, n_kv_heads=32)

# KV parameter counts
mqa_kv = sum(p.numel() for n, p in gqa_as_mqa.named_parameters()
             if "W_K" in n or "W_V" in n)
mha_kv = sum(p.numel() for n, p in gqa_as_mha.named_parameters()
             if "W_K" in n or "W_V" in n)

print(f"GQA(G=1) KV params: {mqa_kv:,}")    # 1,048,576 (= 2 * 4096 * 128)
print(f"GQA(G=32) KV params: {mha_kv:,}")   # 33,554,432 (= 2 * 4096 * 4096)

Uptraining: converting MHA to GQA

Ainslie et al. (2023) showed that an existing MHA model can be converted to GQA without training from scratch. The procedure:

Group the existing H KV heads into G groups.
Within each group, average (mean-pool) the weight matrices of the H / G key heads and H / G value heads.
Replace the per-head KV projections with the mean-pooled group projections.
Fine-tune for a fraction of the original training compute — Ainslie et al. used ~5% of original training FLOPs.

python

def convert_mha_to_gqa(mha_model, n_kv_heads: int):
    """Mean-pool adjacent KV heads to create GQA weights."""
    n_heads = mha_model.n_heads
    d_head = mha_model.d_head
    group_size = n_heads // n_kv_heads

    # Original KV weights: (d_model, n_heads * d_head)
    k_weight = mha_model.W_K.weight.data.view(-1, n_heads, d_head)
    v_weight = mha_model.W_V.weight.data.view(-1, n_heads, d_head)

    # Reshape into groups and mean-pool
    k_grouped = k_weight.view(-1, n_kv_heads, group_size, d_head).mean(dim=2)
    v_grouped = v_weight.view(-1, n_kv_heads, group_size, d_head).mean(dim=2)

    # k_grouped: (d_model, n_kv_heads, d_head) -> (d_model, n_kv_heads * d_head)
    return k_grouped.reshape(-1, n_kv_heads * d_head), \
           v_grouped.reshape(-1, n_kv_heads * d_head)

Ainslie et al. demonstrated this technique on T5 models, recovering near-MHA quality after fine-tuning for roughly 5% of the original training budget. Llama 2 70B was trained from scratch with GQA-8 (Touvron et al., 2023) — the decision to use GQA was made before training began, not retrofitted via uptraining.

Quality: why GQA-8 works

The intuition is dimensional. In MQA, a single KV head of dimension $d_{\text{head}} = 128$ must encode all the contextual information that H query heads need to extract. That's a severe bottleneck for 32 or 64 query heads with diverse learned specializations.

GQA-8 provides 8 independent KV representations, each of dimension d_head. The total KV representational capacity is $8 \cdot 128 = 1024$ dimensions — compared to 128 for MQA and 4096 for MHA (at H = 32). For most tasks, 1024 KV dimensions encode enough diversity that the quality gap relative to MHA is within noise. Ainslie et al. (2023) reported on T5-XXL (13B):

MHA: baseline quality
GQA-8: within 0.2% of MHA on SuperGLUE and summarization
MQA: 0.5–1.0% below MHA on the same tasks

The gap between GQA-8 and MHA is small enough that the 8x KV cache reduction is a clear win. The gap between GQA-8 and MQA is just large enough — especially on tasks with long-range dependencies — that the additional 8x memory cost of GQA-8 over MQA is justified for general-purpose models.

Memory comparison in practice

python

def kv_cache_bytes(seq_len, n_kv_heads, d_head, batch_size, dtype_bytes=2):
    return seq_len * 2 * n_kv_heads * d_head * batch_size * dtype_bytes

seq_len = 8192
d_head = 128
batch = 32

configs = {
    "MHA (H=32)":  kv_cache_bytes(seq_len, 32, d_head, batch),
    "GQA-8":       kv_cache_bytes(seq_len, 8, d_head, batch),
    "GQA-4":       kv_cache_bytes(seq_len, 4, d_head, batch),
    "MQA (G=1)":   kv_cache_bytes(seq_len, 1, d_head, batch),
}

for name, size in configs.items():
    print(f"{name}: {size / 1e9:.2f} GB")

# MHA (H=32): 4.29 GB
# GQA-8:      1.07 GB
# GQA-4:      0.54 GB
# MQA (G=1):  0.13 GB

At batch size 32 with 8K context, GQA-8 keeps the KV cache at ~1 GB — comfortably fitting on any serving GPU alongside the model weights. MHA at 4.3 GB can become the dominant memory consumer on 40 GB GPUs serving 70B+ parameter models.

← Previous