← Back to Blog

Model notes · Chapter 2

Attention: an engineering view of information flow

Understand Q/K/V, self- versus cross-attention, and the memory trade-offs behind MHA, MQA, GQA and MLA.

Muxuan Liu · Prepared from my model-learning notes; examples and equations checked for this edition.

Start with information routing

I find attention easier to understand as a way to route information. At each position, the model forms a query: what information would help here? Keys describe what each available position offers, and values contain the information to combine. Attention compares queries with keys, then uses the resulting weights to mix values.

Capturing linguistic relationships is one application of that operation. It also works on images and other structured inputs. An attention weight alone is not proof of a grammatical relationship or a complete explanation of a model’s decision.

Self-attention and cross-attention do different jobs

Self-attention obtains Q, K and V from the same sequence. Its output has one vector per query position; it does not decide how many translated tokens to produce. Cross-attention takes queries from one sequence and keys/values from another. Its score matrix can therefore have different query and key lengths.

RNN encoder–decoder systems could already handle unequal input and output lengths. Early attention relieved the fixed-vector bottleneck by providing a changing view of encoder states. In an encoder–decoder Transformer, cross-attention connects the sequences; the decoder’s generation process and stopping rule determine output length. Decoder-only models do not necessarily have a separate cross-attention module.

The calculation, one shape at a time

For one head, Q has shape [Tq, Dh], K is [Tk, Dh], and V is [Tk, Dv]. QKᵀ is [Tq, Tk]. Softmax is applied over keys so that every query gets its own mixture of values. Dividing by √Dh helps control the scale of dot products. The scaled formulation belongs to the 2017 Transformer; Bahdanau’s earlier attention uses a different scoring function.

Multiple heads compute separate mixtures in parallel, then concatenate their outputs and apply an output projection. Heads may learn different patterns, but “one head handles verbs, another handles time” should be treated as an illustration, not an assigned or guaranteed role.

Parallel training is not parallel token generation

With known training targets and a causal mask, a Transformer can compute many positions together. In ordinary autoregressive generation, the next token still depends on earlier generated tokens, so decoding proceeds step by step. These are different stages of use.

A KV cache stores the keys and values of the existing prefix so they can be reused. This relies on causal attention and an unchanged prefix/configuration: adding a new token does not change earlier causal states. It reduces repeated computation, but does not remove the need to attend to cached positions.

MHA, MQA, GQA and MLA: what gets stored?

MHA keeps separate K/V heads. MQA shares one K/V head across all query heads. GQA uses g K/V heads for h query heads: g=h recovers MHA, and g=1 recovers MQA. At fixed batch size, length, head width and precision, standard KV storage is proportional to 2·B·T·g·Dh per layer. Actual memory and latency also depend on kernels and layout.

MLA, described in DeepSeek-V2 in 2024, uses a compressed latent representation for keys and values, with additional positional handling. It is not simply GQA with a smaller g. The table below is a structural comparison, not a promise of equal model quality or a universal speed ranking.

Implementation: the transpose you must not skip

The teaching code accepts query and key sequences of different lengths. After weighting values, the tensor is [B, H, Tq, Dh]. It must be transposed back to [B, Tq, H, Dh] before flattening the heads. A direct view would mix the head and sequence layout incorrectly.

The example uses boolean masks with True meaning “allowed.” It rejects rows with no permitted key, which would otherwise make softmax over all −∞ undefined. Mask conventions vary between APIs. This is a transparent reference implementation, not a production KV-cache or fused-kernel implementation.

import math
import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    """Teaching example. Mask: True = allowed, [Tq,Tk] or [B,1,Tq,Tk]."""
    def __init__(self, d_model, num_heads, dropout=0.0):
        super().__init__()
        if num_heads < 1 or d_model % num_heads:
            raise ValueError('d_model must be divisible by num_heads')
        self.h = num_heads
        self.d = d_model // num_heads
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, allowed=None):
        if k.shape[1] != v.shape[1]:
            raise ValueError('Key and value lengths must match')
        b, tq, width = q.shape
        def split(x, projection):
            return projection(x).reshape(b, -1, self.h, self.d).transpose(1, 2)
        Q = split(q, self.q_proj)  # [B,H,Tq,Dh]
        K = split(k, self.k_proj)  # [B,H,Tk,Dh]
        V = split(v, self.v_proj)  # [B,H,Tk,Dh]
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d)
        if allowed is not None:
            if allowed.dtype != torch.bool:
                raise ValueError('Use a boolean mask: True = allowed')
            if not allowed.any(dim=-1).all():
                raise ValueError('Each query must have at least one allowed key')
            scores = scores.masked_fill(~allowed, float('-inf'))
        weights = self.dropout(scores.softmax(dim=-1))
        out = weights @ V  # [B,H,Tq,Dh]
        out = out.transpose(1, 2).contiguous().reshape(b, tq, width)
        return self.out_proj(out)