Multi-Query Attention

Decoding is bound by the bandwidth needed to re-read the KV cache each step. Share one K/V head across all query heads and that cache shrinks by the head count.

[attention][kv-cache]Stale

Standing

StaleLoad-bearing for understanding how the field arrived here, and replaced in practice by something on this list. Worth reading, not worth reaching for.

Kept for the argument it makes rather than for use: grouped-query attention recovers most of the quality it gave up, for most of the same saving.

judged as of 2026-09 · what the labels mean

Theory

At batch size 1 an autoregressive step performs a handful of matrix–vector products and re-reads the entire KV cache. Arithmetic intensity is therefore near 1, and the step time is a division: bytes read over memory bandwidth.

\text{bytes}_{KV} = 2 \cdot L \cdot n \cdot h_{kv} \cdot d_{\text{head}} \cdot b \cdot \text{sizeof(dtype)}

MQA sets $h_{kv} = 1$ ; grouped-query attention interpolates, sharing one K/V head per group of $g$ query heads. GQA at $g = 8$ recovers essentially all of multi-head quality at an eighth of the cache, which is why it, rather than MQA, is what shipped.

\underbrace{h_{kv} = h}_{\text{MHA}} \qquad \underbrace{h_{kv} = h/g}_{\text{GQA}} \qquad \underbrace{h_{kv} = 1}_{\text{MQA}}

Implementation

python · torch ≥ 2.1

import torch
from torch import Tensor
import torch.nn.functional as F


def gqa(q: Tensor, k: Tensor, v: Tensor, mask=None) -> Tensor:
    """q: [B, H, N, D]   k, v: [B, H_kv, M, D]   H % H_kv == 0"""
    h, h_kv = q.size(1), k.size(1)
    if h != h_kv:
        repeat = h // h_kv
        k = k.repeat_interleave(repeat, dim=1)
        v = v.repeat_interleave(repeat, dim=1)
    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

repeat_interleave here is expository — it materialises the expansion. In a real kernel the grouping is expressed as a stride of zero over the head dimension, so the shared K/V head is read once and broadcast, which is the entire point of the technique.

KV cache (MHA)

32.0 GB

GQA, g = 8

4.0 GB

MQA

0.5 GB

70B model, 32k context, batch 16, fp16

flash-attention scaled-dot-product-attention speculative-decoding paged-attention

References

[1]Shazeer — Fast Transformer Decoding: One Write-Head is All You Need (2019)arXiv:1911.02150

[2]Ainslie et al. — GQA: Training Generalized Multi-Query Transformer Checkpoints (2023)arXiv:2305.13245

[3]Kwon et al. — Efficient Memory Management for LLM Serving with PagedAttention (2023)arXiv:2309.06180