NNX Grouped-Query Attention (GQA)

Let query heads share keys and values

Grouped-query attention (GQA) divides query heads into groups. Each group shares one key head and one value head. Queries within a group can still assign different attention weights because their query vectors differ.

Let H be the number of query heads and Hkv the number of key/value heads. Each KV head serves H // Hkv consecutive query heads. Require positive head counts and H % Hkv == 0.

Two useful limits are Hkv = H, which gives ordinary multi-head attention, and Hkv = 1, which gives multi-query attention. Changing this count changes the model's parameterization. Converting an existing trained model is a separate task and can require additional training; this exercise starts with newly initialized parameters.

Projection widths and head order

For input x shaped (T, D), set Dh = D // H, requiring D % H == 0. Q projects from D to D. K and V each project from D to Hkv * Dh. The output projection goes from D to D.

After projection, reshape Q to (T, H, Dh) and transpose it to (H, T, Dh). Reshape K and V to (T, Hkv, Dh) and transpose to (Hkv, T, Dh). These are learned projections of all input features; splitting into heads happens after projection.

To use ordinary batched matrix multiplication, repeat each KV head H // Hkv times along axis zero. For four query heads and two KV heads, the mapping is [K0, K0, K1, K1], with the same pattern for values. Query head h uses KV head h // (H // Hkv).

repeat duplicates each head consecutively. Tiling the head pattern instead gives [K0, K1, K0, K1], which pairs queries with different parameters. Other implementations can compute grouped attention without explicit repetition; this exercise uses repetition to make the mapping easy to inspect.

A small example

Four query heads attend to two positions. Each head has width one, so dividing scores by the square root of the head width changes nothing.

import jax
import jax.numpy as jnp

# Shapes: Q (4, 1, 1), K and V (2, 2, 1).
q = jnp.array([[[0.0]], [[jnp.log(2.0)]],
               [[0.0]], [[jnp.log(3.0)]]])
k = jnp.array([[[1.0], [0.0]], [[0.0], [1.0]]])
v = jnp.array([[[2.0], [6.0]], [[10.0], [14.0]]])
keys = jnp.repeat(k, 2, axis=0)
values = jnp.repeat(v, 2, axis=0)
scores = q @ keys.transpose(0, 2, 1)
weights = jax.nn.softmax(scores, axis=-1)
result = (weights @ values).reshape(4)
assert jnp.allclose(result, jnp.array([4.0, 10.0 / 3.0, 12.0, 13.0]))

The first two heads use values [2, 6]. Their weights are [1/2, 1/2] and [2/3, 1/3]. The last two heads use [10, 14], with weights [1/2, 1/2] and [1/4, 3/4]. This example starts with projected arrays and one query position; the exercise projects a full self-attention sequence.

Where cache savings come from

A decoder can cache earlier keys and values for reuse. Keeping the unexpanded KV arrays requires 2 * T * Hkv * Dh elements per layer and sequence. Ordinary MHA uses 2 * T * H * Dh. At equal head width, sequence length and precision, the ratio is Hkv / H.

For H=64, Hkv=8, Dh=128, T=4096 and two bytes per element, the KV cache is 16 MiB per layer and sequence, compared with 128 MiB for MHA. Multiply by the number of layers and sequences for the full KV cache.

This exercise has no cache. It expands K/V during the forward pass and computes dense (H, T, T) attention scores. Smaller stored KV arrays do not guarantee the same reduction in total memory or runtime. Model weights and temporary arrays also use memory, and trained-model quality depends on the architecture and training setup.

The diagram fixes eight query heads and lets you choose how many KV heads they share. The connections show grouping, not attention weights.