23. Grouped-Query Attention (Fewer KV Heads Than Q Heads)

Problem

A KV-cache's memory footprint scales with the number of key/value heads, since every head needs its own cached K and V tensor per token generated. In ordinary multi-head attention, num_kv_heads == num_q_heads, so cache memory grows linearly with the number of query heads too — for long contexts or many concurrent sessions (as in a real-time serving system handling many simultaneous streams), that memory pressure can be the actual bottleneck, not compute.

Grouped-query attention (GQA) breaks that coupling: keep the full number of query heads (representational capacity unchanged), but use far fewer key/value heads, with each KV head shared across a group of query heads. This shrinks the KV-cache proportionally to num_q_heads / num_kv_heads. (Multi-query attention, MQA, is the extreme case of one shared KV head; ordinary multi-head attention is the other extreme, group size 1.)

Source: src/23_grouped_query_attention.py

def grouped_query_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    num_q_heads: int, num_kv_heads: int,
) -> torch.Tensor: ...

>>> q = torch.randn(2, 5, 8)   # 4 query heads of dim 2
>>> k = v = torch.randn(2, 5, 4)  # 2 kv heads of dim 2
>>> grouped_query_attention(q, k, v, num_q_heads=4, num_kv_heads=2).shape
torch.Size([2, 5, 8])

Step-by-Step Approach

1. Split q into num_q_heads heads and k, v into num_kv_heads heads, all with the same head_dim — same reshape/transpose mechanic as ordinary multi-head attention, just with a different head count per tensor.
2. Compute group_size = num_q_heads // num_kv_heads — the number of query heads that will share each kv head. Assert num_q_heads % num_kv_heads == 0 up front.
3. Expand the kv heads to line up with the query heads using repeat_interleave(group_size, dim=1) on the head axis, not a plain repeat or tile — interleave-repeat keeps consecutive query heads [0, 1, ..., group_size-1] mapped to kv head 0, the next group_size query heads mapped to kv head 1, and so on.
4. Run standard scaled dot-product attention on the now-equal-head-count tensors: qh @ kh.transpose(-1,-2) / sqrt(head_dim), softmax, @ vh.
5. Merge heads back exactly as in plain multi-head attention: transpose the head axis next to sequence, .contiguous(), then view to (batch, seq, num_q_heads * head_dim).
6. Sanity-check the degenerate case: setting num_kv_heads == num_q_heads must reproduce ordinary multi-head attention exactly, since group size 1 means no sharing actually happens.

The key insight is that GQA is not a new attention mechanism, it's multi-head attention with the kv-side head count decoupled from the q-side head count, and the entire implementation delta from plain MHA is one repeat_interleave call that broadcasts each kv head across its group before the identical attention math runs.

Reference solution

import math
import torch


def _split_heads(t: torch.Tensor, num_heads: int) -> torch.Tensor:
    batch, seq_len, dim = t.shape
    head_dim = dim // num_heads
    return t.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)


def grouped_query_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    num_q_heads: int, num_kv_heads: int,
) -> torch.Tensor:
    assert num_q_heads % num_kv_heads == 0
    head_dim = q.shape[-1] // num_q_heads
    group_size = num_q_heads // num_kv_heads  # q heads sharing each kv head

    qh = _split_heads(q, num_q_heads)     # (batch, num_q_heads, seq, head_dim)
    kh = _split_heads(k, num_kv_heads)    # (batch, num_kv_heads, seq, head_dim)
    vh = _split_heads(v, num_kv_heads)

    # broadcast each kv head across its group: heads [0..group_size) of q
    # all attend against kv head 0, the next group against kv head 1, etc.
    kh = kh.repeat_interleave(group_size, dim=1)  # (batch, num_q_heads, seq, head_dim)
    vh = vh.repeat_interleave(group_size, dim=1)

    scores = qh @ kh.transpose(-1, -2) / math.sqrt(head_dim)
    attn_weights = torch.softmax(scores, dim=-1)
    out_h = attn_weights @ vh

    b, h, s, d = out_h.shape
    return out_h.transpose(1, 2).contiguous().view(b, s, h * d)

Key Functions & Tricks

• tensor.repeat_interleave(group_size, dim=1) — the crux of GQA: expands each kv head into group_size consecutive copies along the head axis so it lines up positionally with the query heads that should share it.
• repeat_interleave vs. repeat/tile — repeat would tile the whole sequence of kv heads group_size times ([0,1,0,1]), which pairs the wrong query heads with the wrong kv heads; repeat_interleave gives [0,0,1,1], the correct grouping.
• num_q_heads % num_kv_heads == 0 assertion — GQA requires the query head count to be an exact multiple of the kv head count so every group is the same size.
• Equivalence-to-MHA check when num_kv_heads == num_q_heads — the strongest correctness test, since it confirms the grouping logic collapses to a no-op in the degenerate case rather than silently reshuffling heads.

How to Recognize This Pattern