Inside attention, one tensor at a time

What does attention actually calculate?

Attention lets a token position collect information from other positions. It calculates a weight for each available position, then uses those weights to mix that position’s value vector into the result. A decoder applies this operation repeatedly as it builds the representation used to predict the next token. The useful question is therefore: which vectors are mixed, how are the weights calculated, and where does the result go?

We will follow a small, complete decoder with two layers, a vocabulary of 32 token IDs, hidden width \(D=32\), four query heads, and two key/value heads. Each head has eight channels. These are deliberately small dimensions for inspecting the same kinds of operations used in a Llama-style decoder: RMSNorm, rotary position embeddings, grouped-query attention, SwiGLU, residual additions, and a vocabulary projection. [1]

The accompanying NumPy model uses deterministic random weights. It executes these mechanisms, but it has not learned language. A bright heatmap cell is a result of those demo weights; it does not establish that a trained language model associates two particular words. Token labels, when shown, are reading aids for a small vocabulary rather than the output of a pretrained tokenizer.

From alignment to self-attention

Bahdanau and colleagues’ 2014 preprint, presented at ICLR 2015, introduced an influential translation architecture in which a recurrent decoder forms a different weighted combination of source states at each output step. Luong and colleagues expanded this design space in 2015, studying global and local attention alongside scoring functions such as dot products and learned bilinear comparisons.

The 2017 Transformer made attention the main mechanism for communication between positions, removing recurrence from its proposed encoder–decoder architecture. Its encoder uses self-attention, its decoder uses causal self-attention, and its decoder also reads encoder states through cross-attention. In self-attention, queries and K/V arise from the same sequence; in cross-attention, they arise from different sequences.

Our example follows the causal decoder branch with later additions: RMSNorm, SwiGLU, RoPE, and GQA. These components belong to the modern architecture being inspected, rather than to the original 2017 block.

Start with a token ID, then distinguish its states

An embedding table \(E\in\mathbb{R}^{N_{\mathrm{vocab}}\times D}\) maps an integer token ID to a row of numbers. For six input IDs, the lookup gives \(X^{(0)}\in\mathbb{R}^{1\times6\times32}\). Repeated occurrences of the same ID start with the same embedding. The table itself does not inspect the surrounding sentence.

After a decoder layer, the state at a position includes information collected from its permitted context. The same token can consequently have different states in different prefixes. At layer \(\ell\), attention projects that layer’s normalized input state—not a fresh lookup of the original token embedding. This distinction explains why higher-layer keys and values contain contextual information.

In the causal decoder studied here, position \(t\) can use positions \(0,\ldots,t\). Its final state predicts the token at \(t+1\). Processing all six known input positions in parallel does not reveal future tokens to earlier positions: the attention mask controls which positions can communicate. This is the causal decoder masking used in the original Transformer. [2]

Queries, keys, and values: learned roles and exact shapes

Normalize the current state and apply three learned linear maps:

A query participates in the comparison made by the receiving position. A key participates in the comparison for a candidate source position. A value carries the information that will be mixed. These are numerical roles in an operation; the vectors are learned features rather than literal questions, dictionary keys, or human-readable facts.

Grouped-query attention, or GQA, allows multiple query heads to share one KV head. Here, query heads 0 and 1 use KV head 0; query heads 2 and 3 use KV head 1. They still produce different attention distributions because their queries differ. Ordinary multi-head attention uses as many KV heads as query heads; multi-query attention uses one KV head. [3]

From a dot product to a mixture of values

For query head \(h\), let \(g(h)=\lfloor h/2\rfloor\) select its shared KV head. With position information already applied to the queries and keys, the calculation is:

The causal mask is zero when \(s\le t\), and negative infinity otherwise. A forbidden position therefore receives zero probability. Softmax runs over key positions, so every valid row sums to one. Division by \(\sqrt{d_h}\) controls the scale of dot products as head width changes. [2]

Consider the query at position 2. Only columns 0, 1, and 2 can contribute; columns 3–5 are unavailable. The output is a weighted sum of three eight-dimensional vectors. The first query position has only itself available, so its probability row must be \([1,0,0,0,0,0]\). These conclusions follow from the mask, regardless of the random weights.

An attention probability is specific to one layer, head, and receiving position. Reading it as the model’s overall explanation would ignore the value vectors, other heads, output projection, residual path, and later layers. The heatmap is most useful here as a way to inspect the computation.

How position enters: rotate queries and keys

The demo uses rotary position embeddings, or RoPE. For each pair of channels, position \(p\) determines a rotation angle. Rotating both queries and keys makes their dot product depend on their relative positions. Values remain unrotated in this implementation. [4]

Our explicit convention pairs the first half of a head with its second half: channels \((0,4),(1,5),(2,6),(3,7)\). Other implementations use adjacent pairs with a corresponding weight layout. A pretrained checkpoint’s weights and its RoPE convention must agree. The six prompt positions are \(0,1,2,3,4,5\); this position counter will matter when we introduce a KV cache.

Read the operation line by line in NumPy

The following independent example uses arbitrary random weights to expose each transformation. It is a focused calculation, not a pretrained model or a copy of the interactive demo’s weight initialization. Run the three snippets in order. The complete two-layer model is in tiny_llama.py.

import numpy as np

B, T, D, Hq, Hkv, dh = 1, 6, 32, 4, 2, 8
rng = np.random.default_rng(7)
ids = np.array([[1, 2, 3, 4, 5, 6]])
embedding = rng.normal(0, 0.1, (32, D))
x = embedding[ids]                         # [1, 6, 32]
z = x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)

wq = rng.normal(0, 0.1, (D, Hq * dh))
wk = rng.normal(0, 0.1, (D, Hkv * dh))
wv = rng.normal(0, 0.1, (D, Hkv * dh))
wo = rng.normal(0, 0.1, (Hq * dh, D))
q = (z @ wq).reshape(B, T, Hq, dh).transpose(0, 2, 1, 3)
k = (z @ wk).reshape(B, T, Hkv, dh).transpose(0, 2, 1, 3)
v = (z @ wv).reshape(B, T, Hkv, dh).transpose(0, 2, 1, 3)

The lookup selects six embedding rows. Normalization averages over the last axis, independently for each token; its learned scale is implicitly all ones in this fragment. Each matrix multiplication creates all heads at once. reshape separates the head index from its channels, and transpose moves the head axis before the position axis. These two operations have different jobs: reshaping directly to [B,H,T,dh] would scramble the intended grouping.

positions = np.arange(T)

def rope(a, positions):
    half = a.shape[-1] // 2
    frequencies = 10000.0 ** (-np.arange(half) / half)
    angles = positions[:, None] * frequencies[None, :]
    c = np.cos(angles)[None, None, :, :]
    s = np.sin(angles)[None, None, :, :]
    left, right = a[..., :half], a[..., half:]
    return np.concatenate(
        [left * c - right * s, left * s + right * c], axis=-1
    )

q, k = rope(q, positions), rope(k, positions)

The frequency vector has four entries. Broadcasting supplies one set of angles for each position to every batch and head. The last line rotates Q and K using the same positions, leaving V untouched.

groups = Hq // Hkv
qg = q.reshape(B, Hkv, groups, T, dh)
scores = np.einsum("bgrtd,bgsd->bgrts", qg, k) / np.sqrt(dh)
allowed = positions[None, :] <= positions[:, None]
scores = np.where(allowed[None, None, None, :, :], scores, -np.inf)
scores = scores - np.max(scores, axis=-1, keepdims=True)
weights = np.exp(scores)
weights = weights / np.sum(weights, axis=-1, keepdims=True)
heads = np.einsum("bgrts,bgsd->bgrtd", weights, v)
heads = heads.reshape(B, Hq, T, dh)
merged = heads.transpose(0, 2, 1, 3).reshape(B, T, Hq * dh)
attention_update = merged @ wo             # [1, 6, 32]

The first einsum sums over channels \(d\), producing one score for every query/key pair. Grouping queries lets us share K/V without storing repeated copies. The mask compares key positions along columns with query positions along rows. Subtracting the row maximum stabilizes softmax without changing its result. The second einsum sums over source positions \(s\). Finally, merging four eight-channel outputs gives 32 channels, and \(W_O\) learns how to mix them back into the residual stream.

Place attention inside the complete decoder block

Attention produces an update with the same shape as its input. A residual addition incorporates that update, and a second sublayer processes each position’s channels:

RMSNorm divides each token vector by its root mean square, with a small numerical stabilizer, then applies a learned channel scale. It does not subtract the channel mean. Each normalization has its own scale parameters. [5]

The feed-forward sublayer uses a gated product:

In the tiny model, both expansion matrices have shape \([32,96]\), and the contraction matrix has shape \([96,32]\). Multiplication \(\odot\) is elementwise: one branch gates the other. This sublayer mixes channels separately at each position; attention is where positions exchange information. SwiGLU is one of the gated feed-forward variants studied by Shazeer. [6]

After two blocks, a final RMSNorm and vocabulary projection produce logits with shape \([1,6,32]\). The demo ties that final projection to \(E^\top\). The final 32 counts vocabulary entries; its equality to the hidden width is a demo choice. The last position’s logits define the next-token distribution. Selecting a token is a separate step from computing attention. Meta’s reference implementation provides the full normalization–attention–residual–feed-forward ordering. [1]

Inspect one layer and head

Use the controls to inspect the tiny model’s actual attention probabilities. Check the forbidden upper triangle, compare two query heads sharing a KV head, and follow one row across its permitted sources. The values come from the demo computation; they illustrate mechanics rather than learned linguistic behavior.

What to check when your implementation disagrees

• Wrong axes: inspect the shapes before and after every reshape, transpose, and matrix multiplication.
• Wrong softmax axis: normalize over available keys. Normalizing over heads or queries changes the operation.
• Mask after softmax: this removes probability mass without renormalizing. Apply the mask to scores first.
• Missing self-attention: the current position is allowed. A strict “earlier only” mask leaves the first row empty.
• Incompatible RoPE pairing: use the convention expected by the projection weights and keep position IDs consistent.
• Untrained semantic claims: a working forward pass verifies arithmetic and shape handling; language competence requires trained weights.

Once these operations are clear, the next question becomes precise: which of them must be repeated when one token is appended? That is the starting point for the KV-cache article.

Read the implementation, line by line

The source reader below loads the actual tiny_llama.py used to generate these diagrams. Select an operation to inspect its numbered source lines and the role of each step. The complete implementation includes both decoder layers and a cache for each layer.

Reproduce the reference calculation locally:

python -m pip install numpy
python tiny_llama.py --verify --export-trace trace.json

The verification compares cached single-token and chunked outputs with full-sequence outputs, compares grouped attention with explicitly repeated K/V, checks future masking, and counts the cache payload. The tests use more than one batch item so that accidental broadcasting is less likely to hide an error.

These are measured CPU results for the deterministic NumPy model, not accuracy or speed benchmarks. See trace.json for the full verification record and the data behind the widgets.

Take the same questions to a pretrained model

For a realistic small model, the accompanying GPU probe uses the public TinyLlama-1.1B-Chat-v1.0 checkpoint. Its published configuration has 22 layers, hidden width 2,048, 32 query heads, four KV heads, head dimension 64, and a 5,632-channel feed-forward expansion.

This is a separate pretrained checkpoint. The educational class named TinyLlama above has independent random weights, a 32-label vocabulary, and much smaller dimensions. The downloaded model has a real tokenizer and trained parameters.

Install a CUDA-compatible PyTorch build for your machine, then download requirements-gpu.txt and requirements-numpy.txt into the same directory as the probe. The GPU requirements file includes the NumPy requirements file. Choose an available GPU explicitly before running:

nvidia-smi
python -m pip install -r requirements-gpu.txt
# Replace cuda:0 with the GPU you intend to use.
python probe_pretrained.py \
  --device cuda:0 \
  --dtype float32 \
  --output pretrained_trace.json

The probe runs inference with eager attention so that attention probabilities can be inspected. It records the actual projection and block shapes, cache tensor sizes, software versions, and GPU metadata. It then compares cached and full-forward logits. Eager attention exposes the computation for study; this experiment is not a throughput benchmark.

For this checkpoint, ordinary K/V payload per token at batch one is \(2\times22\times4\times64=11{,}264\) scalars. That is 45,056 bytes in float32, or 22,528 bytes in a two-byte format. These are configuration-based calculations, separate from measured total GPU memory.

Primary references

1. Meta: Llama 3 reference model implementation. Full decoder structure; the tiny implementation is independently written and uses reduced dimensions.
2. Vaswani et al. (2017), Attention Is All You Need. Scaled dot-product attention, multiple heads, and causal masking.
3. Ainslie et al. (2023), GQA. Sharing key/value heads across groups of query heads.
4. Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding. Position-dependent rotations of attention vectors.
5. Zhang and Sennrich (2019), Root Mean Square Layer Normalization. RMS-based normalization.
6. Shazeer (2020), GLU Variants Improve Transformer. Gated feed-forward networks, including SwiGLU.