Grouped-Query Attention (GQA)

Share keys and values within groups

Grouped-query attention (GQA) uses H query heads and K key/value heads. Each group of H // K consecutive query heads shares one key head and one value head. Queries remain separate, so heads in a group can assign different weights to the shared keys.

This reduces the key/value projection size compared with using H independent key/value heads at the same head width. It can also reduce cache storage when a decoder stores the shared heads without duplication. Neither benefit guarantees a particular runtime speed or model quality.

This exercise initializes a custom Linen module and evaluates unmasked self-attention. It creates no cache and performs no training or generation.

Map each query head to its group

Let R = H // K. Query head h uses key/value head h // R. For H = 8 and K = 2, query heads 0 through 3 use KV head 0, and query heads 4 through 7 use KV head 1.

The reference makes this assignment explicit with jnp.repeat along the head axis. This runnable example shows the difference from tiling:

import jax.numpy as jnp

kv = jnp.array([[[10., 11.], [20., 21.]]])  # (1 position, 2 KV heads, 2 features)
repeated = jnp.repeat(kv, 4, axis=1)        # (1, 8, 2)
tiled = jnp.tile(kv, (1, 4, 1))

assert jnp.array_equal(
    repeated[0, :, 0], jnp.array([10., 10., 10., 10., 20., 20., 20., 20.])
)
assert jnp.array_equal(
    tiled[0, :, 0], jnp.array([10., 20., 10., 20., 10., 20., 10., 20.])
)
assert not jnp.array_equal(repeated, tiled)

Tiling in this form changes the assignment to query heads. Preserve the consecutive-group convention when matching the reference. An alternative implementation can index the appropriate KV head directly instead of building repeated arrays, provided it uses the same assignment.

Keep the dimensions distinct

For input x of shape (T, D_in), define D = qkv_features and d_head = D // H.

H, K and D must be positive, H % K == 0, and D % H == 0. D_in may differ from D. K = 1 gives multi-query attention; K = H gives separate keys and values for every query head. Keep both boundary cases valid rather than assuming a particular test configuration.

Implement the compact module

Define a GQA module with num_heads, num_kv_heads and qkv_features fields. In its compact forward method, create these layers in order, keeping their default biases and initializers:

1. nn.DenseGeneral(features=(H, d_head), axis=-1, name="q").
2. nn.DenseGeneral(features=(K, d_head), axis=-1, name="k").
3. nn.DenseGeneral(features=(K, d_head), axis=-1, name="v").
4. After attention, nn.DenseGeneral(features=D_in, axis=(-2, -1), name="out").

Apply the first three projections separately to x. Repeat keys and values H // K times along axis -2. Move the head axis to the front for arrays of shape (H, T, d_head) and calculate:

scores[h]  = queries[h] @ keys[h].T / sqrt(d_head)
weights[h] = softmax(scores[h], axis=-1)
result[h]  = weights[h] @ values[h]

Move the result back to (T, H, d_head) before applying the out layer. This final projection combines heads and returns to D_in. Preserve the layer names and initialization setup to reproduce the seeded reference. Add no causal mask, dropout, normalization, position encoding or residual connection.

What the storage comparison measures

For one sequence and one layer, a compact cache of T positions holds 2 * T * K * d_head key/value elements. With H independent KV heads it holds 2 * T * H * d_head. At equal head width and dtype, GQA therefore uses a fraction K / H of those cached elements. Multiply by bytes per element, batch size and layer count for a total byte count.

This reference repeats keys and values and builds dense scores of shape (H, T, T). It does not implement compact caching or reduce the number of query/key scores. The smaller learned K/V projections and the potential cache savings are separate from the cost of this full-sequence forward pass.