Grouped-Query Attention and What Sharing Key Heads Costs

Multi-head attention gives each head its own query, key, and value projections. The key and value projections are what the cache must hold, so the cache cost scales with head count exactly as it scales with width. Multi-query attention removes that scaling by keeping one key head and one value head for all query heads, cutting the cached tensors by a factor of $h$. Grouped-query attention sits between the two: $h$ query heads are partitioned into $g$ groups, each group shares one key head and one value head, and the cache shrinks by $h/g$. The arithmetic of the saving is trivial. What the shared head gives up is not, and the asymmetry between what a query head loses and what a key head loses is the part worth working through.

The shape of the computation

Let the query projection be full width: $h$ heads of width $d_k$, total $h d_k = d_{\text{model}}$. Let the key and value projections produce $g$ heads of the same width $d_k$, total $g d_k$. Query head $i$ belongs to group $\lfloor i \cdot g / h \rfloor$, and it uses that group’s key and value head. The score for query head $i$ against key position $n$ is

$$s^{(i)}{m,n} = \frac{\left(x_m W_i^Q\right) \left(x_n W{g(i)}^K\right)^\top}{\sqrt{d_k}}$$

and the output of head $i$ is the softmax-weighted sum of the group’s value head. The $h$ outputs are concatenated to $h d_k$ columns and projected by $W^O$ as usual, so the output side of the layer is unchanged. Only the key and value side is narrowed.

The cache saving is exactly $h/g$, with no approximation: $2 L n g d_k s$ bytes instead of $2 L n h d_k s$. At $h = 32$, $g = 8$, $L = 32$, $d_{\text{model}} = 4096$, 16-bit, and 4,096 cached positions, the full multi-head cache is 2 GiB per sequence and the grouped version is 512 MiB. That is the difference between one sequence and four on a fixed device, or between a batch of 8 and a batch of 32.

The query projection is not reduced. The full $h d_k$ query width is kept, so the per-head score matrices still exist at full head count and the softmax is still computed per query head. What is shared is only what gets stored and read.

What a shared key head actually removes

The expressivity argument for multi-head attention rests on the softmax being computed separately per head. That property survives grouped-query attention: each query head has its own score row, its own normalization, and its own attention distribution. A query head can still place its mass wherever it likes, because its own query projection determines the scores.

What changes is the set of directions in which keys can differ. With $h$ independent key projections, the key at position $n$ is represented by $h$ vectors of width $d_k$, and two positions can be distinguished in any of the $h$ subspaces independently. With $g$ shared key heads, all queries in a group see the same $d_k$-dimensional key vector. Two positions whose key representations coincide in that subspace are indistinguishable to every query head in the group, no matter how their queries differ, because the score is an inner product of a query with that single key vector.

The practical consequence is that the head axis stops carrying independent key information. Stack the score matrices for all heads into a $h \times n \times n$ tensor. With independent keys, the $h$ slices are computed from $h$ unrelated key projections and are as independent as the data allows. With $g$ groups, only $g$ distinct key projections exist, so the slices within a group are all inner products against the same key matrix $R_{g(i)}$ and differ only through their queries $q_i$. Every difference between two heads in a group is therefore a difference the query side already had to be able to express.

Whether that matters depends on whether the task needs heads to disagree about which keys are relevant in a way that requires different key subspaces. Heads that differ only in what they do with a common relevance signal, such as one attending to the previous token, one to a syntactic head, and one to a rare token, can share a key projection and still disagree, because their queries are different. Heads that need to detect relevance along a dimension invisible to the shared key cannot.

That is the trade in its narrowest form, and it explains why the reported quality differences are small but not zero. The mechanism is not “fewer heads”; it is “fewer key subspaces”, and the number of heads that genuinely needed their own subspace is usually a small fraction of the total.

Why values tolerate sharing better than keys

Keys enter through the softmax, values enter linearly after it. An error in a key changes the score, which changes the weights through a nonlinearity, and the change is amplified or suppressed depending on where the logit sits relative to the row maximum. An error in a value is averaged with the other values and passes through a linear map.

This asymmetry is why quantization schemes usually treat keys more carefully than values, and it is also why sharing is more damaging on the key side than on the value side. A shared key that cannot separate two positions produces a softmax that must split mass between them; a shared value that conflates two positions produces an output that is a mixture, but the mixture is computed with weights the query head chose, so the query can compensate by weighting one position more heavily.

There is a second reason the value side tolerates sharing, and it comes from the output projection. With independent value heads the layer computes $\sum_i p_i V_i W^O_i$, where $W^O_i$ is the block of $W^O$ belonging to head $i$. Each head’s attention distribution is applied to its own value subspace before the projection, so two heads can route the same distribution through different directions and the projection can combine them differently. With a shared value head the expression collapses to $\sum_i p_i , V W^O_i$, which is a weighted sum of the same $n$ value vectors with an effective weight matrix $\sum_i p_i W^O_i$ of shape $d_k \times d_{\text{model}}$. The layer can still produce a rich output, but every head now reads from one subspace, and the freedom that separate value heads provided is gone.

The extreme case makes the limit concrete. With $g = 1$, the multi-query case, every head reads the same $n \times d_k$ value matrix $V$, so the $h$ blocks of the pre-output concatenation are $p_1V, \dots, p_hV$. Those blocks are not identical, since the attention rows $p_i$ differ, and the layer can still write a full $d_{\text{model}}$-wide vector to the residual stream. What it loses is variety in the value content: every block is a linear combination of the same $d_k$-dimensional value vectors, so there are only $d_k = d_{\text{model}}/h$ independent value directions available for routing, no matter how many heads read them. With $g$ groups the count rises to $g d_k$, which for $g = 8$ and $h = 32$ is a quarter of the model width, and that is why the middle ground is preferred to the extreme.

import numpy as np

def gqa(x, wq, wk, wv, wo, h, g):
    # x: (n, d_model); wq: (d_model, h*d_k); wk, wv: (d_model, g*d_k)
    n, d = x.shape
    d_k = wq.shape[1] // h
    q = (x @ wq).reshape(n, h, d_k)
    k = (x @ wk).reshape(n, g, d_k)
    v = (x @ wv).reshape(n, g, d_k)

    group = np.arange(h) * g // h            # query head -> its key/value group
    kk, vv = k[:, group], v[:, group]         # broadcast shared heads to h
    s = np.einsum('qhd,khd->hqk', q, kk) / np.sqrt(d_k)   # (h, n_q, n_k)
    s = s - s.max(axis=-1, keepdims=True)
    p = np.exp(s)
    p /= p.sum(axis=-1, keepdims=True)
    out = np.einsum('hqk,khd->qhd', p, vv)    # (n_q, h, d_k)
    return out.reshape(n, h * d_k) @ wo

Converting a trained multi-head checkpoint

A model trained with independent key and value heads has $h$ key matrices where a grouped model needs $g$. The conversion has to reduce that set, and the obvious approach of taking the first $g$ heads and discarding the rest throws away parameters that the model was using.

The method that works is to average within each group. For group $j$ covering query heads $j h/g \dots (j+1)h/g - 1$, the shared key matrix is the mean of those heads’ key matrices, and likewise for values. Averaging is the choice that preserves the group’s mean projection exactly: the sum over the group of $q_i^\top K_i$ becomes $(\sum_i q_i)^\top \bar{K}$, so the group’s aggregate behaviour is unchanged at the level of the mean query. It is not exact per head, because individual heads’ key matrices are replaced by their mean and their deviations from it are lost.

Uptraining then recovers most of the gap. The converted model is trained for a small fraction of the original training budget, enough for the query projections and the output projection to adapt to the shared keys and for the shared keys to move toward a compromise that serves all the heads in the group. The published procedure uses this two-stage approach and reports that a modest amount of continued training closes most of the difference on standard benchmarks; the specific figures depend on the model and the benchmark, and the mechanism is what transfers.

The averaging step has a subtlety worth flagging: it is done on the projection matrices, not on the per-head norms. If one head’s key matrix has a much larger norm than its group-mates, which happens when that head learned to produce sharply peaked scores, averaging pulls its norm down and softens its attention. Compensating by rescaling the queries of that head is possible, since only the product $q^\top k$ matters, but the rescaling interacts with the $1/\sqrt{d_k}$ divisor and with layer normalization placement. Implementations that do the conversion usually leave the scale alone and let uptraining fix it, which is why the conversion is described as requiring training rather than as a pure weight transformation.

Where the configuration lands

The choice of $g$ is a memory-versus-quality dial with a specific shape. Cache size falls as $1/g$; the score tensor’s head-axis rank falls from $h$ to $g$; the query-side compute is unchanged; and the key/value projection parameter count falls from $h d_{\text{model}} d_k$ to $g d_{\text{model}} d_k$. Since the projections are a small fraction of total parameters, the parameter saving is minor and the cache saving is the point.

The values that appear in practice cluster around $g$ equal to the number of attention head groups that map naturally onto the number of tensor-parallel partitions, because each partition holds a subset of the key heads. A model sharded across eight devices with $g = 8$ puts one key head on each device, which keeps the cache read local and avoids a collective per decode step. Choosing $g$ smaller than the partition count leaves some devices with no key head to compute, which forces either replication or a different sharding scheme. This is a systems constraint rather than a modelling one, and it is the reason $g$ is often a power of two.

The interaction with rotary embeddings is benign, as noted in rotary position embeddings and relative offsets: the rotation is applied per position and the relative property holds for any pair of query and key vectors, shared or not. The rotation for the shared key head is computed once per position and reused by every query head in the group, which is a real saving at decode time.

The interaction with the cache allocator is less benign. The cache for a grouped model holds $g$ heads instead of $h$, so the block size per position shrinks; block-based allocators that were tuned for a fixed number of heads per block may now pack several positions into one block, changing the indirection granularity. The allocator’s block size is a physical layout decision and does not have to match the head count, but the two interact through alignment: a block whose size is a multiple of the per-position byte count lets the kernel read it as one contiguous run. The arithmetic of the cache itself is in KV cache size arithmetic, and the kernel that consumes it in IO-aware attention tiling.

The remaining question is what happens at very long context, where the cache dominates memory regardless of $g$. At that point sharing has already extracted its factor, and the next lever is eviction or sparsification, which changes the model’s function rather than its representation. That is the subject of sliding-window attention and effective context reach.