NNX Multi-Query Attention (MQA)

Share keys and values across heads

In ordinary multi-head attention, each query head has its own keys and values. Multi-query attention (MQA) keeps separate query heads but gives all of them the same keys and values. It is grouped-query attention with one key/value head.

Sharing does not force the heads to produce identical outputs. Different queries can assign different weights to the same keys, then take different weighted averages of the shared values.

Follow the shapes

Let T be the sequence length, D = d_model, H = num_heads and Dh = D // H. Require positive dimensions and D % H == 0.

K and V have a feature axis of width Dh, but no separate head axis. Broadcast each to (H, T, Dh) to use batched matrix multiplication. Adding a leading singleton axis makes the intent visible; broadcasting directly from (T, Dh) also works. This describes array values and shapes, not a guarantee about views or physical allocations in compiled JAX code.

Compute scores as Q @ K.transpose(0, 2, 1) / sqrt(Dh). Apply softmax along the last axis, which selects among key positions. Multiplying by V gives (H, T, Dh). Transpose to (T, H, Dh) before reshaping to (T, D) so each token's heads stay together, then apply the output projection.

A small example

Here are two query heads attending to shared keys [1, 0] and values [3, 6]. Each head has width one, so the scaling factor is one.

import jax
import jax.numpy as jnp

# Shape: (heads, query positions, head width).
q = jnp.array([[[jnp.log(2.0)], [0.0]],
               [[jnp.log(3.0)], [0.0]]])
k = jnp.array([[1.0], [0.0]])
v = jnp.array([[3.0], [6.0]])
keys = jnp.broadcast_to(k, (2, 2, 1))
values = jnp.broadcast_to(v, (2, 2, 1))
scores = q @ keys.transpose(0, 2, 1)
weights = jax.nn.softmax(scores, axis=-1)
per_head = weights @ values
joined = per_head.transpose(1, 0, 2).reshape(2, 2)
assert jnp.allclose(joined, jnp.array([[4.0, 3.75], [4.5, 4.5]]))

At the first query position, head one uses weights [2/3, 1/3] and head two uses [3/4, 1/4]. At the second position both queries are zero, so both heads average the values equally. This example starts with projected arrays; the exercise also learns the four linear projections.

What memory does MQA save?

During autoregressive decoding, a KV cache stores earlier keys and values so they can be reused. With the same head width, token count and precision, storing one KV head instead of H needs 1/H as many cache elements.

For 64 query heads, head width 128, 4096 tokens and two bytes per element, an ordinary MHA cache takes 2 * 4096 * 64 * 128 * 2 bytes, or 128 MiB per layer and sequence. An MQA cache takes 2 MiB. Across 32 layers that is 4 GiB versus 64 MiB; for 16 sequences, 64 GiB versus 1 GiB.

Those numbers describe the KV arrays only. This exercise implements no cache and still builds dense (H, T, T) scores. A cache reduction is not the same reduction in total memory or execution time. Sharing K/V also changes the model's parameterization; its effect on quality depends on the model and training setup.

The diagram uses eight query heads. Change the KV head count to compare sharing patterns; the connecting lines show which heads share K/V, not their attention weights.