Chapter 3 — Interview Bank

Every answer here follows the same four beats. Say the first two within 20 seconds; the last two are what turn a correct answer into a strong one.

Questions are stacked by difficulty. Each level assumes everything above it.

How questions are written here. Real interviewers rarely ask "define attention". They describe a situation (a model config, a dashboard, a profiler trace, a teammate's claim) and watch whether you spot the concept underneath. So every question below opens with a scenario, then names the concept being tested, then gives the four beats. Practice by reading only the scenario and answering out loud before you look.

Where the questions come from. The themes are the ones widely reported for three kinds of team: research labs (derive \sqrt{d_k}, the softmax Jacobian, the online-softmax recurrence), LLM infrastructure and serving teams (KV-cache sizing, GQA, prefill vs decode), and applied ML teams (masking bugs, padding, attention as explanation). No question is attributed to a named company, and none is presented as a verbatim reported question. Every number comes from the chapter's running example ("the river bank was steep") or from the published configurations of GPT-2 small, Llama 2 7B and 70B, Llama 3 8B and Mistral 7B. "(Section k)" points to the chapter tab, "(Card k)" to the Derivations tab.

Level 1 — Screen

Q1 · "Bank, or bank?"

You've joined the team behind a search assistant. The PM has read that the model's first layer is an embedding table and asks: "So the word 'bank' gets the same vector in 'the river bank was steep' and in 'the bank raised rates'? Then how does the model ever tell them apart, and what does attention add that the table can't?"

• Concept tested: context-free embeddings, and attention as content-based mixing.
• Intuition: the table is a dictionary entry printed before the sentence is read; attention lets the word read its neighbours and rewrite itself.
• Answer: yes, the lookup is the same. Attention then builds each token's new vector as a weighted average of the values of the tokens it can see, with weights set by how well its query matches their keys:

\mathbf o_i

token i's new, context-aware vector (bold = a vector)

\sum_j

add up over every token j that token i may see

a_{ij}

how much of its attention token i gives to token j; positive, and each row sums to 1

\mathbf v_j

token j's value: the message it hands over if attended to

\operatorname{softmax}_j

exponentiate each score, then divide by the total over all keys j

\mathbf q_i

token i's query: what it is looking for

\mathbf k_j

token j's key: what it advertises (^\top = transpose, so \mathbf q_i^\top \mathbf k_j is a dot product)

d_k

the length of each query and key vector (2 in the toy, 128 in Llama 2)

• In the chapter's toy, bank's value starts as pure "money", [0, 1]. After attending to the, river and itself, its output is [0.576, 0.284] in [geography, money] space: mostly geography.
• Why: any distance, all at once, by content. Every pair of tokens is one step apart, every position is computed in one matrix product, and the weights depend on what the tokens say, not where they sit (Section 2).
• Next they'll ask: "Then why did anyone use RNNs?" → They read context through one running summary vector, in order; attention removed that bottleneck (Q28).

Q2 · Query, key, value in plain words

A new teammate is reading a Llama architecture diagram and asks: "Why does every token need three different vectors? Why not just take the dot product of the embeddings with each other and use that as the attention score?"

• Concept tested: the three roles, and why scores must be asymmetric and learned.
• Intuition: a library: the query is your question, the key is the label on each book's spine, the value is what's inside the book.
• Answer: three learned linear maps of the same embedding give three roles. Hook: Ask, Advertise, Deliver.

\mathbf x

one token's embedding row (bank = [1, 1, 0, 0] = [nature, money, function word, describes])

W_Q

the learned query map: turns the embedding into a question (Ask)

W_K

the learned key map: turns it into a label others can match (Advertise)

W_V

the learned value map: turns it into the message it sends (Deliver)

\mathbf q

the query vector ([2, -1] for bank in the toy)

\mathbf k

the key vector ([1, 1] for bank)

\mathbf v

the value vector ([0, 1] for bank: pure "money")

• Why not \mathbf x^\top \mathbf x: raw dot products are symmetric (bank → river would equal river → bank) and every token would favour itself when norms are similar. With separate maps the toy gets score(bank → river) = 2 but score(river → bank) = 1: relevance is directional and learned (Section 3).
• Next they'll ask: "Could you merge W_Q and W_K into one matrix?" → For the scores, only the product W_Q W_K^\top matters (a d_{\text{model}} \times d_{\text{model}} matrix of rank at most d_k). Keeping the two thin factors is cheaper and enforces that low rank.

Q3 · Negative attention?

Your team's attention-visualization dashboard shows, for one token pair, "attention score −3.2" next to "attention weight 0.004". A teammate files a bug: "attention can't be negative." Is it a bug?

• Concept tested: score vs weight, and what softmax does.
• Intuition: the score is a raw vote; softmax turns the votes into shares of a pie.
• Answer: not a bug. The score is the scaled dot product, any real number. The weight is its softmax, always positive, and a row of weights sums to 1:

a_j

the weight on key j: its share of the row's attention

s_j

key j's score, \mathbf q^\top \mathbf k_j / \sqrt{d_k}; can be negative

e

Euler's number, about 2.718; e to any power is positive

\sum_{m}

add up over every key m in the row

s_m

the score of key m

• Why: exponentiating makes every term positive, and dividing by the total makes them sum to 1. A weight of 0.004 just means that score sits well below the row's top scores.
• Next they'll ask: "Does a big score always mean a big weight?" → Only relative to the rest of the row. Softmax is shift-invariant: add 10 to every score and the toy's bank weights stay 0.140 / 0.576 / 0.284 (Section 4, Card 2).

Q4 · The PR that removed \sqrt{d_k}

*You're reviewing a pull request that "simplifies" attention to scores = q @ k.T, dropping the scale. Unit tests at d_k = 8 on six-token inputs still pass. Production uses d_k = 128 and 4,096-token contexts. Do you approve?*

• Concept tested: why the scale exists (Card 3).
• Intuition: a sum of 128 coin-flip-sized terms is about 11 times bigger than one, and softmax magnifies gaps exponentially.
• Answer: no. With independent unit-variance components the dot product's variance is d_k:

\text{Var}

variance: how much something spreads around its average

\mathbf q^\top \mathbf k

the raw (unscaled) dot product of one query and one key

\sum_{c=1}^{d_k}

add up over the d_k components, c = 1 to d_k

q_c\,k_c

the product of the query's and key's c-th components, mean 0 and variance 1 each

• At d_k = 128 the logits have standard deviation 11.3. Over 4,096 keys the median top weight becomes 0.868 instead of 0.00526, and the entropy 0.49 nats instead of 7.83: the row is nearly one-hot. At d_k = 8 the damage is milder (median top weight 0.104 over 4,096 keys), and over six-token test inputs a peaked row still looks plausible, which is why the tests passed.
• Why: variances of independent terms add, so divide by the standard deviation, \sqrt{d_k} (Section 5).
• Next they'll ask: "Why not divide by d_k?" → At initialization that overshoots to near-uniform rows (Q26).

Q5 · "Attention shows the model's reasoning"

A fintech team has built a transformer that classifies loan-application documents. An analyst's model-explanation report shows attention heatmaps and says: "The model declined because it attended to the phrase 'late payment'." The reviewer asks you whether that sentence is a valid explanation.

• Concept tested: attention weights are not explanations.
• Intuition: knowing where someone looked doesn't tell you what changed their mind.
• Answer: not on its own. Jain and Wallace (2019) found attention weights often disagree with gradient-based importance, and that quite different attention patterns can give the same prediction. Wiegreffe and Pinter (2019) answered that whether attention "explains" depends on what you call an explanation, but neither paper licenses reading a heatmap as the reason for a decision.
• Why: three mechanical gaps. A weight multiplies a value, and a big weight on a near-zero value contributes nothing (in the toy, the takes 0.140 of bank's attention and adds nothing, because its value is [0, 0]). After layer 1 each position is already a blend of many tokens. And many heads and layers are summed through W_O and the residual stream.
• Next they'll ask: "What would you use instead?" → Attributions tested against the model's behaviour: integrated gradients or gradient × input, and ablations (remove the phrase and see whether the decision flips). See Q37.

Q6 · Training loss near zero, generations garbage

A team trains a small decoder-only model on support tickets. Training loss drops to 0.05 within an hour, and the validation loss (computed with the same forward pass) is just as low. Sampled text is nonsense. What's your first hypothesis, and how do you test it?

• Concept tested: the causal mask; future leakage.
• Intuition: the model was allowed to read the answer sheet while taking the test.
• Answer: the causal mask is missing or broken. Each position can see the token it is supposed to predict, so it learns to copy. At generation time the future isn't there. The mask adds -\infty to every forbidden score before the softmax:

S_{ij}

the score of query position i on key position j

M_{ij}

the mask entry: 0 if the pair is allowed, -\infty if forbidden

i

the position doing the reading (the query)

j

the position being read (the key)

• Test: change the future, and the past must not move. In the toy, changing steep moves the masked outputs of the first four tokens by exactly 0; without the mask, bank's output moves by 0.217 (Section 6, Card 6).
• Why: e^{-\infty} = 0, so a forbidden key gets weight exactly zero, not "small".
• Next they'll ask: "Why did validation look fine?" → It runs the same leaky forward pass. Only generation, or the change-the-future test, exposes it.

Q7 · Batch of one vs batch of eight

Your inference service returns slightly different logits for the same prompt depending on whether it runs alone or in a batch of 8 padded to the longest prompt. A customer notices that answers change with traffic. What's wrong?

• Concept tested: the key-padding mask. Mask types: Future, Fake, Far (causal, padding, sliding window).
• Intuition: the padded prompt is reading blank pages and averaging them into its answer.
• Answer: padding keys are not masked (or the mask has the wrong polarity, Q27). Real tokens spend attention on PAD tokens. In the toy, "the river" padded to length 5 gives river 49.8% of its attention on three PAD tokens; with a key-padding mask its weights are 0.330 / 0.670, identical to the unpadded sentence (Section 6).
• Why: a PAD key has a score like any other key unless you add -\infty to it. Rows at PAD positions are junk but harmless, as long as the loss ignores them.
• Next they'll ask: "Left or right padding for generation?" → Left-pad decoder-only batches so the last position of every row is a real token, and offset the position ids so padding doesn't shift positions.

Q8 · Do 32 heads cost 32× the parameters?

A finance partner is costing a Llama-2-7B-sized model and asks: "The spec says 32 attention heads. Doesn't that make attention 32 times as expensive as one head?"

• Concept tested: multi-head parameter count (Card 9).
• Intuition: the heads split one wide road into 32 lanes; they don't build 32 roads.
• Answer: no. Each head works in a slice of width d_k = d_{\text{model}}/h, so the parameter count is the same for any number of heads:

h

the number of heads (32 in Llama 2 7B)

d_{\text{model}}

the model width (4,096 in Llama 2 7B)

d_k

each head's width: d_{\text{model}}/h, 128 here

• For Llama 2 7B that is 67,108,864 parameters per layer, with 1 head or with 64.
• Why: the first term is W_Q, W_K, W_V (each d_{\text{model}} \times d_{\text{model}} in total); the second is the output map W_O.
• Next they'll ask: "So what do heads cost?" → Rank, because each head's score matrix has rank at most d_k; and KV-cache memory, unless heads share keys and values (Q29, Section 10).

Q9 · "Dog bites man"

An engineer testing an encoder forgot to add positional encodings. They shuffle the input tokens and find that the outputs are the same vectors, just shuffled. They file it as a bug in the attention kernel. Is it?

• Concept tested: permutation equivariance; why positions are needed.
• Intuition: without positions, attention sees a bag of words: a set, not a sequence.
• Answer: not a bug. Self-attention with no positions and no mask is permutation-equivariant: permute the inputs and the outputs permute the same way.

\text{Attn}

one self-attention layer with no positional information and no mask

\Pi

a permutation matrix: it reorders the rows (tokens)

X

the input, one row per token

• Reversing the toy sentence permutes the outputs to within 1.1 \times 10^{-16} (Section 9). So "dog bites man" and "man bites dog" get the same vectors, rearranged.
• Why: scores depend only on token contents, and a weighted sum doesn't care about the order of its terms (Card 7).
• Next they'll ask: "Is it permutation-invariant?" → No; that word is a common trap. Invariance needs pooling (a mean over tokens). And a causal mask breaks the symmetry, which leaks some position information (Drill 12). Real models add positions (LLM Ch 4).

Q10 · Why do long chats slow down?

Support reports that your Llama 2 7B assistant (fp16, one A100 80GB) streams quickly at first but slows down late in long conversations, and that the first reply after a user pastes a long document takes a while to appear. Explain both.

• Concept tested: the KV cache and memory-bound decoding.
• Intuition: every new word means re-reading the whole notebook of past keys and values, and the notebook keeps growing.
• Answer: each generated token must read all the weights and the whole KV cache from memory. The cache grows by 512 KiB per token, so per-token time grows linearly with the conversation:

t_{\text{token}}

the time to generate one token, as a lower bound

T

the number of tokens already in the conversation

• On an A100 the weights alone take 6.61 ms per token (at most 151 tokens/s at batch 1); a 4,096-token history adds 1.05 ms (130 tokens/s). The long paste is different: its prefill processes all the new tokens at once, and the score part of that work grows with the square of its length.
• Why: decoding is memory-bound, about 1 FLOP per byte read (Section 13).
• Next they'll ask: "What reduces it?" → Sharing K and V across heads (GQA, Q29), a quantized cache, reusing the cache for repeated prompt prefixes, or a sliding window (Q39).

Q11 · Self vs cross

A vision-language team says they "add cross-attention layers so the text can look at the image". Where do the queries, keys and values come from in those layers, and what's different from self-attention?

• Concept tested: self-attention vs cross-attention.
• Intuition: in self-attention a sentence reads itself; in cross-attention the sentence reads another document.
• Answer: the queries come from the text tokens (the readers); the keys and values come from the image encoder's output (the thing being read):

Q

the matrix of queries, one row per text token

K

the matrix of keys, one row per image token

V

the matrix of values, one row per image token

X_{\text{text}}

the text tokens' current vectors

Z_{\text{img}}

the image encoder's output vectors

• Why: the original encoder–decoder translators did the same with the source sentence (Bahdanau, Cho and Bengio, 2014; Vaswani et al., 2017). Self-attention is the special case where all three come from one sequence (Section 2).
• Next they'll ask: "Does cross-attention need a causal mask, and what does the cache look like?" → No mask over the image: all of it exists up front. Its keys and values are computed once per image and reused at every decoding step; the text's own self-attention keeps its causal mask.

Q12 · "Attention is O(n^2)": what exactly?

An executive slide says: "Attention is quadratic, so moving our Llama 2 7B product from 4K to 8K context will quadruple compute." Is that right?

• Concept tested: which parts of a layer are quadratic (Card 11).
• Intuition: only the pairwise comparisons grow with the square; everything that touches one token at a time grows linearly.
• Answer: no. Per layer, the projections and the feed-forward network are linear in T; only the score and weighted-sum products are quadratic:

T

the context length in tokens

d

the model width, d_{\text{model}} (4,096 for Llama 2 7B)

d_{ff}

the feed-forward width (11,008 for Llama 2 7B)

• At T = 4{,}096 that is 550 GFLOP for projections, 275 GFLOP quadratic and 1,108 GFLOP for the FFN: the quadratic part is 14.2% of the layer (7.7% with a causal kernel). Doubling to 8K multiplies the forward pass by 2.28, not 4 (2.15 with causal). The quadratic part takes over only past T^* = 2d + 1.5\,d_{ff} = 24{,}704 tokens; doubling from 32K to 64K costs 3.14×.
• Why: d² for the weights, T² for the pairs (Section 12).
• Next they'll ask: "And memory?" → The T \times T weights, if stored, grow quadratically too: 1 GiB per layer at 4K. Online softmax avoids storing them (Q33, Worked problem 2).

Level 2 — Onsite core

Q13 · Attention by hand

The interviewer writes the chapter's toy on the whiteboard: in "the river bank was steep", bank's query is [2, -1]; the keys of the, river and bank are [0, 0], [1, 0], [1, 1]; their values are [0, 0], [1, 0], [0, 1]; d_k = 2. "With a causal mask, compute bank's output. No calculator for the arithmetic you can do by hand."

• Concept tested: the full computation in order. Hook: Score, Scale, Mask, Softmax, Sum (four S's and an M).
• Intuition: bank asks "geography or money?", finds river's key lines up best, and borrows mostly river's value.
• Answer: scores 0, 2, 1 → divide by \sqrt 2: 0, 1.414, 0.707 → was and steep get -\infty → exponentiate: 1, 4.113, 2.028 (sum 7.141) → weights 0.140, 0.576, 0.284 → output:

\mathbf o_{\text{bank}}

bank's new vector in [geography, money] value space (it started as [0, 1], pure money)

• All rows at once, in matrix form (Card 1):

\text{Attn}

the attention layer's output, one row per token

Q

the queries stacked as rows, T \times d_k

K

the keys stacked as rows, T \times d_k (^\top = transpose, so QK^\top is every query–key dot product)

V

the values stacked as rows, T \times d_v

d_k

the query and key length (2 in the toy)

M

the mask: 0 where allowed, -\infty where forbidden

• Why: each output is a convex combination of the values: weights are non-negative and sum to 1 (Section 4).
• Next they'll ask: "What changes without the \sqrt 2?" → Weights 0.090 / 0.665 / 0.245. The trap: at d_k = 2 the scale softens the row slightly; saturation is a large-d_k problem (Q14).

Q14 · Derive the \sqrt{d_k}

A research-lab interviewer: "Every transformer divides the scores by \sqrt{d_k}. Derive why it's \sqrt{d_k}, not d_k and not \sqrt{d_{\text{model}}}, and state exactly what you assumed."

• Concept tested: the variance of a dot product (Card 3).
• Intuition: add up d_k independent wobbles and the total wobbles \sqrt{d_k} times as much.
• Answer: assume query and key components are independent with mean 0 and variance 1 (true at initialization with standard init and normalized inputs). Each product has variance 1, and variances of independent terms add:

\text{Var}

variance: the average squared distance from the mean

q_c

the query's c-th component

k_c

the key's c-th component

\mathbb E

the average over random draws (the expectation)

\mathbf q^\top \mathbf k

the raw dot product: a sum of d_k such products

• So divide by the standard deviation, \sqrt{d_k}. Hook: variance adds, so divide by the standard deviation. It's d_k, not d_{\text{model}}, because the dot product happens inside one head of width d_k. The simulation gives \text{Var}(\mathbf q^\top \mathbf k) = 128.4 at d_k = 128.
• Why it matters: logits with standard deviation s over T keys have entropy

H

the entropy of one attention row, in nats (\ln T = perfectly uniform)

T

the number of keys

s

the standard deviation of the logits in that row

• With s = 1 that's 7.82 (simulated 7.82). Unscaled at d_k = 128, s = 11.3 is far outside the formula's range: the row is nearly one-hot (median entropy 0.49).
• Next they'll ask: "And the gradients?" → A saturated row sends gradient only to its winner, and that gradient shrinks as the row saturates (22% of rows have a top weight above 0.99 at d_k = 128). Vaswani et al. (2017) themselves wrote that they "suspect" this; say "suspect" too (Section 5, Worked problem 1).

Q15 · Differentiate the softmax

"Here's one attention row: bank, with weights 0.140 / 0.576 / 0.284 on the, river and bank, and values [0, 0], [1, 0], [0, 1]. We want bank to read more like geography, so take a loss L that falls as bank's geography coordinate rises (L = minus that coordinate), with upstream gradient \mathbf g = [-1, 0]. Derive the gradient of L with respect to the three scores, compute it, say which way a gradient-descent step moves them, and tell me when it vanishes."

• Concept tested: the softmax Jacobian and the attention backward pass (Cards 4–5).
• Intuition: a descent step raises a key's score if its value is better, for the loss, than the average the row currently returns.
• Answer: the softmax derivative is "one-hot minus p":

\frac{\partial a_i}{\partial s_j}

how fast weight a_i changes as you nudge score s_j, holding the other scores fixed: a slope, not a fraction to divide

a_i

the weight on key i

a_j

the weight on key j

\delta_{ij}

1 if i = j, else 0

J

the softmax Jacobian: the table of all those slopes

\mathbf a

the row of weights as a vector

• Chain it with the output \mathbf o = \sum_j a_j \mathbf v_j:

\frac{\partial L}{\partial s_j}

how fast the loss L (minus bank's geography coordinate) changes as you nudge score s_j

\mathbf g

the upstream gradient \partial L / \partial \mathbf o, here [-1, 0]: the loss falls as geography rises

\mathbf v_j

key j's value

\mathbf o

the row's current output, [0.576, 0.284]

• Numbers: \mathbf g^\top \mathbf o = -0.576, so the score gradients are +0.081 (the), −0.244 (river), +0.164 (bank); they sum to 0. Descent moves against them, so river's score rises and the other two fall. Through \mathbf q_{\text{bank}} they give \partial L / \partial \mathbf q = [-0.057, 0.116]; the descent step goes along [0.057, -0.116], which turns bank's query toward river's key. A finite-difference check agrees to 2.4 \times 10^{-12}. The trap: the gradient points uphill. With \mathbf g = [1, 0] (a loss that grows with geography) every sign flips, and descent turns the query away from river.
• When it vanishes: when the row is one-hot (J \to 0; saturation), or when every value scores the same against \mathbf g. Hook: push values by weight; pull scores by surprise, a(v - o) (Section 8).
• Next they'll ask: "The biggest the Jacobian can be?" → Its largest eigenvalue is at most 1/2, reached by two keys at 50/50 (Drill 3).

Q16 · Code causal attention

"Implement batched, multi-head, causal self-attention in numpy. Inputs (B, T, d), four d \times d weight matrices, h heads. No Python loops over heads or positions. Then show me a one-line test that the mask works."

• Concept tested: shapes, the split/merge reshape, the mask, and a stable softmax.
• Intuition: split, attend, concat, project; the test is "change the future, the past must not move".
• Answer:

import numpy as np

def mha_causal(X, Wq, Wk, Wv, Wo, h):
    """X: (B, T, d); W*: (d, d). Multi-head causal attention with no loop over heads."""
    B, T, d = X.shape
    dk = d // h
    def split(M):                                        # (B, T, d) -> (B, h, T, dk)
        return M.reshape(B, T, h, dk).transpose(0, 2, 1, 3)
    Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
    S = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(dk)        # Score, Scale: (B, h, T, T)
    blocked = np.triu(np.ones((T, T), dtype=bool), k=1)  # key j > query i
    S = np.where(blocked, -np.inf, S)                    # Mask the scores, not the weights
    S = S - S.max(axis=-1, keepdims=True)                # stable Softmax
    P = np.exp(S)
    P = P / P.sum(axis=-1, keepdims=True)
    O = (P @ V).transpose(0, 2, 1, 3).reshape(B, T, d)   # Sum, then merge heads
    return O @ Wo

rng = np.random.default_rng(0)
B, T, d, h = 2, 6, 16, 4
X = rng.standard_normal((B, T, d))
Ws = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(4)]
out = mha_causal(X, *Ws, h)
X2 = X.copy()
X2[:, -1] += 5.0                                         # change the future: the last token
out2 = mha_causal(X2, *Ws, h)
print(out.shape, np.abs(out2[:, :-1] - out[:, :-1]).max())  # earlier rows must not move: 0.0

• Why each line matters: k=1 keeps the diagonal (Q17); the max-subtraction keeps exp finite (Q23); reshape(B, T, h, dk) then transpose keeps each head's features together (Q20).
• Next they'll ask: "Add a KV cache." → Keep K and V for past tokens; each step computes only the new token's q, k, v, and the test is that cached decoding matches this full causal pass (Section 14).

Q17 · The off-by-one mask

*A model's loss is NaN, and only the first position of each sequence is affected. The mask is built as blocked = np.triu(np.ones((T, T)), k=0). A teammate "fixes" it by filling with −1e9 instead of −inf; the NaN disappears, but the first-token loss stays poor. What happened, twice?*

• Concept tested: a fully masked row; silent leakage.
• Intuition: the first token is told it may look at nothing, not even itself.
• Answer: k=0 blocks the diagonal too, so row 0 has no allowed key: the softmax computes 0/0 = NaN. With a −1e9 fill, row 0's scores are all equal, so it becomes a uniform 0.2 over all five toy tokens: it now reads the future (Section 6). Fix: k=1.
• Why: every row of a causal mask must keep at least the diagonal. A finite fill hides a fully masked row instead of failing loudly.
• Next they'll ask: "How would a test have caught it?" → Change the future and check that no earlier output moves (Q6), and assert that no row of the mask is all-blocked.

Q18 · NaN after adding padding

Your loss turns NaN, but only on batches that contain an empty document, which the loader pads to full length. Masked positions are already excluded from the loss. Why is it still NaN, and what's the robust fix?

• Concept tested: fully masked rows, and where to mask.
• Intuition: a row with nothing to read divides zero by zero, and one NaN spreads.
• Answer: in the all-padding sequence every key is masked, so every row is fully masked: a hand-written softmax returns NaN there, and so does nn.MultiheadAttention with its default need_weights=True. Excluding those positions from the loss doesn't help: in the backward pass 0 \times \text{NaN} = \text{NaN}, so the NaN reaches the shared weights. (Checked on PyTorch 2.14.1: F.scaled_dot_product_attention on CPU instead returns zeros for a fully masked row, as does nn.MultiheadAttention with need_weights=False, which routes through it. That is silent, and it has varied by version and backend, so don't rely on either.)
• Fix: drop empty sequences in the loader; fill masks with the dtype's most negative finite value (torch.finfo(dtype).min) rather than -\infty when rows can be fully masked; mask keys, and ignore padded query rows only in the loss.
• Next they'll ask: "Why mask keys, not queries?" → A padded query's output is junk nobody reads; a padded key would leak into real tokens' outputs (Q7).

Q19 · Count the parameters

"GPT-2 small has d = 768 and 12 layers, with biases on every linear layer. How many parameters are in one layer's attention? How much of the 124M model is attention? Then do Llama 2 70B, which shares keys and values."

• Concept tested: the parameter count (Card 9).
• Intuition: four square maps (Q, K, V, O) plus a bias vector each.
• Answer:

d

the model width (768 for GPT-2 small)

• That's 2,359,296 weights plus 3,072 biases. Over 12 layers: 28,348,416, or 22.8% of 124,439,808.
• With grouped-query attention only H_{kv} key and value heads exist:

d

the model width (8,192 for Llama 2 70B)

H_{kv}

the number of key/value heads (8)

d_h

the width of each head (128)

• Llama 2 70B: 150,994,944 per layer, against 268,435,456 if it were full multi-head attention, 43.75% less.
• Next they'll ask: "And the number of heads?" → It cancels: 4d² for any h (Q8).

Q20 · Heads scrambled by a reshape

*A teammate reimplements GPT-2 and loads the pretrained weights. Perplexity is far worse than the reference, though no shape errors are raised. The head split is q.view(B, h, T, d_k). What's wrong?*

• Concept tested: memory layout of the head split.
• Intuition: you cut a stack of index cards the wrong way: each "head" gets whole cards instead of the same column from every card.
• Answer: a (B, T, d) tensor is laid out token by token. Viewing it directly as (B, h, T, d_k) reinterprets the memory, so "head 0" gets token 0's entire vector, then part of token 1. The correct split is view(B, T, h, d_k).transpose(1, 2), and the merge must undo it with transpose(1, 2).reshape(B, T, d):

import numpy as np

B, T, h, dk = 1, 3, 2, 2
X = np.arange(B * T * h * dk).reshape(B, T, h * dk)    # token t holds features 4t ... 4t + 3
right = X.reshape(B, T, h, dk).transpose(0, 2, 1, 3)   # (B, h, T, dk)
wrong = X.reshape(B, h, T, dk)                         # same shape, scrambled contents
print(right[0, 0].tolist())  # head 0 = the first dk features of every token
print(wrong[0, 0].tolist())  # "head 0" = all of token 0, then half of token 1

• Why: both tensors have the same shape, so nothing errors; the pretrained W_O expects each head's features in its own slice (Section 10).
• Next they'll ask: "How would you catch it?" → Compare one layer's output against the reference implementation on a fixed input before training anything.

Q21 · Entropy collapse on the dashboard

During pre-training, your dashboard shows the attention entropy of several heads in the last layer sliding from about 7 nats toward 0.3, while the maximum logit climbs. A few thousand steps later the loss spikes. What's happening, and what would you change?

• Concept tested: logit growth and entropy collapse (the C in S-C-A-L-E).
• Intuition: the query and key weights keep growing, so the softmax keeps getting colder until each head stares at one token.
• Answer: if W_Q and W_K grow by a factor c, every logit grows by c^2:

c

the factor by which W_Q and W_K have grown during training

\mathbf q

a query at initialization

\mathbf k

a key at initialization

• In the simulation (4,096 keys, d_k = 128, scaled) the median entropy goes 7.82 → 3.24 → 0.339 for c = 1, 2, 4. Zhai et al. (2023) link this entropy collapse to training instability.
• Fixes: QK-norm (normalize \mathbf q and \mathbf k, then use a learned scale; Henry et al., 2020; Dehghani et al., 2023), \sigmaReparam (Zhai et al., 2023), logit soft-capping (Gemma 2), and a lower learning rate or longer warmup. Soft-capping squashes each logit smoothly:

s

a raw attention logit

c_{\text{cap}}

the cap: logits can never exceed it in size

• With c_{\text{cap}} = 50, logits of 10, 50 and 200 become 9.87, 38.1 and 49.97.
• Next they'll ask: "What would you monitor?" → Per-head entropy, the max logit per layer, and the share of attention on token 0 (Q34).

Q22 · KV cache on a napkin

An infra interviewer: "Llama 2 7B: 32 layers, 32 heads, head dimension 128, fp16. How much memory does the KV cache take for one 4,096-token conversation? Show the formula."

• Concept tested: KV-cache size (Card 13).
• Intuition: K and V, every layer, every KV head, every dimension, every token.
• Answer:

L

the number of layers (32)

H_{kv}

the number of key/value heads (32: Llama 2 7B shares nothing)

d_h

the head dimension (128)

b

bytes per number (2 for fp16)

• That's 512 KiB per token, 2 GiB per 4,096-token sequence. The leading 2 is for K and V.
• Why: under a causal mask, past keys and values never change, so they are computed once and kept; the query is never cached (Section 13).
• Next they'll ask: "Llama 3 8B?" → 8 KV heads, so 128 KiB per token: 16 GiB at 131,072 tokens (Llama 3.1 8B's context).

Q23 · fp16 softmax overflow

A custom fused attention kernel works in fp16. It produces inf and NaN on long prompts but never on short ones. What's the likely bug, and why does length matter?

• Concept tested: numerically stable softmax (the E in S-C-A-L-E).
• Intuition: e^x outgrows fp16 fast, and a long row gives more chances for one big logit and a bigger sum.
• Answer: the kernel exponentiates raw logits. fp16's largest number is 65,504, so e^x overflows for x > 11.09. Subtract the row max first; softmax is unchanged:

e

Euler's number, about 2.718

s_j

key j's logit

\sum_{l}

add up over every key l in the row

s_l

key l's logit

s_{\max}

the largest logit in the row

• Now every exponent is at most 0. Also accumulate the sum in fp32: the denominator of a long row can pass 65,504 even when each term is fine. For a sense of scale: unscaled d_k = 128 logits exceed 11.09 about 16% of the time, and a naive fp16 softmax over 4,096 keys gave NaN in every simulated row; the stable version, none.
• Why length matters: more keys means a larger maximum logit and a larger sum (Section 11).
• Next they'll ask: "Where else does this trick appear?" → Online softmax keeps a running max for exactly this reason (Worked problem 2).

Q24 · Attention as kernel regression

A statistician on your team says: "Attention is just Nadaraya–Watson kernel regression from 1964 with a new name." Defend or refute, precisely.

• Concept tested: attention as a kernel smoother (Card 8).
• Intuition: both estimate "what's here" by averaging what's nearby, weighted by similarity.
• Answer: mostly defend, with one precise correction. Nadaraya–Watson averages known outputs with kernel weights:

\hat f(\mathbf x)

the estimate at a new point \mathbf x (hat = estimated)

\mathbf x

the point you want a prediction for: plays the query

\sum_j

add up over the stored examples j

K_h

the kernel: a similarity score that falls with distance, bandwidth h

\mathbf x_j

stored example j's input: plays the key

y_j

stored example j's output: plays the value

• A Gaussian kernel expands into a dot product plus two norm terms:

\mathbf q

the query vector

\mathbf k

a key vector

h

the kernel bandwidth: bigger means smoother, more uniform weights

\lVert \mathbf q - \mathbf k \rVert^2

the squared distance between query and key (\lVert\cdot\rVert = length of a vector)

\lVert \mathbf q \rVert^2

the query's squared length: the same for every key, so it cancels in the softmax

\lVert \mathbf k \rVert^2

the key's squared length: acts as a per-key bias

• So a Gaussian NW smoother is dot-product attention with temperature h^2 and a bias -\lVert\mathbf k\rVert^2/2h^2 on each key (the simulation matches to 5.6 \times 10^{-16}; plain \mathbf q^\top \mathbf k without the bias is off by up to 1.09). With equal-norm keys they coincide exactly.
• What attention adds: the kernel is learned (through W_Q and W_K), the values are learned, there are many heads, and layers stack. Tsai et al. (2019) develop this kernel view.
• Next they'll ask: "What does the kernel view predict?" → Outputs stay in the convex hull of the values; a uniform row is a plain mean with variance 1/T (Drill 5); a low d_k limits what can be selected (Drill 7).

Q25 · Is attention linear?

A colleague argues: "Attention is a weighted sum, so it's a linear layer. That's why transformers need the feed-forward network for nonlinearity." Half right?

• Concept tested: linear in V, nonlinear in Q and K.
• Intuition: a linear layer's weights are fixed; attention's weights are recomputed from the input every time.
• Answer: half right. For fixed weights the output is linear in the values. But the weights are a softmax of products of the input with itself, so as a function of X attention is strongly nonlinear:

X

the input, one row per token

W_Q W_K^\top

the only combination of the query and key maps that affects the scores

W_V

the value map

• Why the FFN still matters: each attention output is a convex combination of the values, so it stays inside their convex hull and can't create a new direction per token; the FFN transforms each position on its own. Without FFNs and residuals, stacked attention collapses (Q41).
• Next they'll ask: "Where does X appear?" → Three times: twice inside the softmax, once in the values. That data-dependent weighting is what a linear layer can't do (Section 7).

Q26 · Why not 1/d_k?

You're reading code that follows μP (Yang et al., 2022) and divides the scores by d_k, not \sqrt{d_k}. A teammate says: "So the original Transformer paper was wrong." Which scale is right?

• Concept tested: the independence assumption behind \sqrt{d_k} (Card 16).
• Intuition: random vectors add their wobbles (\sqrt{d_k}); aligned vectors add their lengths (d_k).
• Answer: both, under different assumptions. With independent \mathbf q and \mathbf k (initialization), \mathbf q^\top \mathbf k grows like \sqrt{d_k}, so 1/\sqrt{d_k} keeps it O(1). If training aligns them, the dot product grows like d_k:

\mathbb E[\mathbf q^\top \mathbf k]

the average dot product when the key points the same way as the query

d_k

the head width (128 here)

• Simulated at d_k = 128 with \mathbf k = \mathbf q: \mathbf q^\top \mathbf k/\sqrt{d_k} averages 11.3, \mathbf q^\top \mathbf k/d_k averages 1.00. At initialization, 1/d_k overshoots: entropy 8.314 against the uniform 8.318, so no key is selected.
• Why μP uses it: μP changes the initialization and per-layer learning rates together so that hyperparameters transfer across widths; 1/d_k is part of that package, not a drop-in swap.
• Next they'll ask: "How do production models handle late-training growth instead?" → They bound the logits: QK-norm or soft-capping (Q21).

Q27 · The SDPA mask polarity bug

*A team ports a model from nn.MultiheadAttention to F.scaled_dot_product_attention and passes the same boolean padding mask (True = padding). After the port, short sequences produce garbage, as if they attend only to padding. What happened?*

• Concept tested: PyTorch mask conventions. Hook: Polarity, Placement, Phase.
• Intuition: the same True means "keep out" in one API and "come in" in the other.
• Answer: Polarity. In SDPA, a boolean mask's True means may attend. In nn.MultiheadAttention's attn_mask and key_padding_mask, True means blocked. Hugging Face's attention_mask uses 1 = keep. Passing the padding mask unchanged flips it: real tokens are blocked and padding is allowed. Fix: pass ~mask, or convert to a float mask (0 allowed, a large negative number blocked), which SDPA adds to the scores. All of this was checked on PyTorch 2.14.1.
• The other two P's: Placement: is_causal=True with fewer queries than keys uses a top-left-aligned mask, wrong for cached decoding (Q36). Phase: dropout_p is applied whenever you pass it; model.eval() doesn't turn it off, so pass dropout_p=0.0 at inference. Also, is_causal=True together with an attn_mask raises an error.
• Next they'll ask: "What test catches all three?" → Run a padded batch and each sequence alone; outputs at real positions must match to float tolerance.

Q28 · The translator that forgot long sentences

You inherit a legacy translation system: a 2014-style LSTM encoder–decoder. It's fine on short sentences but its output degrades on sentences past about 30 words. A new hire asks: "Why did attention fix this, and why did the 2017 Transformer then throw away the LSTM entirely?"

• Concept tested: why attention exists: the fixed-vector bottleneck, then the sequential bottleneck.
• Intuition: the old system had to summarize a whole paragraph on one index card before translating; attention let it keep the paragraph open on the desk.
• Answer: three steps, each fixing the last one's flaw.

• Bahdanau attention, in their notation, scores each source position j for target step t and builds a context vector:

e_{tj}

the score for source word j when writing target word t

t

the decoder step: which target word is being written

\mathbf v_a

a learned vector that turns the hidden layer into one score

W_a

a learned matrix applied to the decoder state

\mathbf s_{t-1}

the decoder's state before step t: plays the query

U_a

a learned matrix applied to the encoder states

\mathbf h_j

the encoder's state at source word j: plays both key and value

\mathbf c_t

the context vector: a weighted average of the encoder states

\sum_j

add up over every source position j

\alpha_{tj}

the softmax of e_{tj} over j: the attention weights

• Why: the fixed vector was a capacity bottleneck, and attention removed it. Recurrence was a time bottleneck (sequential computation, long paths between distant words), and the Transformer removed it at the price of O(T^2) pairs (Section 1).
• Next they'll ask: "What's the bottleneck now?" → The T^2 memory (solved exactly by online softmax) and the KV cache at inference (Q29).

Q29 · Why share keys and values?

An LLM-infra interviewer: "MQA and GQA make a model slightly worse at the same size. Why were they invented? Make it concrete for Llama 2 70B serving 8 conversations of 4,096 tokens on A100s."

• Concept tested: why MQA and GQA exist: decoding is memory-bound, and the KV cache is what grows.
• Intuition: 64 readers each carrying their own copy of the whole library is wasteful; let groups of readers share one copy.
• Answer: Shazeer (2019, "Fast Transformer Decoding: One Write-Head is All You Need") observed that incremental decoding is limited by memory bandwidth, because every step reloads the cached keys and values. Multi-query attention (MQA) keeps one key/value head for all query heads. Ainslie et al. (2023) noted that MQA can degrade quality, and that training a separate model just for faster inference is unattractive, and proposed grouped-query attention (GQA): H_{kv} groups between 1 and h. Llama 2 70B uses 8 KV heads for 64 query heads (Touvron et al., 2023).

H_{kv}

the number of key/value heads (1 = MQA, h = ordinary multi-head attention)

h

the number of query heads (64 for Llama 2 70B)

• Concrete: each decode step for 8 sequences of 4,096 tokens reads the 138 GB of weights plus the cache. With full MHA the cache is 85.9 GB, 38% of the bytes, and the step takes at least 110 ms at one A100's 2,039 GB/s. With Llama 2 70B's GQA it is 10.7 GB (7.2%, 72.9 ms); with MQA, 1.34 GB (1.0%, 68.3 ms). (Split across n GPUs with tensor parallelism, each time falls by about n; the shares don't change.) Per token: 2.5 MiB (MHA), 320 KiB (GQA), 40 KiB (MQA).
• Why GQA won: most of MQA's saving at a fraction of its quality cost; and an existing MHA checkpoint can be converted by mean-pooling key/value heads within each group and continuing training with about 5% of the original pre-training compute (Ainslie et al., 2023). Hook: share K and V, keep Q (Section 13).
• Next they'll ask: "Would you convert our 7B model?" → Q31.

Level 3 — Senior

Q30 · How many chats fit?

You serve Llama 2 70B in fp16 on two A100 80GB GPUs. Chats run up to 4,096 tokens. Product wants 50 concurrent users at launch. How many can you actually hold, and what are your options?

• Concept tested: KV-cache budgeting with GQA.
• Intuition: the GPUs are a suitcase; the weights go in first, and every open conversation needs its own bag of keys and values.
• Answer: the weights take 138 GB, leaving 22.05 GB (20.53 GiB). With GQA each 4,096-token conversation needs 1.25 GiB of cache, so:

N_{\max}

the most full-length conversations that fit at once

T

tokens per conversation (4,096)

• 16 users, not 50. If Llama 2 70B used full multi-head attention (10 GiB per conversation), it would be 2. These are upper bounds: they ignore activations, fragmentation and framework overhead.
• Options: an int8 KV cache halves the per-token bytes (32 users); paged KV allocation (vLLM, Kwon et al., 2023) only spends memory on tokens that exist, which helps a lot when most chats are shorter than 4K; or add GPUs. Fifty full-length chats need 62.5 GiB of cache, 205.1 GB in all: 3 GPUs by bytes. In practice tensor parallelism splits the 64 query heads and 8 KV heads evenly, so a degree that divides 8 (4 GPUs) is the realistic choice.
• Why: the cache, not the weights, is what scales with users × length (Section 13).
• Next they'll ask: "What if the average chat is 1K tokens?" → With paging, roughly four times as many fit, as long as you cap or evict the long ones.

Q31 · Convert to GQA or retrain?

Your company's 7B model uses full multi-head attention with Llama 2 7B's shape. Long-context serving is too expensive. Leadership asks whether to retrain from scratch with grouped-query attention, or whether there's a cheaper path.

• Concept tested: GQA uptraining (Ainslie et al., 2023) and its cost-benefit.
• Intuition: you don't rebuild the house to share a kitchen; you knock down a few walls and let people adjust.
• Answer: convert, then uptrain. Ainslie et al. (2023) mean-pool the key and value heads within each group and continue pre-training with about 5% of the original compute, reporting quality close to the multi-head model and speed close to MQA. For 8 KV heads (4 query heads per group):

d

the model width (4,096)

H_{kv}

the number of key/value heads after conversion (8)

d_h

the head width (128)

• The cache drops from 512 to 128 KiB per token; parameters fall by 805M (6.74B → 5.93B), so the weights are 11.87 GB; full-length 4K conversations on one 80 GB GPU go from 30 to 126.
• Choosing the group count: keep H_{kv} divisible by your tensor-parallel degree so each GPU holds whole KV heads; evaluate long-context retrieval specifically, since sharing keys costs most there.
• Next they'll ask: "Why not MQA (H_{kv} = 1)?" → More quality loss for little extra saving once the cache is a small share of the bytes (Q29).

Q32 · Prefill fast, decode slow

A profile of your Llama 2 7B server shows the GPU's FLOP utilization is high while processing prompts (prefill) and only a few percent while generating tokens (decode). Is something misconfigured?

• Concept tested: arithmetic intensity; compute-bound vs memory-bound.
• Intuition: prefill is a chef cooking for a banquet from one trip to the pantry; decode is walking to the pantry for every single plate.
• Answer: no; that's expected at small batch. Compare the work done per byte read with the hardware's balance point:

I

arithmetic intensity: FLOPs done per byte read from memory

• Prefill at T = 4{,}096 reuses every weight for all 4,096 tokens: about 1,365 FLOP/byte for the projections, far above the A100's ridge of 153, so it is compute-bound. Decode at batch 1 uses each 2-byte weight for one multiply-add: about 1 FLOP/byte, so it is memory-bound. Attention over the cache does h/H_{kv} FLOPs per byte: 1 for Llama 2 7B, 8 for Llama 2 70B.
• Fixes: batch many sequences per decode step (continuous batching), share K and V (GQA), quantize weights and cache (fewer bytes), and speculative decoding (LLM Ch 5).
• Next they'll ask: "Why doesn't batching help the attention part as much?" → Each sequence has its own cache, so attention bytes grow with the batch; only the weight reads are shared.

Q33 · 128K context: where does memory go?

A team extends Llama 3 8B to 131,072-token training sequences. With their hand-written attention, training runs out of memory even at batch size 1 on 80 GB GPUs. Where is the memory going?

• Concept tested: the T^2 memory of materialized attention, and the exact fix (Card 12).
• Intuition: they are writing down every pair's weight, and the number of pairs is the square of a very large number.
• Answer: a naive implementation stores the weight matrix for the backward pass:

P

the attention weight matrix of one layer, one T \times T block per head

h

the number of heads (32)

T

the sequence length (131,072)

b

bytes per number (2 for bf16)

• That is 32 GiB per head and 1 TiB per layer. The fix is exact, not approximate: online softmax (the math inside FlashAttention; Dao et al., 2022) keeps one log-sum-exp per row, 16 MiB per layer, and recomputes the weights in the backward pass (Worked problem 2). After that, the remaining memory is ordinary activations that grow linearly in T, handled with activation checkpointing; at inference the KV cache is 16 GiB at this length.
• Why: d² for the weights, T² for the pairs (Section 12).
• Next they'll ask: "Does FlashAttention reduce FLOPs?" → No; it reduces memory and memory traffic. The speedup comes from fewer reads and writes (LLM Ch 5).

Q34 · The first token hogs attention

Interpretability work on your chat model finds heads putting around 80% of their attention on the first token, which carries no meaning. Separately, ops ships a streaming cache policy, "keep only the last 4,096 tokens", and quality collapses the moment the first token is evicted. Connect the two.

• Concept tested: attention sinks (the A in S-C-A-L-E).
• Intuition: a softmax row must spend all its money; a head with nothing to buy parks it in the same safe place every time.
• Answer: the first token is an attention sink (Xiao et al., 2023). Because weights must sum to 1, a head that has nothing useful to read puts its mass on an always-visible token, usually the first. That's learned behaviour, not a bug. Evicting the sink forces that mass onto real tokens, which changes every such head's output. StreamingLLM keeps the first 4 tokens plus a recent window. The logit lead a sink needs to absorb a share w against T - 1 equal keys:

w

the share of the row the sink takes (0.9 here)

e

Euler's number, about 2.718

a

the sink's logit lead over every other key

T

the number of keys in the row

• For 90% at T = 1{,}024 that's a lead of 9.13: easy to learn.
• Next they'll ask: "Other places sinks show up?" → BERT heads attending to [SEP] and punctuation as a no-op (Clark et al., 2019). Fixes that give the softmax another outlet change the architecture, so measure before you adopt one.

Q35 · Quality drops on long documents

Your document-QA accuracy falls as context grows, even well inside the trained length. Logged attention shows the heads that should find the answer becoming flatter on longer inputs. Explain the mechanism with a number, and list fixes.

• Concept tested: dilution: a fixed logit advantage buys less as T grows.
• Intuition: a voice that stands out in a room of 100 is lost in a stadium.
• Answer: a relevant key with a fixed logit lead a over T - 1 distractors gets weight e^a/(e^a + T - 1). With a = 5: 0.539 at T = 128, 0.127 at 1,024, 0.035 at 4,096 and 0.0045 at 32,768. To keep a 50% share, the lead must grow with the log of the length:

a_{1/2}

the logit lead a key needs to take half the attention

T

the number of keys in context

• That is 4.84, 6.93, 8.32 and 10.4 at those lengths.
• Fixes: scale logits with \log T (Chiang and Cholak, 2022, proposed this kind of length-dependent scaling); train on long sequences so heads learn larger leads; shorten the context with retrieval and reranking; put the question near the evidence.
• Next they'll ask: "Is that the only long-context failure?" → No: position encodings may not extrapolate past the trained length (LLM Ch 5), and a sliding window can't see far tokens directly (Q39).

Q36 · The cached decoder disagrees with itself

*Your speculative-decoding verifier sends 4 draft tokens through the target model in one call, on top of a 1,000-token KV cache. Its logits don't match step-by-step decoding. The attention call is F.scaled_dot_product_attention(q, k, v, is_causal=True) with q of length 4 and k, v of length 1,004.*

• Concept tested: Placement: causal-mask alignment for non-square attention.
• Intuition: the mask was drawn from the top-left corner, as if the 4 new tokens were the first 4 tokens of the conversation.
• Answer: with is_causal=True and unequal query and key lengths, SDPA uses a top-left-aligned mask: new query i may see only keys 0 to i, not the cache (checked on PyTorch 2.14.1). Cached decoding needs the bottom-right alignment:

i

the index of a new query, 0 to T_q - 1

j

the index of a key, 0 to T_k - 1

T_k

the number of keys: cache plus new tokens (1,004)

T_q

the number of new queries (4)

• Pass an explicit boolean mask, or torch.nn.attention.bias.causal_lower_right(T_q, T_k). Also check that rotary positions start at the cache length, not 0.
• Why: the full-sequence pass and the cached pass must compute the same function; only the mask's alignment differs here (Section 14).
• Next they'll ask: "What test do you keep in CI?" → Decode with the cache and compare against one full causal forward pass on the same tokens; they must agree to float tolerance.

Q37 · The regulator wants attention maps

A bank's model-risk team is reviewing a transformer credit model. To produce the specific reasons required when an applicant is declined, they propose showing the top-attended tokens as the reasons. What do you advise?

• Concept tested: attention vs attribution, in a governance setting.
• Intuition: attention shows where the model looked, not what moved the score.
• Answer: don't use raw attention as reasons. The weights aren't faithful attributions (Jain and Wallace, 2019; Wiegreffe and Pinter, 2019); they mix across heads and layers, ignore what the values carry, and skip the residual stream and the FFNs. Reasons should be tested against the model's behaviour: input attributions (integrated gradients, SHAP-style methods) checked with ablations or counterfactuals ("with this feature changed, does the decision flip?"). Where the reasons must be robust, a simpler, constrained model on structured features is often the better answer.
• Why: a reason has to be something that, if different, would change the decision; attention weights don't promise that (Section 4).
• Next they'll ask: "Is attention useless for interpretability then?" → No: circuit analysis (the QK and OV circuits; Elhage et al., 2021) and induction heads (Olsson et al., 2022) study it carefully, with interventions, not heatmaps alone.

Q38 · Prune heads to save latency?

Ops proposes removing 8 of the 32 heads in every layer of your Llama 2 7B chat model (batch 1) to cut latency, citing papers showing many heads can be removed. What do you expect it to buy?

• Concept tested: head redundancy vs where decode time goes.
• Intuition: you'd be cutting 25% of the attention, which is a small slice of what each token has to read.
• Answer: Michel, Levy and Neubig (2019) found many heads can be removed at test time with little loss, and Voita et al. (2019) found a few specialized heads do most of the work, so quality may hold. But decode at batch 1 is memory-bound:

\text{tokens/s}

the decode speed at batch 1

\text{weight bytes}

every weight is read once per generated token

• Pruning removes 537M parameters, 8.0% of the weights (1.07 GB), so the weight-only bound rises from 151 to 164 tokens/s, about 9%. The cache shrinks from 512 to 384 KiB per token, and prefill FLOPs at 4K fall by 10.7%. Ragged head counts per layer also complicate kernels.
• Better levers: quantize weights and cache (halves the bytes), GQA uptraining (Q31), or speculative decoding. If you do prune, re-evaluate on long-context tasks, where specialized heads matter most.
• Next they'll ask: "When would pruning pay?" → Compute-bound settings (large-batch prefill), where FLOPs, not bytes, are the limit.

Q39 · Sliding window or full attention?

You're building a 32K-context document assistant on a Mistral-7B-sized model (32 layers, 8 KV heads, head dimension 128). One option is full causal attention; the other is Mistral's 4,096-token sliding window. Compare them.

• Concept tested: a mask as a set of allowed pairs (the "Far" in Future, Fake, Far); receptive field vs cost.
• Intuition: each layer can only see 4,096 tokens back, but information can hop back another window at each layer, like a relay.
• Answer: the window allows pairs with

W

the window size (4,096 tokens)

i

the query position

j

the key position

• Cost: a rolling cache holds only W tokens: 0.5 GiB instead of 4 GiB at 32K (128 KiB per token), and the quadratic FLOPs per token per layer drop from 537M to 67M. Reach: stacked layers give a theoretical receptive field of 32 \times 4{,}096 = 131{,}072 tokens (Jiang et al., 2023), but far information arrives second-hand and compressed.
• Trade-off: windows are good when most dependencies are local; exact retrieval of a fact 20K tokens back suffers. Hybrids interleave local and global layers (Gemma 2 does this). LLM Ch 4 covers the rolling cache.
• Next they'll ask: "What breaks with sinks?" → A pure window evicts the first tokens; keep a few sink tokens (Q34).

Q40 · Replace softmax with linear attention?

A paper promises O(T) attention by replacing the softmax kernel with a feature map (Katharopoulos et al., 2020). Your product does needle-in-a-haystack retrieval over long contracts. Should you switch?

• Concept tested: what the softmax kernel buys, seen through the kernel view.
• Intuition: linear attention keeps a fixed-size running summary of the past, which is the RNN's notebook again.
• Answer: replace \exp(\mathbf q^\top \mathbf k) with \phi(\mathbf q)^\top \phi(\mathbf k) and reorder the products:

\mathbf o_i

token i's output

\phi

a fixed feature map applied to queries and keys (Katharopoulos et al. used \text{elu}(x) + 1)

\mathbf q_i

token i's query

\sum_j

add up over the keys j visible to token i

\mathbf k_j

key j

\mathbf v_j^\top

value j, as a row

• The sums don't depend on i, so they're computed once (or kept as a running state when decoding): cost linear in T, constant memory per step.
• What you lose: the exponential kernel can make one key out of thousands dominate; a low-dimensional feature map can't concentrate that sharply, and a fixed-size state must compress the whole past, which reintroduces the bottleneck attention removed (Q28). For needle retrieval, that's exactly the wrong trade. Hybrids with some full-attention layers are the usual compromise.
• Next they'll ask: "What about state-space models?" → Same family of trade-offs: a learned recurrent state instead of a cache (LLM Ch 15).

Q41 · "Is attention all you need?"

A research interviewer: "Defend or attack: a stack of self-attention layers alone, with no feed-forward networks and no residual connections, is a good language model."

• Concept tested: rank collapse; the roles of the FFN and residuals (Card 17).
• Intuition: averaging averages, again and again, makes every token the same.
• Answer: attack. Dong, Cordonnier and Loukas (2021) show that pure attention loses rank doubly exponentially with depth. Measure how far the tokens are from all being equal:

r(X)

the relative residual: 0 means every token's vector is identical

X

the token vectors, one row per token

\mathbf 1

a column of ones, one per token

\bar{\mathbf x}^\top

the average token vector, as a row (bar = average)

\lVert X - \mathbf 1\,\bar{\mathbf x}^\top \rVert_F

how far the tokens are from their average, all together (subscript F = Frobenius norm: a matrix's length treated as one long vector)

\lVert X \rVert_F

the overall size of all the token vectors together

• In the chapter's simulation (32 tokens, width 64, random weights), pure attention goes 0.78 → 0.27 → 0.013 → 2.2 \times 10^{-6} over layers 1–4. With residual connections it is still 0.75 at layer 5 and 0.10 at layer 12: residuals mitigate collapse; they don't abolish it at random initialization. The FFN adds per-token nonlinearity and holds most of the parameters (attention is 31.9% of Llama 2 7B). And without positions, the stack sees a set (Q9).
• Defend the title, though: Vaswani et al. (2017) meant no recurrence and no convolution, not no FFN.
• Next they'll ask: "Why does pre-norm help deep stacks?" → It keeps a clean residual path from input to output; LLM Ch 4 covers pre- vs post-norm.

Level 4 — Research

These questions give you a mathematical setup (a distribution for the queries and keys, a model configuration, a row of scores) and ask for an exact number. There's one method, and every problem below uses it.

The recipe: Count it, scale it, normalize it.

1. Count it. Write every shape. FLOPs are 2mnk for an (m \times k)(k \times n) product; bytes are elements × bytes per element. Most "how big" and "how fast" questions end here.
2. Scale it. For anything random, write the logit as a sum of independent terms and use "variances add" (Card 3). A row whose logits have standard deviation s has entropy about \ln T - s^2/2 while s is small, and goes nearly one-hot once s is well above \sqrt{2 \ln T}.
3. Normalize it. Use softmax algebra: factor out the max (shift invariance), one-hot minus p (the Jacobian), and the log-sum-exp merge (online softmax). These are the S-O-F-T moves of the Derivations tab.
4. Check and break it. Check against the toy or a simulation, then say what breaks at scale (saturation, memory, bandwidth) and the fix.

Worked problem 1 — Saturation at d_k = 128

A classic research-lab whiteboard theme ("why \sqrt{d_k}?"), written here as a composite; it is not a reported question from a named interview.

Setup: one Llama-2-shaped head: d_k = 128, 4,096 keys, query and key components i.i.d. N(0, 1) at initialization. Derive the distribution of the logits with and without the 1/\sqrt{d_k}. Predict the top weight and the entropy of a row. What happens to learning, and what changes once training aligns queries and keys?

Intuition first: a sum of 128 coin-flip-sized terms is about 11 times bigger than one term, and softmax turns gaps into ratios exponentially, so the biggest logit takes almost everything.

Solution.

1. The logits. Each score is a sum of d_k independent products with mean 0 and variance 1:

s_j

the unscaled logit of key j in this row

\mathbf q

the row's query vector

\mathbf k_j

key j's vector

\sum_{c=1}^{d_k}

add up over the d_k = 128 components

q_c

the query's c-th component, N(0, 1)

k_{jc}

key j's c-th component, N(0, 1)

\text{Var}

variance: average squared distance from the mean

d_k

the head width, 128

The simulation gives variance 128.4. Scaled by 1/\sqrt{d_k}, the standard deviation is 1.

1. Two regimes. For logits with standard deviation s over T keys:

H

the entropy of the row in nats; \ln T = 8.318 is perfectly uniform

T

the number of keys, 4,096

s

the standard deviation of the row's logits

Scaled (s = 1): H \approx 8.318 - 0.5 = 7.82; the simulation gives a median of 7.83, about 2,511 effective keys (e^H). Unscaled (s = 11.3, far above 4.08): the formula no longer applies. The top logit sits about 3.6 standard deviations above the mean, and gaps of that size are multiplied by e^{\text{gap}}, so one key takes almost everything: median top weight 0.868 (mean 0.765), 40.5% of rows above 0.9, median entropy 0.49 nats, about 1.63 effective keys.

1. Two keys, exactly. The top weight exceeds 0.99 when the two logits differ by more than \ln 99 = 4.595. Given \mathbf q, the difference is exactly normal:

P

"the probability that"

a_1

the softmax weight on key 1

a_2

the softmax weight on key 2

s_1

key 1's logit

s_2

key 2's logit

\sim

is distributed as

N

a normal distribution with the given mean and variance

\lVert \mathbf q \rVert^2

the query's squared length: itself random, chi-squared with d_k degrees of freedom

c

the scale applied to the logits: 1 unscaled, 1/\sqrt{d_k} scaled

Averaging over \lVert \mathbf q \rVert^2 gives 0.773 unscaled vs 0.00135 scaled at d_k = 128. (Treating s_1 - s_2 as plain N(0, 2d_k) gives 0.774 and 0.00116; a million-draw simulation gives 0.773 and 0.00142.)

What to say out loud: at initialization, \mathbf q^\top \mathbf k has standard deviation \sqrt{d_k} = 11.3. Unscaled, a 4,096-key row is nearly one-hot (median top weight 0.868, entropy 0.49); scaled, it is broad (top weight 0.005, entropy 7.83). The scale keeps logits O(1) so the softmax starts in its sensitive range. Variance adds, so divide by the standard deviation.

Follow-ups that separate senior from research-level answers:

• "So unscaled gradients vanish?" Not simply. The softmax Jacobian's size is non-monotone in d_k without scaling: small for a near-uniform row at d_k = 1 (0.020), largest at d_k = 32 (0.30), then falling as rows saturate (0.181 at 128, 0.041 at 512). At d_k = 128 the unscaled Jacobian is even larger than the scaled one (0.181 vs 0.025), but it routes the gradient into about 1.6 keys instead of about 2,510, and it goes to 0 as a row saturates (22% of rows have a top weight above 0.99 at d_k = 128, 46% at 512). Vaswani et al. (2017) wrote that they "suspect" this effect; say the same (Card 4).
• "Does the argument hold after training?" It assumes independence. If training aligns \mathbf q and \mathbf k, \mathbf q^\top \mathbf k/\sqrt{d_k} averages \sqrt{d_k} = 11.3; 1/d_k keeps it at 1. That's the μP argument (Yang et al., 2022; Drill 6, Card 16).
• "What if the weights grow?" Scaling W_Q and W_K by c scales logits by c^2: entropy 7.82 → 0.339 at c = 4. Entropy collapse and instability follow (Zhai et al., 2023; Q21).
• "How does QK-norm bound it?" With unit-norm \mathbf q and \mathbf k and a logit scale g, logits live in [-g, g], so one key's weight is at most e^g/(e^g + (T-1)e^{-g}): 0.843 for g = 5 at T = 4{,}096 (Drill 18).
• "Inputs not unit variance?" The variances multiply into the logit spread and act like a temperature (Drill 1).

How to identify it in real data: log per-head entropy and the max logit per layer. Entropy near 0 or a top weight near 1 across most rows at initialization means the scale is wrong; the same pattern appearing mid-training means logit growth.

How to mitigate it:

1. Scale by 1/\sqrt{d_k} (or by the measured standard deviation if inputs aren't unit variance).
2. Bound the logits: QK-norm (Henry et al., 2020; Dehghani et al., 2023) or soft-capping (Gemma 2).
3. Train gently: learning-rate warmup and a lower peak rate when entropy starts to fall.
4. Monitor entropy and max logit per head, and alert on fast drops.

Worked problem 2 — Exact attention without the T \times T matrix

A common research-engineering theme ("derive the FlashAttention recurrence"), written as a composite; not a reported question from a named interview.

Setup: one query row over T = 131{,}072 keys (Llama 3.1 8B's context), value width 128. You may keep only O(T) numbers per head. Derive a one-pass algorithm that is exactly softmax attention, prove it, give its memory, and show how to combine results from two GPUs that each saw half the keys.

Intuition first: softmax only needs two running quantities, the largest score so far (for safety) and the sum of exponentials (for the denominator). When a bigger score turns up, everything computed so far is off by one common factor, and one multiplication fixes it.

Solution.

1. The recurrence. Process the keys in blocks B_1, B_2, \dots. After block n keep a running max m_n, a running normalizer l_n and an unnormalized output \mathbf o_n (Card 15):

m_n

the largest score seen after n blocks

m_{n-1}

the largest score seen before block n

B_n

the n-th block of keys

s_j

key j's scaled score for this query

j

a key index

n

the block counter

l_n

the running sum of exponentials, measured relative to m_n

l_{n-1}

the same sum after the previous block

e

Euler's number, about 2.718

e^{m_{n-1} - m_n}

the rescale factor: 1 if the max didn't change, less than 1 if it grew

\sum_{j \in B_n}

add up over the keys in block n

\mathbf o_n

the running weighted sum of values, not yet divided by l_n

\mathbf o_{n-1}

the same sum after the previous block

\mathbf v_j

key j's value vector

Start from m_0 = -\infty, l_0 = 0, \mathbf o_0 = \mathbf 0; the output is \mathbf o_N / l_N after the last block.

1. Proof by induction. Claim: after block n, l_n and \mathbf o_n equal the direct sums over every key seen so far, measured relative to the current max:

B_1 \cup \dots \cup B_n

every key in the first n blocks

It holds for n = 1. If it holds for n - 1, multiplying by e^{m_{n-1} - m_n} turns every old term e^{s_j - m_{n-1}} into e^{s_j - m_n}, and the new block adds its own terms. After the last block, \mathbf o_N / l_N = \sum_j e^{s_j}\mathbf v_j / \sum_j e^{s_j}: the factor e^{-m_N} cancels, so the result is softmax attention, not an approximation. Every exponent is \le 0, so nothing overflows.

1. Two GPUs. Each half returns (m_1, l_1, \mathbf o_1) and (m_2, l_2, \mathbf o_2); merging is the same rescaling:

m_1

the max score on GPU 1's half

m_2

the max score on GPU 2's half

l_1

GPU 1's normalizer, relative to m_1

l_2

GPU 2's normalizer, relative to m_2

\mathbf o_1

GPU 1's unnormalized output

\mathbf o_2

GPU 2's unnormalized output

m

the overall max

l

the merged normalizer

\mathbf o

the merged unnormalized output; the answer is \mathbf o / l

1. Memory. Storing this head's T \times T weights in fp16 at T = 131{,}072 takes 32 GiB. The online form keeps, per row, one log-sum-exp m + \ln l for the backward pass: one fp32 number per row, 16 MiB for a 32-head layer.
2. Check. On a 4,096-key row with value width 128, in 32 blocks of 128: maximum error against direct softmax 7.8 \times 10^{-16} in fp64 and 1.7 \times 10^{-7} in fp32. Block sizes from 1 to 4,096 all agree to within 2.1 \times 10^{-15}, and merging two halves agrees to 1.0 \times 10^{-15}. On the toy's bank row, one key at a time: (m, l) = (0, 1) → river's score 1.414 raises the max, rescale e^{0 - 1.414} = 0.243, (m, l) = (1.414, 1.243) → bank's key: (1.414, 1.736) → output [0.576, 0.284], exactly the direct answer.

What to say out loud: keep a running max, rescale the past. Three numbers per row (max, sum, partial output) make attention exact with O(T) memory, and the same rescaling merges partial results from different devices. FLOPs don't change; the speedup comes from never writing the T \times T matrix to memory.

Follow-ups that separate senior from research-level answers:

• "How does the backward pass work without the weights?" Store the per-row log-sum-exp, then recompute each block's weights e^{s_j - \text{lse}} on the fly during the backward pass; extra FLOPs, no T^2 memory (Card 12).
• "Is this an approximation, like sparse or low-rank attention?" No; it computes the same function, up to floating-point rounding. Sparse and low-rank methods change the function.
• "Why is it faster if FLOPs are the same?" Attention at long T is limited by memory traffic; the fused kernel keeps blocks in fast on-chip memory (Dao et al., 2022; how it does that is LLM Ch 5).
• "How does decoding use it?" Split the cached keys across thread blocks or devices, compute partials, merge with step 3 (Drill 11).
• "fp16?" The max subtraction keeps every exponent \le 0; accumulate l and \mathbf o in fp32 (Drill 17).

How to identify the need in real data: memory that grows with the square of sequence length; out-of-memory errors that move with T, not with batch size; a profiler showing softmax and the score matrix bound by memory bandwidth.

How to mitigate it:

1. Use a fused kernel (PyTorch's F.scaled_dot_product_attention dispatches to FlashAttention or memory-efficient backends when dtype and mask allow).
2. Never materialize P; return attention weights only for debugging.
3. Checkpoint activations for the remaining linear-in-T memory.
4. For decoding, split the cache and merge partials.

Remember: Count it, scale it, normalize it. Logit spread \sqrt{d_k} → one-hot rows unless you scale. A T^2 matrix you can't store → keep a running max and rescale the past.

Level 4 drills

Twenty problems. Drills 1–8 are about the softmax and the scale, 9–12 about masks and the online algorithm, 13–20 about cost. Cover the answer, run the recipe on paper, then check. Each answer is either exact (computed) or, where marked, verified by simulation; every one was also recomputed independently from scratch.

Drill 1 · Hotter than you think

A refactor removed the normalization before the attention projections, and at initialization the query and key components now have variance 2 each. The head still divides by \sqrt{d_k} with d_k = 128, over T = 4{,}096 keys. "What's the variance of the scaled logits, what entropy do you expect, and what's the one-line fix?"

Concept tested: variances multiply into the logit spread.

• Answer: \text{Var}(\mathbf q^\top \mathbf k) = 512, so the scaled logits have standard deviation 2. The small-s formula gives \ln T - 2 = 6.32; the simulated median entropy is 6.43, against 7.83 with unit variances (verified by simulation). Fix: scale by 1/(\sqrt{d_k}\,\sigma_q\sigma_k) = 0.0442, or restore the normalization.
• Intuition: wider inputs make the softmax colder, exactly like lowering a temperature.
• Why:

\mathbf q^\top \mathbf k

the raw dot product

\sigma_q^2

the variance of each query component (2)

\sigma_k^2

the variance of each key component (2)

At s = 2 the formula \ln T - s^2/2 is at the edge of its range, which is why the simulation lands slightly above 6.32.

Drill 2 · Two keys, one winner

To make Worked problem 1 doable on a whiteboard, the interviewer shrinks it to two keys: "d_k = 128, components N(0, 1). What's the probability that the softmax puts more than 0.99 on one key, with and without 1/\sqrt{d_k}?"

Concept tested: saturation, exactly.

• Answer: 0.773 unscaled vs 0.00135 scaled (exact, averaging over \lVert \mathbf q \rVert^2); a million-draw simulation gives 0.773 and 0.00142 (verified by simulation).
• Intuition: a top weight above 0.99 needs a logit gap of only \ln 99 = 4.6; unscaled, the typical gap between two logits is several times that.
• Why: treating the gap as plain normal gives the quick version:

P

the probability that one of the two weights exceeds 0.99

\bar\Phi

the upper tail of the standard normal: P(Z > x)

c

the logit scale: 1 unscaled, 1/\sqrt{128} scaled

That gives 0.774 and 0.00116; the small gap from the exact values comes from the randomness of \lVert \mathbf q \rVert^2, which fattens the tails.

Drill 3 · Jacobian of a confident row

A two-way attention head (think: "attend to the previous token or to the sink") sits at weights p = (0.9, 0.1). "Write the softmax Jacobian, its eigenvalues, and the largest the top eigenvalue can ever be. What is it at p = 0.99?"

Concept tested: the softmax Jacobian \operatorname{diag}(\mathbf p) - \mathbf p\mathbf p^\top (Card 4).

• Answer: J = \begin{bmatrix} 0.09 & -0.09 \\ -0.09 & 0.09 \end{bmatrix}, eigenvalues 0 and 0.18 = 2p(1-p). The maximum is 0.5, at p = 1/2; at p = 0.99 it is 0.0198 (exact).
• Intuition: a confident row barely reacts to nudges in its scores; a 50/50 row reacts most.
• Why:

J

the softmax Jacobian: how each weight responds to each score

p

the weight on the first key (0.9)

The zero eigenvalue is shift invariance (adding the same amount to both scores changes nothing); the other direction is "move mass from one key to the other".

Drill 4 · One gradient by hand

Back to the toy's bank row (weights 0.140 / 0.576 / 0.284 on the, river, bank; keys [0, 0], [1, 0], [1, 1]; values [0, 0], [1, 0], [0, 1]; d_k = 2). The loss falls as bank reads more geography (L = minus bank's geography coordinate), so \mathbf g = [-1, 0]. "Give the gradient with respect to the scores, the query and the values, and say which way a descent step turns the query."

Concept tested: the three backward paths (Card 5).

• Answer: scores: [0.081,\ -0.244,\ 0.164], summing to 0. Query: [-0.057,\ 0.116], so the descent step moves it along [0.057,\ -0.116], toward river's key [1, 0] and away from bank's own [1, 1]. Values: 0.140\,\mathbf g, 0.576\,\mathbf g, 0.284\,\mathbf g, so descent moves each toward geography. A finite-difference check agrees to 2.4 \times 10^{-12} (exact).
• Intuition: values are pushed in proportion to their weight; scores are pulled by surprise, a(v - o); the descent step turns the query toward keys whose scores should rise.
• Why: the score gradients come from Q15's formula with \mathbf g^\top \mathbf o = -0.576. Then:

\frac{\partial L}{\partial \mathbf q}

how fast the loss changes as you nudge each component of bank's query: a slope, not a fraction to divide; points uphill

L

the loss: minus bank's geography coordinate

\frac{\partial L}{\partial s_j}

the score gradient for key j

\mathbf k_j

key j (the key of the is [0, 0], so it drops out)

Drill 5 · The averaging head

A head's query is zero (like the toy's the), so every score is equal. Its T values are i.i.d. with variance 1. "What's the variance of the head's output, and what does that say about a head that never learns to focus?"

Concept tested: uniform attention is a plain mean.

• Answer: 1/T: 2.44 \times 10^{-4} at T = 4{,}096; simulated 2.45 \times 10^{-4} (verified by simulation).
• Intuition: averaging thousands of independent messages leaves almost nothing but their mean.
• Why:

\sum_{j=1}^{T}

add up over all T values

v_j

one component of value j, variance 1

A head stuck near uniform passes on a nearly constant vector: it carries no information about which token is which.

Drill 6 · Aligned at the end of training

Late in training a head's keys have aligned with its queries. Model it as \mathbf k = \mathbf q with components N(0, 1) and d_k = 128. "What's the average of \mathbf q^\top \mathbf k/\sqrt{d_k}? Of \mathbf q^\top \mathbf k/d_k?"

Concept tested: when 1/\sqrt{d_k} is the wrong scale (Card 16).

• Answer: \sqrt{128} = 11.3 vs 1 (simulated 11.3 and 1.00; verified by simulation).
• Intuition: aligned vectors add their lengths, not their wobbles.
• Why:

\mathbb E

the average over random draws

\mathbf q^\top \mathbf q

the query's squared length

So 1/\sqrt{d_k} leaves logits growing like \sqrt{d_k}; 1/d_k keeps them O(1), which is μP's choice (Yang et al., 2022).

Drill 7 · A one-dimensional head

A pruning experiment shrinks one head to d_k = 1: every query and key is a single number. "Across all queries, which keys can ever be a row's top choice?"

Concept tested: the rank limit on what a head can select (Card 10).

• Answer: only two: the key with the largest k (for positive queries) and the one with the smallest k (for negative queries). In 2,000 random heads with 32 queries each, 100% of rows picked one of those two, never more than 2 distinct targets (verified by simulation).
• Intuition: a one-number query can only say "more" or "less", so it can only point at an extreme.
• Why:

s_{ij}

query i's score on key j

q_i

query i, a single number

k_j

key j, a single number

\arg\max_j

the key j that makes what follows largest

\arg\min_j

the key j that makes what follows smallest

With d_k = 2, keys spread on a circle can each be selected (all 12 in the simulation).

Drill 8 · Low-rank logits

"A head has T = 4{,}096 and d_k = 128. What's the largest possible rank of QK^\top? Is the softmax weight matrix limited the same way?"

Concept tested: the rank bound, and what the softmax undoes.

• Answer: \text{rank}(QK^\top) \le 128. The softmax weights are not limited: the exponential is applied entrywise and can raise the rank to full. In a smaller simulated case (T = 64, d_k = 8) the scores have rank 8 and the weights rank 64 (exact bound; verified by simulation).
• Intuition: the logits are built from 128 directions; exponentiating each entry bends them out of that subspace.
• Why:

Q

the T \times d_k matrix of queries

K

the T \times d_k matrix of keys

\text{rank}

the number of independent directions in a matrix

The rank bound still bites at the top: what a head can sharply select depends on the geometry of d_k-dimensional keys (Drill 7).

Drill 9 · What causal masking saves

A team claims their causal attention kernel is "twice as fast as bidirectional, for free". "At T = 4{,}096, what fraction of (query, key) pairs does a causal mask allow, and when does the saving actually happen?"

Concept tested: counting the causal triangle.

• Answer: (T + 1)/2T = 0.500122. The saving is real only if the kernel skips fully masked blocks; a kernel that computes every score and then masks does all the work (exact).
• Intuition: a triangle is half a square, plus the diagonal.
• Why:

T

the sequence length (4,096)

Drill 10 · Online softmax by hand

"Scores [1, 3 \mid 2, 5] arrive in two blocks; the values are the scalars [1, 2 \mid 3, 4]. Trace the running max m, normalizer l and unnormalized output o, then compare with the direct softmax average."

Concept tested: keep a running max, rescale the past (Card 15).

• Answer: after block 1: (m, l, o) = (3,\ 1.1353,\ 2.1353). Block 2 raises the max to 5, so rescale by e^{-2} = 0.13534: (m, l, o) = (5,\ 1.2034,\ 4.4383). Output o/l = 3.6881, the same as the direct answer (exact).
• Intuition: when a bigger score appears, everything so far was measured against too small a max; one multiplication corrects it.
• Why: block 1 gives l = e^{1-3} + e^{3-3} and o = 1\cdot e^{-2} + 2. Block 2:

l

the running sum of exponentials, relative to the current max 5

o

the running weighted sum of values, not yet divided by l

Drill 11 · Merge two GPUs

The same four scores are split across two GPUs. GPU 1 returns (m_1, l_1, o_1) = (3,\ 1.1353,\ 2.1353) and GPU 2 returns (m_2, l_2, o_2) = (5,\ 1.0498,\ 4.1494). "Combine them. Also give the row's log-sum-exp, which the backward pass stores."

Concept tested: the merge is the same rescaling.

• Answer: m = 5; l = e^{-2}(1.1353) + 1.0498 = 1.2034; o = e^{-2}(2.1353) + 4.1494 = 4.4383; output 3.6881; log-sum-exp m + \ln l = 5.1852 (exact).
• Intuition: whoever saw the bigger max sets the reference; the other side is rescaled down to it.
• Why: Worked problem 2, step 3. The order of merging doesn't matter, so partial results can be combined in any tree.

\text{lse}

the log-sum-exp of the row's scores: \ln \sum_j e^{s_j}

Drill 12 · Position from nothing

A decoder has no positional encoding. Its input starts with a BOS token whose value is 1; every other token is identical, with value 0, and all scores are equal. "What does the causal attention output at position t? And without the causal mask?"

Concept tested: the causal mask leaks position (Card 7).

• Answer: 1/t: position 1 sees only BOS (1.0), position 2 gets 0.5, position 8 gets 0.125. Without the mask every position gets 1/T, so no position information (exact).
• Intuition: each token can count how many tokens it is allowed to see, and the mask makes that number its position.
• Why:

o_t

the output at position t

t

the position (BOS is position 1)

\sum_{j=1}^{t}

add up over the t visible tokens

v_j

token j's value: 1 for BOS, 0 for the rest

This is why decoders without positional encodings can still learn order (Haviv et al., 2022; Kazemnejad et al., 2023), though the signal is weak.

Drill 13 · Batch KV budget

A serving team batches 8 sequences of 4,096 tokens in fp16. "KV-cache memory for Llama 2 7B, for Llama 2 70B as shipped (GQA, 8 KV heads), and for 70B if it used full multi-head attention?"

Concept tested: the KV formula (Card 13).

• Answer: 16 GiB, 10 GiB, 80 GiB (exact).
• Intuition: the 70B model is 10× bigger but has fewer KV heads per layer than the 7B (8 vs 32), so its cache is smaller.
• Why:

B

the number of sequences in the batch (8)

T

tokens per sequence (4,096)

L

layers (32 for 7B, 80 for 70B)

H_{kv}

key/value heads (32 for 7B; 8 for 70B; 64 if 70B were full MHA)

d_h

head width (128)

b

bytes per number (2)

Drill 14 · Arithmetic intensity of decode attention

One decode step attends one new query per head over T cached tokens. "How many FLOPs per byte of cache read does attention do, for Llama 2 7B and for Llama 2 70B? Compare with the A100's balance point."

Concept tested: why GQA speeds up decode attention.

• Answer: h/H_{kv}: 1 FLOP/byte for Llama 2 7B, 8 for Llama 2 70B. The A100's ridge is 153, so both are memory-bound (exact).
• Intuition: with sharing, each cached key is reused by several query heads before it leaves the chip.
• Why: scores and the weighted sum cost 2 \times 2\,T d_h FLOPs per query head; the cache read is 2\,T\,d_h\,b bytes per KV head:

I

FLOPs per byte read from the cache

h

query heads (32 for 7B, 64 for 70B)

Drill 15 · When does T^2 win?

"Llama 2 7B: d = 4{,}096, SwiGLU FFN width d_{ff} = 11{,}008, full multi-head attention. At what context length do the quadratic attention FLOPs equal everything else in the layer?"

Concept tested: the crossover (Card 11).

• Answer: T^* = 2d + 1.5\,d_{ff} = 24{,}704 tokens; 49,407 with a causal kernel that skips masked blocks, from 2Td(T + 1) = 8Td^2 + 6Td\,d_{ff}; "double it" gives 49,408, the large-T shortcut (exact).
• Intuition: the quadratic term grows by d per extra token per token; the rest is a fixed cost per token.
• Why:

d

the model width (4,096)

d_{ff}

the FFN width (11,008)

T^*

the crossover context length

Drill 16 · What the backward pass keeps

"Llama 2 7B, T = 4{,}096, batch 1, fp16. If you store the attention weights P of every head in every layer for the backward pass, how much memory is that? How much if you keep only one fp32 log-sum-exp per row?"

Concept tested: naive vs online attention memory (Card 12).

• Answer: 32 GiB vs 16 MiB (exact).
• Intuition: a square of 4,096 per head, against a single column.
• Why:

L

layers (32)

h

heads per layer (32)

T

sequence length (4,096)

Drill 17 · fp16 without the max trick

*A kernel computes exp(scores) directly in fp16 on unscaled d_k = 128 logits (standard deviation 11.3) over T = 4{,}096 keys. "Above what logit does exp overflow? What fraction of logits do that? What fraction of rows end up NaN?"*

Concept tested: why the max subtraction is not optional.

• Answer: \ln 65{,}504 = 11.09. About 16% of logits exceed it (0.164 by the normal approximation; 0.161 simulated). With 4,096 keys, essentially every row has at least one, so the naive fp16 softmax gave NaN in every simulated row; the max-subtracted version gave none (verified by simulation).
• Intuition: fp16 runs out of room fast; a long row almost surely contains one logit that's too big.
• Why:

s

one unscaled logit, standard deviation 11.3

\bar\Phi

the upper tail of the standard normal

One inf in the denominator turns every weight into 0 or NaN. Subtracting the max makes every exponent \le 0.

Drill 18 · QK-norm sets a bandwidth

A model uses QK-norm: queries and keys are scaled to unit length and the logit is g\,\mathbf q^\top \mathbf k. "Show that this is a Gaussian kernel and give its bandwidth. If someone sets g = 1/\sqrt{128} out of habit, what's the most weight any key can get at T = 4{,}096? What g allows 0.9?"

Concept tested: with normalized vectors, the scale is the only temperature.

• Answer: bandwidth 1/\sqrt g: 3.36 for g = 1/\sqrt{128}, so the most any key can get is 2.9 \times 10^{-4}: the head can't select anything. A top weight of 0.9 needs g = 5.26 (exact). This is why QK-norm models learn the scale rather than fixing it at 1/\sqrt{d_k}.
• Intuition: unit vectors live on a sphere; a tiny g makes every point look equally close.
• Why: for unit vectors \lVert \mathbf q - \mathbf k \rVert^2 = 2 - 2\,\mathbf q^\top \mathbf k, so

g

the logit scale applied after normalizing

\lVert \mathbf q - \mathbf k \rVert^2

the squared distance between query and key on the unit sphere

w_{\max}

the best case for one key: it sits at \mathbf q^\top \mathbf k = 1, every other key at -1

e

Euler's number, about 2.718

The constant g cancels in the softmax, leaving a Gaussian kernel with bandwidth h = 1/\sqrt g. Solving w_{\max} = 0.9 gives g = \tfrac12 \ln\!\big(0.9 \cdot 4{,}095 / 0.1\big) = 5.26.

Drill 19 · How loud must a sink be?

Interpretability shows a head that puts 90% of its attention on the first token in 1,024-token contexts, with the rest spread evenly. "How large a logit lead over the other keys does the sink need? What about at 128 and 4,096 tokens?"

Concept tested: attention-sink logits (Card 18).

• Answer: \ln(0.9 \cdot 1{,}023 / 0.1) = 9.13; 7.04 at 128 tokens, 10.51 at 4,096 (exact).
• Intuition: to outshout T - 1 equal voices you need a lead that grows only with \ln T, which is easy to learn.
• Why: set the sink's weight to w and solve:

a

the sink's logit lead over each other key

w

the sink's share (0.9)

T

the number of keys

Drill 20 · The decode speed limit

"Llama 2 7B in fp16 at batch 1 on an A100 80GB with 2,039 GB/s of memory bandwidth. Upper bound on tokens per second with an empty cache, and with a 4,096-token cache?"

Concept tested: decode is memory-bound.

• Answer: 151 tokens/s (6.61 ms per token for 13.5 GB of weights); with a 4K cache (2 GiB), 130 tokens/s (1.05 ms more) (exact).
• Intuition: every token is one full pass over the weights and the cache through a fixed-width pipe.
• Why:

\text{bandwidth}

memory bandwidth, 2,039 GB/s

\text{KV bytes}

the cache read at every step: 0, or 2 GiB at 4,096 tokens

Batching shares the weight read across sequences, which is why serving systems batch decode steps (Q32).

Failure-mode cheat sheet

Rapid fire — the day before

One line each. Cover the right column and answer out loud; anything you hesitate on, go back to its full question above.

Review schedule: go through this table the day after you study, then 3 days later, a week later, and 3 weeks later. Any row you miss goes back to the start.