28. Grouped-Query Attention (GQA)

Problem

Multi-head attention stores one K/V head per Q head; multi-query attention (MQA) goes to the other extreme and shares a single K/V head across every Q head, which is cheap to cache but hurts quality. Grouped-query attention sits in between: Q heads are split into groups, and every head within a group shares one K/V head. Mistral's 7B and Mixtral models both use GQA (alongside sliding-window attention) specifically because it shrinks the KV-cache — the dominant memory cost during long-context inference — while staying much closer to full multi-head quality than MQA.

Given q already split into num_q_heads heads, and k/v split into fewer num_kv_heads heads (num_q_heads must be an exact multiple of num_kv_heads), compute attention where each group of num_q_heads // num_kv_heads query heads reuses the same k/v head.

Source: src/28_grouped_query_attention.py

def grouped_query_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    mask: torch.Tensor | None = None,
) -> torch.Tensor: ...

>>> q = torch.randn(1, 4, 5, 4)
>>> k = v = torch.randn(1, 2, 5, 4)
>>> grouped_query_attention(q, k, v).shape
torch.Size([1, 4, 5, 4])

Step-by-Step Approach

1. Read num_q_heads from q.shape[1] and num_kv_heads from k.shape[1]; compute group = num_q_heads // num_kv_heads.
2. Expand k and v from num_kv_heads to num_q_heads with k.repeat_interleave(group, dim=1) — this duplicates each kv head group times contiguously, so heads [0, group) all map to kv head 0, heads [group, 2*group) to kv head 1, and so on.
3. Run ordinary scaled dot-product attention on the now-matching-width tensors: q @ k_exp.transpose(-1, -2) / sqrt(head_dim), softmax, then @ v_exp.
4. If a mask is given as (batch, seq_q, seq_k), unsqueeze(1) it so it broadcasts identically across every query head, exactly as in ordinary multi-head attention.
5. Sanity-check with num_kv_heads == num_q_heads: group becomes 1, repeat_interleave(1, ...) is a no-op, and the result must equal ordinary multi-head attention exactly.

The key insight is that GQA needs no new attention math at all — once k/v are expanded back up to num_q_heads via repeat_interleave, every subsequent line is identical to ordinary multi-head attention. The entire technique lives in how few distinct k/v heads get stored and cached, not in what happens to them once expanded.

Reference solution

import math
import torch


def grouped_query_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    mask: torch.Tensor | None = None,
) -> torch.Tensor:
    batch, num_q_heads, seq, hd = q.shape
    num_kv_heads = k.shape[1]
    group = num_q_heads // num_kv_heads

    # repeat_interleave (not repeat/tile) so heads [0, group) all map to kv
    # head 0, heads [group, 2*group) all map to kv head 1, etc. -- matching
    # the "contiguous group shares one kv head" definition of GQA
    k_exp = k.repeat_interleave(group, dim=1)  # (batch, num_q_heads, seq, hd)
    v_exp = v.repeat_interleave(group, dim=1)

    scores = q @ k_exp.transpose(-1, -2) / math.sqrt(hd)
    if mask is not None:
        scores = scores.masked_fill(~mask.unsqueeze(1), float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return weights @ v_exp

Key Functions & Tricks

• tensor.repeat_interleave(group, dim=1) — duplicates each kv head group times contiguously; this is the piece that must not be confused with tensor.repeat, which would tile the whole sequence of heads instead of duplicating each one in place.
• q.shape[1] // k.shape[1] — recovers the group size directly from tensor shapes rather than requiring it as a separate argument.
• mask.unsqueeze(1) — inserts the query-heads axis into a per-batch mask so it broadcasts across all num_q_heads, unchanged from ordinary MHA masking.
• Batched @ over 4D tensors — once k/v are expanded to num_q_heads, the matmul shapes are identical to plain multi-head attention.

How to Recognize This Pattern