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]

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