Grouped-Query Attention

Autoregressive decoding is bound by memory bandwidth, and the memory in question is the KV cache. Multi-query attention shrinks it by a factor of $h$ — one key–value head for every query head — which is a large win and a measurable quality loss. Grouped-query attention puts a dial between the two.

The query projections are untouched: there are still $h$ of them, and each head still computes its own attention pattern. What is shared is the thing that gets cached. With $g$ groups the cache is $2 \cdot g \cdot n \cdot d_h$ per layer instead of $2 \cdot h \cdot n \cdot d_h$, a factor of $h/g$.

Why the loss is small

A head’s attention pattern is set by $W^Q_i {W^K_{g(i)}}^{\top}$. Sharing $W^K$ within a group does not force the patterns to agree, because each head keeps its own query projection and can still rotate its queries into a different part of the shared key space. What is genuinely lost is the ability of two heads in a group to read different value subspaces of the same position — the $OV$ paths are constrained in a way the $QK$ paths are not.

Ainslie et al. report that uptrained GQA at $g = 8$ recovers most of the quality gap to MHA while retaining most of the MQA speedup, which is the shape of result that made it the default. The comparison is at fixed $h$: GQA is not free, it is cheap.

Interaction with the cache

Because a group’s keys and values are one tensor, decoding becomes a batched matmul with a broadcast rather than a gather. That matters for paged caches, where the page is now shared by $h/g$ query heads and the eviction unit is correspondingly coarser.

import torch
from torch import Tensor, nn


class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int) -> None:
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.h, self.g = n_heads, n_kv_heads
        self.d_head = d_model // n_heads
        self.q = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.kv = nn.Linear(d_model, 2 * n_kv_heads * self.d_head, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: Tensor, causal: bool = True) -> Tensor:
        b, n, _ = x.shape
        q = self.q(x).view(b, n, self.h, self.d_head).transpose(1, 2)
        k, v = (
            t.view(b, n, self.g, self.d_head).transpose(1, 2)
            for t in self.kv(x).chunk(2, dim=-1)
        )

        # Broadcast each KV head to the query heads of its group.
        k = k.repeat_interleave(self.h // self.g, dim=1)
        v = v.repeat_interleave(self.h // self.g, dim=1)

        out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, is_causal=causal
        )
        out = out.transpose(1, 2).reshape(b, n, self.h * self.d_head)
        return self.proj(out)

repeat_interleave allocates, which defeats the purpose during decoding. The version that matters lives inside the attention kernel, where the KV head index is computed from the query head index and the cache is read once rather than copied — the code above is the readable statement of what that kernel does.