Multi-Query Attention (MQA)

Standard multi-head attention (MHA) projects the input into H independent sets of queries, keys, and values. Each head has its own learned weight matrices W_Q_h, W_K_h, W_V_h, each of shape (d_model, d_head), where $d_{\text{head}} = \frac{d_{\text{model}}}{H}$. Multi-query attention (Shazeer 2019) keeps all H query projections but collapses keys and values to a single shared head: one W_K and one W_V, each of shape (d_model, d_head), broadcast to all query heads.

The weight structure

In MHA with $d_{\text{model}} = 4096$ and H = 32 heads ($d_{\text{head}} = 128$), the projection parameters are:

• Q: 32 matrices of shape (4096, 128) — or equivalently one (4096, 4096) matrix. 16.8M parameters.
• K: 32 matrices of shape (4096, 128) — 16.8M parameters.
• V: 32 matrices of shape (4096, 128) — 16.8M parameters.

Total QKV parameters: ~50.3M.

In MQA, Q stays the same — 32 heads, 16.8M parameters. But K and V each have a single (4096, 128) projection: 0.52M parameters each. Total QKV parameters: ~17.9M. The K and V parameter count drops by 32x.

During the forward pass, each query head q_h attends to the same shared keys and values:

# MQA forward pass (simplified)
# q: (batch, seq, H, d_head) — per-head queries
# k: (batch, seq, 1, d_head) — single shared key head
# v: (batch, seq, 1, d_head) — single shared value head
# k and v are broadcast across the H dimension

scores = einsum("b s h d, b t 1 d -> b h s t", q, k) / sqrt(d_head)
attn = softmax(scores, dim=-1)
out = einsum("b h s t, b t 1 d -> b s h d", attn, v)

The single key and value head is broadcast — every query head sees identical K and V tensors. The different learned W_Q_h matrices are the only thing differentiating heads.

The KV cache problem MQA solves

During autoregressive decoding, the model generates one token at a time. At each step, it needs the keys and values for every previous token — the KV cache. In MHA, each new token appends H key vectors and H value vectors to the cache.

KV cache size per token in MHA:

2 * H * d_head * bytes_per_element

For a 32-head model with $d_{\text{head}} = 128$ in float16:

At a sequence length of 8192 tokens:

Serving 64 concurrent sequences: $128 \text{MB} \cdot 64 = 8.2 \text{GB}$ of GPU memory consumed by KV cache alone.

MQA stores only 1 key head and 1 value head per token:

That's a 32x reduction. The same 64 concurrent sequences at 8192 tokens now cost 256 MB instead of 8.2 GB. On a 40 GB A100, that frees ~8 GB for larger batch sizes, directly increasing serving throughput.

The KV cache also dominates the memory-bandwidth cost of each decode step. At each token generation, the model reads the entire KV cache to compute attention. With MHA, reading 128 MB of KV cache per sequence through HBM (A100: ~2 TB/s bandwidth) takes $128 \frac{\text{MB}}{2} \frac{\text{TB}}{s} = 64 μs$ — just for the memory read, before any compute. With MQA, the read is $4 \frac{\text{MB}}{2} \frac{\text{TB}}{s} = 2 μs$. At high batch sizes where decode is memory-bound, this directly translates to lower per-token latency.

PyTorch implementation

Here's a complete implementation of both MHA and MQA in the same framework, showing the structural difference:

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, S, _ = x.shape
        q = self.W_Q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        k = self.W_K(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        v = self.W_V(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.W_O(out)


class MultiQueryAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        # Single shared K and V projections
        self.W_K = nn.Linear(d_model, self.d_head, bias=False)
        self.W_V = nn.Linear(d_model, self.d_head, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, S, _ = x.shape
        q = self.W_Q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        k = self.W_K(x).view(B, S, 1, self.d_head).transpose(1, 2)
        v = self.W_V(x).view(B, S, 1, self.d_head).transpose(1, 2)

        # k and v broadcast from (B, 1, S, d_head) to (B, n_heads, S, d_head)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.W_O(out)

The only structural change is the K and V linear layers: nn.Linear(d_model, d_model) becomes nn.Linear(d_model, d_head). Everything else — the attention computation, the softmax, the output projection — is identical. PyTorch's broadcasting handles the rest: when q has shape (B, 32, S, 128) and k has shape (B, 1, S, 128), torch.matmul automatically expands k across all 32 heads.

Parameter count comparison

d_model = 4096
n_heads = 32

mha = MultiHeadAttention(d_model, n_heads)
mqa = MultiQueryAttention(d_model, n_heads)

mha_params = sum(p.numel() for p in mha.parameters())
mqa_params = sum(p.numel() for p in mqa.parameters())

print(f"MHA: {mha_params:,} parameters")   # 67,108,864 (64M)
print(f"MQA: {mqa_params:,} parameters")   # 34,603,008 (33M)
print(f"Reduction: {1 - mqa_params/mha_params:.1%}")  # 48.4%

The 48% parameter reduction comes entirely from the K and V projections. The Q and output projections — which together account for half the attention parameters — are unchanged.

Quality cost

The single KV head is a compression bottleneck. In MHA, each query head attends to its own specialized key-value representation — head 3 might learn to encode syntactic relationships while head 17 encodes coreference. In MQA, all 32 query heads compete to extract different information from the same shared key-value embedding.

Shazeer (2019) reported MQA quality on WMT14 En-De translation:

• MHA baseline: 28.4 BLEU
• MQA: 28.1 BLEU (−0.3)
• MQA inference speedup: 1.7x (encoder-decoder Transformer)

A 0.3 BLEU drop on translation is minor, but the gap widens on tasks that require the model to simultaneously track multiple types of information across long contexts — multi-step reasoning, complex code generation, long-document summarization. The single KV head must encode everything every query head needs, and some information is inevitably lost in the compression.

Empirical evidence from deployed models

Falcon-7B (Technology Innovation Institute, 2023) uses MQA with $d_{\text{model}} = 4544$, 71 query heads, $d_{\text{head}} = 64$, and 1 KV head. Its KV cache per token is $2 \cdot 1 \cdot 64 \cdot 2 = 256 \text{bytes}$ versus $2 \cdot 71 \cdot 64 \cdot 2 = 18, 176 \text{bytes}$ for an equivalent MHA model — a 71x reduction.

PaLM (Chowdhery et al., 2022) uses MQA across all model sizes. The 540B variant has 48 heads with $d_{\text{head}} = 256$ and 1 KV head. At 2048-token sequence length in float16, MQA KV cache: $2048 \cdot 2 \cdot 1 \cdot 256 \cdot 2 = 2 \text{MB}$. MHA would need $2048 \cdot 2 \cdot 48 \cdot 256 \cdot 2 = 96 \text{MB}$. Per sequence, not per batch.

StarCoder (Li et al., 2023) uses MQA in its 15.5B parameter code model. The choice was motivated by code generation latency: developers expect near-instant completions, and the KV cache reduction lets StarCoder serve longer contexts (8192 tokens) without saturating GPU memory.

When MQA is the right choice

MQA makes sense under specific deployment constraints:

• Edge deployment — Mobile or embedded inference where GPU memory is measured in single-digit gigabytes. A 7B MQA model can serve 4K-token sequences on a 6 GB GPU; the MHA equivalent cannot.
• Extreme batch sizes — When you need to serve hundreds of concurrent requests and the KV cache is the memory bottleneck. MQA lets you scale batch size 32x before hitting the same memory wall.
• Latency-critical serving — The KV cache reduction directly speeds up the memory-bound decode step. On a single A100, MQA can cut per-token latency by 30–50% at high batch sizes.

MQA is the wrong tradeoff when quality on complex reasoning tasks is paramount and memory is not the bottleneck. For most modern large-scale deployments, grouped query attention (GQA) has replaced MQA as the default — it recovers most of the quality while keeping most of the memory savings. Llama 2 70B, Llama 3, Mistral, and Gemma all chose GQA over MQA. The shift happened because GQA-8 (8 KV heads) recovers most of MQA's memory advantage — KV cache drops to 25% of MHA versus MQA's ~3% — while keeping 8 independent KV representations that better serve diverse query heads.