Mid-level interview questionAsked at Anthropic · interview write-up
Implement scaled dot-product attention in PyTorch, then tell me where it can go wrong numerically.
Model answer
def sdpa(q, k, v, mask=None): # mask True = may attend
s = q @ k.transpose(-2, -1) / q.size(-1) ** 0.5
if mask is not None:
s = s.masked_fill(~mask, float("-inf"))
s = s - s.amax(-1, keepdim=True)
p = s.exp()
return (p / p.sum(-1, keepdim=True)) @ v
Scale by √d so scores don't saturate softmax. Subtract the row max before exp so nothing overflows; torch's softmax does this internally. A fully masked row gives NaN, since max minus max is inf minus inf, so guard it. In FP16, compute scores and softmax in FP32. FlashAttention does the same blockwise with a running max.
Practise the rest
20 Attention questions, 492 in all, each with a model answer. Free.
More from Anthropic interview questions.
More Attention interview questions
- Walk me through scaled dot-product attention. What are queries, keys and values, and why divide by the square root of d?
- What's inside one transformer block, and where do most of its parameters sit?
- What does the causal mask do, and what does it imply for a decoder model reading 'The ladder did not reach the window because it was too short'?
- A teammate shows you a GPT-2 head that puts most of a pronoun's attention on the right noun and concludes that the head resolves pronouns. How do you respond?
- Product wants to double the context you support, from 1,000 to 2,000 tokens. From the attention side, what gets more expensive, and what families of fixes would you weigh?
- Implement batched multi-head self-attention with a causal mask.