Explain the transformer and attention at whiteboard level.
- 2 min — draw the stack: embed → blocks → logits
- 2 min — attention formula, causal mask, multi-head
- 2 min — cost shape: O(n²) compute, O(n) memory, KV per token
- 1 min — modern variants and why (GQA, MoE, sliding window)
Each token becomes a vector (embedding plus position, today usually RoPE). A stack of identical blocks (32 for Llama-3.1-8B, 80 for 70B) refines those vectors. Each block has attention — every token forms a query, scores it against the keys of earlier tokens, and takes a softmax-weighted mix of their values — and a feed-forward network (MLP) applied per token, which holds ~2/3 of the parameters and much of the stored knowledge. After the last block, the final vector is projected to logits.
attention(Q,K,V) = softmax(QKᵀ / √d) · V, with a causal mask so tokens see only the past.
Cost shape: attention compute is O(n²) in sequence length. Attention memory is O(n) in practice because FlashAttention tiles the computation in on-chip SRAM and never materialises the n×n score matrix (at 128k tokens that matrix would be 32 GB per head per layer in BF16). Decode memory is dominated by weights plus the KV cache, which grows linearly: 320 KB/token for Llama-3.1-70B in BF16.
A developer-docs assistant moves from 8k to 128k context and TTFT goes from 0.6 s to 9 s. On the whiteboard: prefill MLP cost grows linearly (16× more tokens), attention cost grows quadratically (256× more score computations), and at 128k attention is now a large share of prefill FLOPs. Meanwhile each live 128k session holds 40 GB of KV cache on a 70B model. Conclusion: long context is a capacity-planning decision, so the team caps context at 32k and retrieves the rest.
- The why: attention lets any token pull information from any earlier token in one layer, and the whole sequence trains in parallel — the two things RNNs could not do.
- Modern variants exist to cut serving cost: MoE feed-forward layers (compute), grouped-query attention (KV memory 8× smaller with 8 KV heads vs 64 query heads), sliding-window layers, and FlashAttention kernels (memory).
- Org scope: application teams rarely touch architecture; they own the context-length budget. The serving or platform team owns the kernel stack (FlashAttention version, KV dtype) and publishes the per-token memory number product teams plan against.
- Where it breaks: presenting attention weights as an explanation of the model's reasoning. They are a mechanism and a weak interpretability signal.
KV bytes/token = 2 × layers × kv_heads × head_dim × bytes. Llama-3.1-70B: 2 × 80 × 8 × 128 × 2 = 327,680 B ≈ 320 KB. 128k tokens → 40 GB. Without GQA (64 KV heads) it would be 2.5 MB/token, 320 GB per session — why GQA is standard.
| The | cap | is | 12 | |
|---|---|---|---|---|
| The | 1.00 | · | · | · |
| cap | 0.30 | 0.70 | · | · |
| is | 0.10 | 0.60 | 0.30 | · |
| 12 | 0.05 | 0.55 | 0.15 | 0.25 |
What the answer sounds like at each level
Draws embeddings → N blocks of attention + MLP → logits, and writes the softmax(QKᵀ/√d)V formula with the causal mask.
Explains O(n²) compute vs O(n) memory with FlashAttention, KV cache growth, and GQA's effect on memory.
Turns the cost shape into a context-length policy and capacity plan, and says who owns kernels vs context budgets.
Follow-ups the interviewer will push
1Why divide by √d?
Dot products of d-dimensional vectors grow with √d; without scaling, softmax saturates and gradients vanish. Scaling keeps logits in a trainable range.
2Does FlashAttention make attention linear?
Memory, yes — O(n) by tiling and recomputation. Compute stays O(n²); it is faster because it cuts HBM reads and writes, not FLOPs.
3Now 10× more users at 128k context. What runs out first?
KV cache memory, not compute: 40 GB per 70B session in BF16 means ~4 concurrent sessions per spare 160 GB. FP8 KV halves it; GQA is already baked in; the next lever is context caps per tier.
Red flags interviewers listen for
- Says FlashAttention changes attention's compute complexity.
- Cannot name where most parameters live (MLP).
- Treats attention maps as faithful explanations.
How it's scored
| Dimension | What strong looks like |
|---|---|
| Correctness | Q/K/V, softmax, √d scaling, causal mask, residual stream, MLP role. |
| Systems | Separates compute vs memory complexity and computes KV bytes/token for a real model. |
| Judgement | Links architecture to a context-length or hardware decision. |
Sources
- Attention Is All You Need — Vaswani et al., 2017paper
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al., 2022paper
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — Ainslie et al., 2023paper