Grouped-Query Attention

Back in Chapter 2 a block was said to hide one thing an assembly line never allows: at each station, a token reaches sideways and reads every earlier token in the sequence. That sideways reach is attention, and it is the reason word order and context mean anything at all. RMSNorm set each vector’s scale; rotary embeddings stamped each vector with its position. Now the vectors finally meet. This is the densest box in the block, so it gets a whole chapter — but only this one job: how one position reads the others, and the specific trick, grouped-query attention, that forge uses to make that reading cheap in memory.

Start with the shape of a read. Every position turns itself into three vectors: a query (what am I looking for?), a key (what do I offer as a label?), and a value (what do I actually carry?). A position compares its query against every key, turns those comparisons into weights, and takes the weighted blend of the matching values. A token about to predict what follows “the capital of France is” issues a query that its own earlier keys — “capital”, “France” — answer strongly, and pulls their values in. Attention is that lookup, run for every position at once.

5.1 Queries, keys, and the shared group

Here is forge’s whole attention forward, from model.py. Read it for the shapes; the actual dot-product is one call near the bottom, which the next section opens.

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor | None,
        cos: torch.Tensor,
        sin: torch.Tensor,
        start_pos: int,
        past: KVEntry | None,
    ) -> tuple[torch.Tensor, KVEntry]:
        b, t, _ = x.shape
        q = self.W_query(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.W_key(x).view(b, t, self.n_kv_groups, self.head_dim).transpose(1, 2)
        v = self.W_value(x).view(b, t, self.n_kv_groups, self.head_dim).transpose(1, 2)

        # qk-norm runs per head BEFORE the rotation — order is weight-critical
        if self.q_norm is not None and self.k_norm is not None:
            q, k = self.q_norm(q), self.k_norm(k)

        q = apply_rope(q, cos, sin, offset=start_pos)
        k = apply_rope(k, cos, sin, offset=start_pos)

        if past is not None:
            k = torch.cat([past[0], k], dim=2)
            v = torch.cat([past[1], v], dim=2)
        if self.kv_quant is not None:
            k, v = self.kv_quant(k), self.kv_quant(v)
        entry = (k, v)

        k = k.repeat_interleave(self.group_size, dim=1)
        v = v.repeat_interleave(self.group_size, dim=1)

        context = self.backend(q, k, v, scale=self.head_dim**-0.5, is_causal=mask is not None)
        context = context.transpose(1, 2).reshape(b, t, self.d_out)
        return self.out_proj(context), entry

Listing 5.1 · forge’s grouped-query attention. Three projections, the per-head norm and rotation from the last two chapters, the group expansion, then one call to the attention backend.

The first three lines project x into queries, keys, and values, each reshaped into heads — parallel attention computations that each look at a different head_dim-sized slice, so the model can attend several ways at once. Notice the asymmetry: W_query produces n_heads heads, but W_key and W_value produce only n_kv_groups — fewer. In the shipped 0.6B model that is 16 query heads against 8 key/value heads. That gap is the whole idea in the name.

Why keep fewer key/value heads? Because in a running engine the keys and values are what you store — every token ever seen keeps its key and value around so later tokens can attend to it (the KV cache, next chapter). Queries are used once and discarded; keys and values live for the whole conversation. Halving the number of key/value heads halves that stored memory. Grouped-query attention is the compromise: keep all the query heads, but let each key/value head serve a whole group of them. With 16 query heads and 8 key/value heads, group_size is 2 — every key/value head is shared by two queries.

The line that makes the sharing concrete is repeat_interleave(self.group_size, dim=1): it copies each key/value head group_size times so the head counts line up again before the dot-product. Head 0 and head 1 of the queries both read the same copied key/value head; heads 2 and 3 share the next; and so on. The storage stays small, but the arithmetic sees a full set of heads.

Two lines in between belong to earlier chapters: the q_norm/k_norm call is RMSNorm applied per head, and apply_rope is the rotation — done here, on queries and keys, right before they meet. And three things belong to later chapters: start_pos and past are how the cache feeds in tokens already seen, and kv_quant is a memory optimization that stays off until Part II. With no cache, start_pos is 0, past is None, and those two if blocks do nothing — so for this chapter, ignore them.

5.2 The read itself

The call self.backend(q, k, v, ...) runs one function, eager_attention. It is the entire dot-product read, in four lines:

def eager_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    scale: float,
    is_causal: bool,
) -> torch.Tensor:
    scores = q @ k.transpose(-2, -1) * scale
    if is_causal:
        scores = scores + _causal_bias(q.shape[-2], k.shape[-2], q.device, scores.dtype)
    return torch.softmax(scores, dim=-1) @ v

Listing 5.2 · the scaled dot-product. Scores, a causal mask, a softmax, and a weighted sum of the values.

q @ k.transpose(-2, -1) is every query dotted with every key: for a sequence of t tokens it produces a t × t grid where entry (i, j) is how strongly token i’s query matches token j’s key. The * scale divides those scores by the square root of head_dim; without it the dot-products grow with the head size, the softmax that follows sees enormous numbers, and it collapses onto a single token. The scale keeps the scores in a range the softmax can spread across.

Then the mask. _causal_bias builds a t × t grid that is 0 on and below the diagonal and −∞ above it, and adds it to the scores. A token may read itself and everything earlier, never anything later — position 3 sees positions 0 through 3 and nothing beyond. This is what makes the model a predictor: at training time it must guess token 4 from tokens 0–3 alone, so at inference it can generate left to right. Add −∞ to the future scores and the softmax on the next line turns them into exactly zero weight.

torch.softmax(scores, dim=-1) turns each row of scores into weights that sum to one — a distribution over “how much of each earlier token should I pull in.” The final @ v uses those weights to blend the value vectors. Each token walks out of attention holding a weighted mixture of every earlier token’s value, weighted by how well their keys answered its query.

5.3 Correct, to a rounding

Attention gets the same treatment as every mechanism before it: an independent second implementation of the same read, the same weights loaded into both, one seeded batch fed to each, and a required match. The companion driver builds forge’s attention and the reference’s, runs one (2, 9, 32) batch through each with a causal mask, and prints the largest disagreement:

Listing 5.3 · forge’s attention against an independent implementation.

Not exactly zero this time — a rounding-sized 1.19e-07, and for the same reason Chapter 2’s full forward pass landed at 5.96e-07. forge scales the scores after the matmul, q @ k.transpose(-2, -1) * scale; the reference scales the queries before it. The two are equal in exact arithmetic, but floating-point multiplication and addition do not perfectly commute, so the last digits differ. The gap sits more than eighty times under the 1e-5 gate, so the two implementations agree — and any real mistake would move the gap far past a rounding error, as the checkpoint now shows.

Attention is the third of a block’s four mechanisms, and the one place tokens read one another. The fourth runs next: a transformation applied to each position on its own, with no sideways reach, before the block’s final residual add. That is the next box — the feed-forward network.