05 — Foundations

Multi-Head Attention#

MHA · MQA · GQAMathematical 1 min read
◆ The PatternRunning multiple attention computations in parallel

Instead of one big attention operation, we split into h heads, each with dimension d_k = d_model / h. Each head learns different patterns — one might track syntax, another coreference, another positional relationships.

MultiHead(Q,K,V) = Concat(head₁, ..., headₕ) · W_O
Run h parallel attention ops with different projections, concat results, project back to d_model.
MQA: all heads share K,V — only Q varies per head
Multi-Query Attention. Massive KV-cache savings (÷ h). Used in PaLM, Falcon.
GQA: groups of heads share K,V — middle ground
Grouped-Query Attention. G groups, each with h/G query heads. LLaMA 2 70B uses 8 KV heads for 64 query heads.

Standard MHA with 32 heads and d_model=4096 → d_k=128 per head. GQA is now the default for large models: it trades a tiny quality hit for huge inference speedups via smaller KV caches.

Interactive — compare MHA, MQA, GQA head layouts

Python — grouped-query attention#

Pattern bridge: Multiple attention heads capture different relationship types in parallel — like running several correlation analyses simultaneously. In markets, combining RSI, MACD, and volume is multi-headed analysis.
Worked example

Llama 2 70B has 64 query heads. If each had its own keys and values, the KV-cache for one 4k-token request would be 10.7 GB. Grouped-query attention shares them across groups — 8 KV heads — for 1.34 GB; a single shared KV head (MQA) would need 0.17 GB. Fewer KV heads mean more requests per GPU, at a small cost in quality.

Where it misleads
  • Not every head matters: many can be pruned after training with little loss (Michel, Levy & Neubig 2019), so head count is not a measure of capability.
  • Stories about what a head “does” come from a few probes; most heads are hard to interpret.
  • Converting a trained multi-head model to GQA needs some extra training (“uptraining”).
In code
layers, d_k, bytes_per, ctx = 80, 128, 2, 4_096      # Llama-2-70B shapes, fp16, 4k context
def kv_gb(kv_heads):
    return 2 * layers * kv_heads * d_k * bytes_per * ctx / 1e9   # K and V, every layer, every token
mha, gqa, mqa = kv_gb(64), kv_gb(8), kv_gb(1)         # 64 query heads; 64, 8 or 1 KV heads
Python, in your browser — edit the code and run it again

Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.

The same idea with a library
import torch, torch.nn as nn, torch.nn.functional as F

class GQA(nn.Module):
    def __init__(self, d, n_heads, n_kv_heads):
        super().__init__()
        self.n_heads, self.n_kv = n_heads, n_kv_heads
        self.d_k = d // n_heads
        self.W_q = nn.Linear(d, n_heads * self.d_k, bias=False)
        self.W_k = nn.Linear(d, n_kv_heads * self.d_k, bias=False)
        self.W_v = nn.Linear(d, n_kv_heads * self.d_k, bias=False)
        self.W_o = nn.Linear(n_heads * self.d_k, d, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_kv, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_kv, self.d_k).transpose(1, 2)
        # Repeat KV heads to match query heads
        r = self.n_heads // self.n_kv
        k = k.repeat_interleave(r, dim=1)
        v = v.repeat_interleave(r, dim=1)
        attn = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.W_o(attn.transpose(1,2).reshape(B, T, -1))
Sources
  • N. Shazeer, “Fast Transformer Decoding: One Write-Head Is All You Need”, arXiv:1911.02150, 2019
  • J. Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, EMNLP, 2023
  • P. Michel, O. Levy & G. Neubig, “Are Sixteen Heads Really Better than One?”, NeurIPS, 2019

Last reviewed October 2026 · How this site checks its content · Report a mistake

Linked from 1
← Previous
Self-Attention
Open in the full reader, with the topic sidebar