Skip to content

← Flax NNX step 37 of 100

Hard Primitives

NNX Multi-Query Attention (MQA)

Jump to the task

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.

Projection Input width Output width Shape before attention
Query D D (H, T, Dh) after reshaping
Key D Dh (T, Dh)
Value D Dh (T, Dh)
Output D D (T, D) after joining heads

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.

Loading visualization…

Problem

Write mha_mqa(seed, x, num_heads, d_model) using an MQAMHA(nnx.Module).

  1. Convert the seed and dimensions to Python integers. Create one nnx.Rngs(int(seed)) and pass it to all four projections.
  2. Construct nnx.Linear layers in this order: q_proj from D to D, k_proj from D to Dh, v_proj from D to Dh, and out_proj from D to D. Keep the default initializers and biases. Construction order matters because each layer consumes the random stream.
  3. Project x, split Q into heads and share K/V as described above. Compute scaled dot-product attention without a mask or dropout.
  4. Join the heads in token order and apply out_proj. Return the result flattened to shape (T * D,).

x is a finite float32 array of shape (T, D), with positive T and D. num_heads is positive and divides D. Scalar inputs arrive as floats representing integers. When using JIT, keep seed, head count and model width static. Every entrypoint call constructs a fresh model; a forward call on a retained model should leave its parameters unchanged.

All positions can attend to all positions here. MQA specifies sharing across heads; it does not itself impose causal masking.