Cruxly

Cruxly › AI engineering › LLM Foundations

Transformers and Attention, From the Whiteboard

one attention row by hand, then the parameter and KV-cache arithmetic

Open the interactive lesson

Prompt: the interviewer hands you a marker and says: "Walk me through what happens to one token in one transformer block, then count the parameters of a 7B model from its config, then tell me why the KV cache exists and who pays when we shrink it."

40 minutes first. Do one attention row on paper with three tokens, then count a Llama-2-7B-shaped config. The trap is reciting "attention is all you need" vocabulary without being able to produce a single number; every later serving and training lesson in this track assumes you can.

Attention is a soft lookup, and you can do one row by hand

Attention is a dictionary lookup where the match is fuzzy and the answer is a weighted blend. Each token's vector is projected three ways: a query (what I am looking for), a key (what I advertise) and a value (what I hand over if you attend to me). Token i scores every visible token j with q_i · k_j, turns the scores into weights with a softmax, and returns the weighted sum of the value rows.

Do one row with three tokens. Query q = [1, 0], keys k1 = [1, 0], k2 = [0, 1], k3 = [1, 1], values v1 = [1, 0], v2 = [0, 2], v3 = [2, 2]:

scores     = [q.k1, q.k2, q.k3]          = [1, 0, 1]
scaled     = scores / sqrt(d_k = 2)      = [0.707, 0, 0.707]
softmax    = e^x / sum e^x               = [0.401, 0.198, 0.401]
output     = 0.401*v1 + 0.198*v2 + 0.401*v3
           = [0.401 + 0.802, 0.396 + 0.802] = [1.203, 1.198]

The weights sum to 1, so the output is a point inside the cloud of value rows. Tokens 1 and 3 match the query equally and share most of the weight; token 2 still gets 20 percent, because softmax never gives exactly zero. That is the "soft" in soft lookup, and it is why attention can blend several sources instead of choosing one.

Why divide by sqrt(d_k)? If query and key components are roughly independent with unit variance, the dot product of two d_k-dimensional vectors has variance d_k, so its standard deviation is sqrt(d_k). A simulation of random Gaussian vectors gives a standard deviation of 2.0 at d_k = 4, 8.0 at d_k = 64 and 23 at d_k = 512. With raw scores [8, 0, -8] the softmax is [0.9997, 0.0003, 0.0000]: nearly one-hot, and the softmax gradient in that regime is close to zero, so training stalls. Dividing by sqrt(64) = 8 turns the same scores into [1, 0, -1] and weights [0.665, 0.245, 0.090], a distribution that still has gradient. The scaling is a variance fix, not a tuning knob, and "so the softmax does not saturate" is the one-sentence answer.

:::key In matrix form the whole layer is softmax(Q K^T / sqrt(d_k)) V. Q K^T is an n x n table of who-looks-at-whom; that table is where the O(n squared) cost lives. :::

One block, and where the parameters live

A block is attention plus an MLP, each wrapped in a residual connection, and about two thirds of the weights are in the MLP. For one token the sequence is:

@@ 1 | Normalise Apply RMSNorm (or LayerNorm) to the token's vector so its scale is predictable.

@@ 2 | Attention Project to Q, K, V, mix values from visible tokens as above (several heads in parallel), project back with W_o, and add the result to the original vector. That addition is the residual stream.

@@ 3 | Normalise again, then MLP Normalise, expand to a wider hidden size, apply a nonlinearity (SwiGLU in Llama-style models), project back, and add to the stream again.

The stream is a shared whiteboard: each block reads it, writes a small update, and passes it on. Modern models use pre-norm, where normalisation sits before each sub-layer, inside the residual branch. That leaves the identity path from input to output untouched, so gradients flow through dozens of layers without being rescaled at each step; the original post-norm layout needed careful warm-up to train deep stacks. Keep it to that one paragraph in an interview.

Counting parameters is a four-matrix, three-matrix exercise. Take a Llama-2-7B-shaped config: d_model = 4096, 32 blocks, MLP hidden size 11,008, vocabulary 32,000.

attention per block = 4 * d^2            = 4 * 4096^2     = 67.1M   (W_q, W_k, W_v, W_o)
MLP per block       = 3 * d * 11,008     = 135.3M          (gate, up, down in SwiGLU)
block total         = 202.4M             (rule of thumb: 12 * d^2 = 201.3M)
32 blocks           = 6.476B
embeddings + output = 2 * 32,000 * 4096  = 262M            (untied input and output matrices)
total               = 6.74B  -> "7B"

The rule of thumb 12 d^2 per block works because attention is 4 d^2 and a classic MLP with a 4x expansion is 8 d^2. SwiGLU uses three matrices at a smaller hidden size to land at nearly the same count. The MLP share is 135.3 / 202.4 = 67 percent. Interviewers like to ask where the parameters are because it tells you what fine-tuning methods touch: LoRA can target only the attention projections, a third of the block, but current practice applies it to every linear layer, MLP included, because attention-only adapters fall short of full fine-tuning.

Two consequences to state. Decode speed is governed by bytes read per token, which is why the parameter count converts directly into the memory-bandwidth arithmetic in GPU Memory Math: Weights, KV Cache, and What Actually Fits. And embeddings are a rounding error at 7B (4 percent) but not at small scale: a 1B model with a 128K vocabulary and d = 2048 spends 0.5B on embeddings and output together if they are untied, half its parameters (Llama 3.2 1B ties them, so it spends 0.26B, about a fifth), which is why vocabulary size is a design decision (see the next lesson).

Which mask makes it a chat model, an encoder or a translator

The only structural difference between the three families is which tokens may attend to which. A causal mask sets scores for future positions to minus infinity before the softmax, so token i sees only tokens 1..i.

So "why can a decoder-only model do classification while an encoder cannot chat?" has a mechanical answer. Classification is a text-in, label-out task and any generator can emit the label as its next token; an encoder has no output head that produces tokens in sequence, so chatting would need a decoder bolted on. The reverse also holds: for pure classification or embedding, a small encoder is often far cheaper than a chat model (measure it on your workload), which is the practical reason retrieval stacks use them (see Embeddings and Similarity later in this topic).

Training exploits the mask; inference pays for it. With a causal mask, one forward pass over a 1,000-token document yields 1,000 next-token predictions at once, because position i is not allowed to see i+1. That parallelism is why pretraining is efficient. At inference you do not know the next token yet, so you generate one at a time and each new token must read all earlier ones. That asymmetry is the whole reason the KV cache exists, and why it does not exist in training: training sees the entire sequence at once and computes keys and values for every position in the same pass. Inference reuses them across steps.

Prefill and decode are two different workloads from this one fact. Prefill processes the prompt in parallel (compute-bound); decode produces one token per step reading the whole cache (memory-bound). The serving lessons in this track (TTFT, TPOT, and the Latency Budget and Continuous Batching and the Scheduler) are built on that split.

The KV cache: O(n) memory, O(n squared) compute, and who pays for GQA

Attention costs O(n squared) compute but the KV cache costs only O(n) memory, and both numbers matter for different reasons. The score table has n x n entries per head: at n = 8,000 that is 64 million scores (128 MB per head in fp16 if you materialised it), and going from 8K to 32K context multiplies attention compute by 16 while multiplying the cache by only 4. Fused kernels such as FlashAttention avoid ever writing the table to memory, which changes the memory story but not the FLOP count.

The cache stores one key and one value vector per token per layer per KV head. For the 7B-shaped config (32 layers, head size 128, fp16):

multi-head (32 KV heads):  2 * 32 * 32 * 128 * 2 bytes = 524,288 B = 512 KB per token
                           at 4,096 tokens = 2.0 GB per request
grouped-query (8 KV heads): 2 * 32 * 8 * 128 * 2       = 131,072 B = 128 KB per token
                           at 4,096 tokens = 0.5 GB per request
multi-query (1 KV head):    2 * 32 * 1 * 128 * 2       =  16,384 B =  16 KB per token

What changes between MHA, MQA and GQA is how many key/value heads the query heads share. Multi-head attention gives every query head its own K and V. Multi-query attention (one shared K/V head) cuts the cache 32x but costs quality. Grouped-query attention, from the GQA paper (Ainslie et al., 2023), uses an intermediate number of groups and reports quality close to multi-head with speed near multi-query; the paper also shows you can convert an existing multi-head checkpoint by uptraining on a small fraction of the original compute. Four-times cache reduction, as in the example above, is a 4x concurrency gain on the same GPU, the arithmetic worked in GPU Memory Math: Weights, KV Cache, and What Actually Fits.

Who pays? The model owner pays a small quality risk at training time (and the K/V projection weights shrink, from 4 d^2 to about 2.5 d^2 of attention weights at 8 KV heads here). The operator gains cache memory and decode bandwidth. The user gets longer contexts and cheaper tokens. MLA, in DeepSeek-V2 (arXiv 2405.04434), goes further by caching a compressed latent per token instead of per-head keys and values; the paper reports a 93.3 percent KV-cache reduction against its own earlier dense model, which is the authors' claim for their architecture rather than a general rule.

:::key Verified October 2026 Llama-2-7B itself uses multi-head attention (32 KV heads); Llama-2-70B and most newer open models use GQA. The cache numbers above use a 7B-shaped config for illustration. Check any specific model's num_key_value_heads in its config before you size a fleet. :::

Positions: why context extension is a position-encoding problem

Attention by itself is order-blind, so position has to be injected, and the way it is injected decides how far the model can read. Shuffle the keys and values and the weighted sum is unchanged. The original transformer added fixed sinusoidal vectors to the input embeddings; learned absolute embeddings were the other early option. Both tie a token to an absolute slot, and neither extrapolates cleanly past the trained length.

RoPE (rotary position embedding, RoFormer 2021) rotates the query and key vectors by an angle proportional to position. In 2D: rotate q by m * theta and k by n * theta; their dot product then depends only on m - n. A numeric check with q = [1, 0.5], k = [0.4, 0.9], theta = 0.3 gives a score of 1.0767 for positions (5, 2) and the identical 1.0767 for positions (105, 102). The model sees relative distance, which is what language mostly needs, and nothing is added to the embeddings. Real models rotate many 2D pairs at a spread of frequencies, fast ones for local order and slow ones for long range. Most open models and many frontier models use RoPE or a close variant.

Context extension is then a question of what happens at positions the model never saw. Run a model trained to 4,096 tokens at position 16,000 and the high-frequency rotations land at angles never seen in training; attention scores go out of distribution and quality falls off a cliff, not gradually. The fixes rescale positions rather than retrain from scratch. Position interpolation (Chen et al., 2023) squeezes position indices back into the trained range: to go from 4,096 to 16,384 tokens, multiply every position by 4096/16384 = 0.25. The paper reports extending LLaMA models to 32,768 tokens with about 1,000 fine-tuning steps or fewer, and argues that interpolating is far safer than extrapolating. YaRN and similar methods refine this by treating frequencies differently.

What this means in the room:

How to say it in the interview

"A block does two things to each token's vector: attention lets it read from earlier tokens by scoring its query against their keys and mixing their values, and the MLP transforms it on its own. Both write back into a residual stream. Roughly 12 times d-model squared parameters per block, two thirds of them in the MLP; attention compute grows with the square of the sequence, but the KV cache only grows linearly, and grouped-query attention shrinks that cache by sharing key and value heads."

Checklist

Practice this interactively
Share this lessonWhatsAppLinkedInX

Related lessons