Grouped Query Attention

MHA vs MQA vs GQA

Three variants of multi-head attention differ in how many key-value (KV) head sets they maintain:

• Multi-Head Attention (MHA): Original Transformer attention. Each of H query heads has its own K and V projection. KV cache grows with H.
• Multi-Query Attention (MQA): All query heads share a single KV head. Extreme KV cache savings but slight quality drop — used in PaLM, early Falcon.
• Grouped Query Attention (GQA): Query heads are split into G groups; each group shares one KV head. Interpolates between MHA (G=H) and MQA (G=1). Llama 3 70B uses H=64 query heads and G=8 KV heads.

Why KV cache size matters

During autoregressive generation, you cache the K and V tensors for every previously generated token to avoid recomputing them. For a model with H attention heads, d_model dimensions, and context length L, the KV cache size per layer is:

2 × L × d_head × H_kv × bytes_per_element

For Llama 3 70B (bf16) with 80 layers, 4096 d_head, 8 KV heads, at context length 8192:

2 × 8192 × 128 × 8 × 80 × 2 bytes ≈ 21 GB

With full MHA (64 KV heads) this would be 168 GB — impossible to fit on a single GPU. GQA makes long-context 70B inference feasible on 2–4 GPUs.

GQA implementation

import torch, torch.nn as nn, 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, "n_heads must be divisible by n_kv_heads"
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads   # query heads per KV group
        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(d_model, d_model, bias=False)

    def forward(self, x, mask=None):
        B, T, _ = x.shape
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_kv_heads, self.d_head).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_kv_heads, self.d_head).transpose(1, 2)

        # Repeat KV heads to match query head count
        K = K.repeat_interleave(self.n_rep, dim=1)   # (B, n_heads, T, d_head)
        V = V.repeat_interleave(self.n_rep, dim=1)

        scale = math.sqrt(self.d_head)
        attn = torch.softmax(torch.matmul(Q, K.transpose(-2, -1)) / scale, dim=-1)
        if mask is not None:
            attn = attn.masked_fill(mask == 0, 0)

        out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out)

Memory savings calculation

def kv_cache_gb(n_layers, seq_len, n_kv_heads, d_head, bytes_per_elem=2):
    # bytes_per_elem: 2 for bf16/fp16, 4 for fp32
    bytes = 2 * n_layers * seq_len * n_kv_heads * d_head * bytes_per_elem
    return bytes / 1e9

# Llama 3 70B specs
print("MHA (64 KV heads):", kv_cache_gb(80, 8192, 64, 128), "GB")
print("GQA (8 KV heads): ", kv_cache_gb(80, 8192,  8, 128), "GB")
# MHA (64 KV heads): 168.95 GB
# GQA (8 KV heads):   21.12 GB

GQA in Llama 3 and Mistral

All modern efficient LLMs use GQA. Llama 3 8B uses 8 KV heads and 32 query heads (4:1 ratio). Llama 3 70B uses 8 KV heads and 64 query heads (8:1 ratio). Mistral 7B uses 8 KV heads and 32 query heads. Gemma 2 uses GQA with varying group sizes by layer size.

The configuration is in the model's config.json: look for num_key_value_heads vs num_attention_heads. If they differ, it's GQA (or MQA if num_key_value_heads == 1).

Training and conversion

GQA models must be trained (or fine-tuned) with GQA from the start — you can't simply remove KV heads from a trained MHA model without significant quality loss. However, the Ainslie et al. (2023) GQA paper showed that you can convert an MHA checkpoint to GQA by mean-pooling the KV heads of each group and then fine-tuning for a fraction of the original training compute (<5%), recovering most of the original quality.

Gotchas

• repeat_interleave vs expand: In inference code, you can avoid materialising repeated KV heads by using expand (zero-copy view) instead of repeat_interleave (copies data). Flash Attention 2 handles GQA natively without expansion.
• Flash Attention 2 GQA: Use flash_attn_func with the gqa argument — it handles the group broadcasting internally and is significantly faster than the naive repeat approach.
• Checkpoint loading: When loading GQA weights with HuggingFace, verify config.num_key_value_heads matches the checkpoint. Mismatches cause silent shape errors.