Multi-Query Attention: Fast Transformer Decoding: One Write-Head is All You Need

Let all attention heads share one key and one value, and the cache a decoder rereads at every step shrinks by the number of heads.

A Transformer generates text one token at a time, and each step reads back the stored keys and values of every earlier token. Noam Shazeer kept eight query heads but gave them a single shared key and value, which cut a translation model's decoder time from 46 to 3.8 microseconds per token for a 0.2 BLEU drop.

A November 2019 arXiv paper with three tables and a code change you can describe in one sentence: delete the letter h from the key and value projections.

Multi-query attention (MQA) is a variant of the attention layer in the Transformer. In standard multi-head attention, each of the $h$ heads computes its own query, its own key and its own value from the input. In multi-query attention the heads keep their own queries but share a single key and a single value. The paper explains why that matters when a model generates text: it stores the keys and values of every past token in a cache and reads that cache back at every step, and the time of a step is set by how many bytes it reads, not by how much it multiplies. Sharing K and V makes the cache $h$ times smaller.

The sections go in this order: first the shapes of multi-head attention, then what a decoder stores and rereads when it generates one token at a time, then the paper's count of bytes moved per arithmetic operation, which shows why that rereading dominates. Then the one-letter change, what it does to the count, and the paper's measurements of quality and speed. The last section follows the idea into grouped-query attention, the 2023 generalization used by Mistral 7B and Llama 2 70B.

Multi-head attention, in shapes

An attention function takes one query vector $q$ and $m$ key-value pairs, stacked as matrices $K$ and $V$. It scores the query against every key, turns the scores into weights that sum to one, and returns the weighted sum of the values:

A key that points the same way as the query gets a large dot product, so its value gets a large weight. The original Transformer also divides the scores by $\sqrt{d_k}$; the paper leaves that constant out of its code because it can be folded into the projection that produces $q$ or $K$.

Multi-head attention runs $h$ of these in parallel. The paper writes it with einsum, where each letter names a tensor axis: b the batch of sequences, n the query positions, m the memory positions being attended to, d the model width, h the heads, and k, v the width of one head's keys and values. Any letter missing from the output of an einsum is summed over.

# multi-head attention, batched (paper, section 2.3)
# X: [b, n, d] queries' inputs;  M: [b, m, d] memory;  P_q, P_k: [h, d, k]
Q = einsum("bnd,hdk->bhnk", X, P_q)     # h query heads
K = einsum("bmd,hdk->bhmk", M, P_k)     # h key heads
V = einsum("bmd,hdv->bhmv", M, P_v)     # h value heads
logits  = einsum("bhnk,bhmk->bhnm", Q, K)
weights = softmax(logits + mask)
O = einsum("bhnm,bhmv->bhnv", weights, V)
Y = einsum("bhnv,hdv->bnd", O, P_o)     # per-head output projection, summed

In the paper's translation model, $d = 1024$, $h = 8$ and $k = v = 128$, so $h k = d$. Each projection tensor P_q, P_k, P_v, P_o has shape $[8, 1024, 128]$, about a million weights. For one memory position, the eight key heads together produce 8 × 128 = 1024 numbers and the eight value heads another 1024. The head axis h appears in K and V, and that axis is what the paper removes.

Decoding one token at a time: the KV cache

During training the whole target sentence is known, so all $n$ query positions go through the layer in one call, as in the listing above, with a mask that stops each position from attending to later ones. Generation cannot do this. The token at position $i$ is an input at position $i+1$, so the model has to run one position per call, $n$ calls in a row.

At each call, the new position needs keys and values for every earlier position. Those were computed at earlier calls and do not change, because a position's key depends only on the tokens up to it. So the decoder keeps them in a KV cache: each call computes one new key and one new value, appends them, and then attends over the whole cache. The query is used once and discarded.

The cache holds, for each of $L$ layers and each cached token, one key and one value per head:

For the paper's decoder ($h = 8$, $d_k = 128$, $L = 6$) at its full target length of 128 tokens, the self-attention cache is 2 × 8 × 128 × 6 × 128 = 1,572,864 numbers per sentence. The decoder's encoder-decoder attention also reads keys and values of the 128 source tokens at every step, the same amount again, and the timing runs decode 1024 sentences at once. For a larger model with 32 layers, 32 heads of width 128 and a 32,000-token context, (2) gives 8,388,608,000 numbers, 16.8 GB at two bytes each, for a single sequence.

The paper's step function for multi-head attention takes prev_K of shape $[b, h, m, k]$, appends the new position along the m axis, and computes the logits with einsum("bhk,bhmk->bhm", q, K). Every number in that cache is read once per step and enters exactly one multiply-add: head 3's query multiplies head 3's keys and nothing else does.

Why decoding is memory-bound

A chip can be slow for two reasons. It can be busy multiplying, or it can be idle waiting for operands to arrive from main memory. Which one applies depends on the arithmetic intensity of the work, the number of operations done per byte read. A recent datacenter GPU can do on the order of 150 floating-point operations in the time it reads one byte from its memory (the paper says compute capacity can be "two orders of magnitude higher than the memory bandwidth"). Work that does fewer operations per byte than that leaves the arithmetic units waiting, and its speed is set by bytes moved.

Section 2.3.1 of the paper counts both quantities for the batched layer, under three simplifying assumptions: as many memory positions as queries ($m = n$), heads that split the model width ($k = v = d/h$), and sequences no longer than the width ($n \le d$). Every einsum costs $O(bnd^2)$ operations, for $\Theta(bnd^2)$ in total. Memory traffic is the size of every tensor touched, $O(bnd + bhn^2 + d^2)$: the activations, the logits and weights, and the four projection tensors. Dividing:

With $k = 128$, that is under one number moved per hundred operations, because each key is loaded once and used by all $n$ queries. Section 2.4.1 redoes the count for the step-by-step decoder. Across $n$ calls the arithmetic is the same $\Theta(bnd^2)$, but the memory is $\Theta(bn^2 d + nd^2)$:

Each term comes from one tensor. At every call the decoder rereads $K$ and $V$, which have $bhmk$ = $bnd$ numbers, while it does about $bd^2$ operations on the projections; that ratio is $n/d$. It also reloads the four projection matrices, $d^2$ numbers each, which are shared by the $b$ sequences in the batch; that ratio is $1/b$. The activations $x, q, o, y$ add $bd$ numbers per call, a $1/d$ term that is smaller than $n/d$ and absent from (4) for that reason.

The two terms respond to different fixes. A bigger batch shrinks $1/b$, since one load of the weights then serves more sequences; the paper calls this term "the easier one". A bigger batch does nothing for $n/d$, because every sequence has its own cache. With the paper's numbers, $n = 128$ and $d = 1024$, the cache term is 0.125, sixteen times the training ratio $1/k$. The paper notes that when $n \approx d$ or $b \approx 1$ the ratio is close to 1, and that efficient generation needs both terms far below 1.

Figure 2 counts one decode step of a 32-layer model with 32 heads of width 128, attention only (the cache reads and the multiplies on them). Drag the cache length and watch the two bars grow together, so the intensity gauge does not move:

The 1 FLOP per byte follows from the step function above. For one cached token and one head, the step reads 128 key numbers and 128 value numbers at two bytes each, 512 bytes, and does 128 multiply-adds for the score and 128 for the weighted value, at two FLOPs each, 512 FLOPs. No batch size or cache length changes that ratio, because both sides scale with the number of tokens and the number of sequences.

How multi-query attention works

If the cache sets the speed, there are two ways to read less of it. Earlier work attended to fewer positions: local attention keeps only a window of recent tokens, other methods compress the memory. Shazeer takes the other factor in $bhmk$ and removes the $h$: the heads keep separate queries and share one key and one value. In the paper's words, the code is identical "except that we remove the letter ‘h’ from the tf.einsum equations where it represents the ‘heads’ dimension of K, V, P_k, or P_v”.

# multi-query attention, batched (paper, section 3)
# P_q, P_o keep the head axis: [h, d, k];  P_k: [d, k];  P_v: [d, v]
Q = einsum("bnd,hdk->bhnk", X, P_q)     # h query heads, as before
K = einsum("bmd,dk->bmk",   M, P_k)     # no h: one key per position
V = einsum("bmd,dv->bmv",   M, P_v)     # no h: one value per position
logits  = einsum("bhnk,bmk->bhnm", Q, K)
weights = softmax(logits + mask)
O = einsum("bhnm,bmv->bhnv", weights, V)
Y = einsum("bhnv,hdv->bnd", O, P_o)

Compare it line by line with the multi-head listing. P_k goes from $[h, d, k]$ to $[d, k]$, so K goes from $[b, h, m, k]$ to $[b, m, k]$, and the same for V. The logits einsum bhnk,bmk->bhnm still produces one score matrix per head, because Q still has the head axis. P_q and P_o are unchanged.

Why keep the head axis on the query side? The keys and values are what the decoder stores and rereads; the query is computed fresh at each step and used once, so per-head queries cost almost nothing at decode time. The step function from Section 3, with three typos in the printed listing fixed (see the Provenance panel):

# multi-query self-attention, one decode step (paper section 3, typos fixed)
# x: [b, d] the newest position;  prev_K: [b, m, k];  prev_V: [b, m, v]
q = einsum("bd,hdk->bhk", x, P_q)               # h queries for this position
k_new = einsum("bd,dk->bk", x, P_k)             # one key for this position
v_new = einsum("bd,dv->bv", x, P_v)             # one value
K = concat([prev_K, k_new[:, None, :]], axis=1) # cache grows to [b, m+1, k]
V = concat([prev_V, v_new[:, None, :]], axis=1)
logits  = einsum("bhk,bmk->bhm", q, K)          # all h heads read the same K
weights = softmax(logits)
o = einsum("bhm,bmv->bhv", weights, V)          # ... and the same V
y = einsum("bhv,hdv->bd", o, P_o)               # [b, d]
return y, K, V                                  # K, V become next step's cache

With the paper's timing batch ($b = 1024$, 128 cached positions, $k = v = 128$, $h = 8$), the multi-head step reads 2 × 1024 × 8 × 128 × 128 = 268,435,456 cached numbers per layer, and the multi-query step reads 33,554,432. Both do the same 268,435,456 multiply-adds on them, so each cached number now feeds $h = 8$ multiply-adds instead of one: the logits einsum bhk,bmk->bhm multiplies every element of K by eight different queries.

Toggle Figure 3 between the three settings. The query heads stay; the key/value heads they read from collapse from one per query to one for all:

Sharing K and V restricts what the layer can represent. In multi-head attention, head $i$ projects each memory position to its own 128-number key and 128-number value, so the eight heads together can read 1024 numbers' worth of different features from every position. In multi-query attention all eight heads read the same 128-number key and the same 128-number value. They can still weight different positions, because their queries differ, and they can still use what they read differently, because each has its own slice of P_o. What disappears is the variety in what is stored about each position. Whether that matters is an empirical question, and Section 4 answers it for two tasks.

Dropping the head axis also removes weights. Each of P_k and P_v shrinks from 8 × 1024 × 128 to 1024 × 128 (917,504 fewer weights each), so 1,835,008 fewer weights per attention layer. The translation model has 18 attention layers (6 in the encoder, 12 in the decoder counting self-attention and encoder-decoder attention), so multi-query loses 33,030,144 weights. The paper widens the feed-forward hidden layer from 4096 to 5440 to put them back: 1344 extra units, times two matrices of width 1024, times 12 feed-forward layers, is also 33,030,144. The quality comparison below is between models of equal size, 211 million parameters.

What sharing K and V does to the cost

Section 3.1 repeats the count with the new shapes. The arithmetic is unchanged at $\Theta(bnd^2)$, since the eight query heads do all the multiplies they did before. The memory across $n$ calls becomes $\Theta(bnd + bn^2k + nd^2)$, where the $bnd$ is $x, q, o, y$, the $bn^2k$ is the shared $K$ and $V$, and the $nd^2$ is the projection weights. Dividing, with $k = d/h$:

Against (4), the cache term $n/d$ became $n/(dh)$, and the paper puts it as "We have reduced the offensive $n/d$ by a factor of $h$." The other terms stay as they were. The $1/b$ term is the same weight loading, and the $1/d$ term is the same $x, q, o, y$ traffic that (4) absorbed into $n/d$; with $n/d$ divided by $h$, it is no longer negligible beside the cache term, so (5) writes it out.

At the paper's timing setup, $d = 1024$, $n = 128$, $b = 1024$, $h = 8$:

The cache term falls by exactly 8 and the total by about 7.2. An exact count of one layer's step gives a similar factor: 276,824,064 numbers read under multi-head (97% of it the cache) against 40,108,032 under multi-query, 6.9 times fewer. Set the sliders in Figure 4 to see where the factor $h$ helps and where it does not:

Two cases follow from the sliders. For long sequences at large batch, the regime of a busy serving system, nearly all traffic is the cache and the saving approaches $h$. With one sequence at a time and a short context, the step is dominated by loading the weights, which multi-query attention does not touch; the paper's own statement is conditional ("given large batch size b, this should dramatically improve performance").

What it costs in quality

Section 4 trains on two tasks. The first is WMT 2014 English-to-German translation, scored by BLEU (0 to 100, word-overlap with reference translations, higher is better) and by perplexity on the development set. The baseline is an encoder-decoder Transformer with 6 layers, $d_{\text{model}} = 1024$, $d_{ff} = 4096$, $h = 8$, $d_k = d_v = 128$. In the multi-query model every attention layer is multi-query: encoder self-attention, decoder self-attention and encoder-decoder attention.

Multi-query attention scores 26.5 dev BLEU against 26.7, and dev $\ln(\text{PPL})$ 1.439 against 1.424. On the test set it gets 27.5 with greedy decoding against 27.7, and 28.5 with beam search of width 4 against 28.4, the highest beam-4 score in the table; the paper describes the two models as performing "similarly" and does not claim an improvement. The abstract's summary is "only minor quality degradation".

A model with fewer heads would also have a smaller cache, so the paper trains that too. Four multi-head variants keep $h \cdot d_k = 128$ (1 head of 128, 2 of 64, 4 of 32, 8 of 16), which gives them the same 256 key and value numbers per token per layer as multi-query, against the baseline's 2048. Their Q and O projections shrink as well, so their feed-forward layers widen further, to 6784, to keep 211 million parameters (the same accounting as above, with 66,060,288 weights moved). The paper compares them on quality only and reports no decode times for them.

At the same K, V size, multi-query attention is 0.3 BLEU above the best variant (2 heads of 64, 26.2) and 0.7 above the worst two (1 head of 128, and 8 heads of 16, both 25.8). The second task, the Billion-Word language modeling benchmark with a 6-layer decoder-only model of 192 million parameters, has the same order: dev perplexity 29.9 for multi-head, 30.2 for multi-query, and 30.9 to 31.2 for the variants. The gap of 0.3 perplexity points between multi-head and multi-query is about 0.01 in log-perplexity (nats), since ln 30.2 − ln 29.9 ≈ 0.010. The paper calls multi-query "slightly worse than the baseline, but significantly better than any of the alternatives involving decreasing h, d_k and d_v".

Measured speed

Section 4.3 times the translation models on one TPUv2 with 8 cores. Training runs batches of 32,768 source and 32,768 target tokens; a step took 433 ms for multi-head and 425 ms for multi-query, 13.2 and 13.0 microseconds per token. Inference decodes a batch of 1024 sentences (128 per core), source and target length 128. The multi-head decoder took 47 ms per step for all 1024 sentences, or 46 microseconds per output token; the multi-query decoder took 3.9 ms per step, or 3.8 microseconds. Translating the whole batch, 128 decoder steps, takes about 6.0 seconds with the first and 0.50 with the second.

The greedy speedup of 12.1 is larger than $h = 8$, and larger than the factor of about 7 that the count in (5) predicts. The Θ analysis drops constant factors, so it predicts the direction and the rough size of the change, and the paper does not break the measured time down further. Beam search with 4 beams gains 6.3x (203 to 32 microseconds); the paper reports the numbers without explaining why that gain is smaller.

The local-attention rows test the earlier approach to the same cost. Restricting decoder self-attention to a window of 32 positions halves the multi-head greedy decoder (46 to 23 microseconds) and costs 0.1 dev BLEU. Adding multi-query on top brings it to 3.3 microseconds, 14 times faster than the full multi-head decoder, which is the paper's evidence that the two methods are orthogonal: one reduces the number of positions read, the other the size of each.

A footnote describes the timing implementation. The TPU implementation needed fixed tensor shapes, so the cache was padded to the maximum length of 128 (or to the window of 32 for local attention) and masked, which makes every decoding step take the same time. The paper notes that growing the tensors step by step could save time near the start of the sequence. The microsecond figures are therefore costs at full cache length, averaged over tokens.

Grouped-query attention and after

Multi-query attention has one key/value head; multi-head has $h$. Ainslie et al. (2023) introduced grouped-query attention (GQA), which uses $G$ key/value heads, each shared by a group of $h/G$ query heads, so multi-head is $G = h$ and multi-query is $G = 1$. Their abstract starts from multi-query's two practical problems: it "can lead to quality degradation", and a model has to be trained with it from the start. They show that a multi-head checkpoint can be converted with 5% of its original pre-training compute, and that grouped-query models come close to multi-head quality at close to multi-query speed. Mistral 7B uses 32 query heads and 8 key/value heads, a cache 4 times smaller than multi-head; Llama 2 70B uses 64 and 8, 8 times smaller. Both are smaller savings than multi-query's 32x and 64x would have been.

Other work attacks other parts of the same decode cost. FlashAttention (Dao et al., 2022) computes exact attention in tiles held in on-chip memory, so the matrix of scores is never written to GPU main memory; it reduces the traffic of one attention call, while multi-query attention reduces the size of the cache that call reads. PagedAttention (Kwon et al., 2023) allocates the cache in fixed-size blocks for a server handling many requests at once, after measuring that existing systems used only 20.4% to 38.2% of their KV-cache memory for actual token states. A smaller cache per token, as in multi-query and grouped-query attention, means more sequences fit in a GPU's memory, which is the large batch that the $1/b$ term in (5) needs.