GenAI and LLM interview question · Self-attention math · mid level
Write out scaled dot-product attention with the shape of every tensor.
Free to read in full. This is 1 of the 150 questions in the GenAI and LLM Engineering interview pack, with 173 diagrams. Get the pack, ₹999 · $20
Why it is asked
Tests precision: can you write the formula and track every shape, including heads and GQA, the way you would when debugging a kernel?
The answer
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k) + M) V, where M is the mask (0 where allowed, minus infinity where blocked).
Each tensor (a multi-dimensional array of numbers) has a shape. Use batch B, sequence n, model width d, h query heads of size d_k (d = h x d_k). Example: Llama 3 8B, B = 1, n = 8,192, d = 4,096, h = 32, 8 key-value heads, d_k = 128.
- Input X: [B, n, d] = [1, 8192, 4096].
- Q = X W_Q, reshaped to [B, h, n, d_k] = [1, 32, 8192, 128].
- K = X W_K and V = X W_V: [B, 8, n, 128]. With GQA (grouped-query attention) each KV head is read by its 4 query heads (the snippet copies it with np.repeat for clarity; real kernels do not copy).
- Scores Q K^T: [B, h, n, n] = [1, 32, 8192, 8192]. Multiply by 1/sqrt(128) = 0.0884 and add the causal mask.
- Softmax over the keys (last axis): each row sums to 1.
- Weights x V: [B, h, n, d_k] = [1, 32, 8192, 128].
- Transpose and merge heads: [B, n, h x d_k] = [1, 8192, 4096], then the output projection W_O keeps [B, n, d].
The shape that matters for cost is the score tensor. Stored in BF16 (2 bytes per number) it is 32 x 8,192 x 8,192 x 2 bytes = 4.29 GB for one layer of one sequence, which is why kernels such as FlashAttention, which works in small on-chip tiles, never store it.
What the interviewer listens for
- softmax(QK^T / sqrt(d_k) + mask) V
- Q is [B, h, n, d_k]; K and V use the number of KV heads
- Scores are [B, h, n, n]: quadratic in sequence length
- Softmax is over the key axis, one distribution per query
- Heads are merged back to [B, n, d] and projected by W_O
Follow-ups the interviewer may ask
- What changes for decode? Hint: Q has one row; scores are [B, h, 1, context].
- Where does d_v enter? Hint: V's head size; usually equal to d_k but it does not have to be.
- How large is the score tensor for 32 layers if kept for backward? Hint: about 137 GB at 8K, which is why it is recomputed.
Common mistakes
- Softmax over the query axis instead of the key axis
- Scaling by d_k instead of sqrt(d_k)
- Giving K and V the query head count in a GQA model
Written by Roni Das, who builds distributed systems and AI systems for a living.
One question is not enough for the interview.
This is 1 of the 150 questions in the GenAI and LLM Engineering pack
Each question has why it is asked, a full answer, what the interviewer listens for, the follow-ups, the common mistakes and a diagram (173 diagrams in all). A PDF you keep, junior to staff level.
India: UPI, cards or netbanking through Razorpay. Elsewhere: PayPal. Download at once from My packs.
More free questions from this pack
- Why are LLMs bad at counting letters or reversing a word? (junior)
- How does speculative sampling keep exactly the target model's distribution? (senior)
- What is GRPO, and how do verifiable rewards train reasoning models? (senior)
- Design a system that answers questions from 300-page contracts accurately. (staff)
Free questions from the other packs
- What do temperature and top-p control, and what values would you use for data extraction versus brainstorming? (AI Engineering)
- What is hybrid search, and how do you merge keyword and vector results with Reciprocal Rank Fusion? (AI Engineering)
- Users ask "list all X" questions and RAG returns only five items. Why, and what do you do? (AI Engineering)
- How do you check that an LLM judge agrees with human reviewers? (AI Engineering)
- Cut the cost of an LLM feature serving 10 million requests a day by 60% without hurting quality. Walk through your plan. (AI Engineering)
- What happens, step by step, when an LLM server handles one request? (Inference Engineering)
- Derive the KV cache size formula and work it out for Llama 3 70B. (Inference Engineering)
- What is PagedAttention? (Inference Engineering)
- When does speculative decoding stop helping or make things worse? (Inference Engineering)
- Design a service to serve Llama 3 70B at 10 QPS with latency targets. (Inference Engineering)
- Compute the cosine similarity of a = (3, 4) and b = (4, 3) by hand. What does the result mean? (Maths for ML)
- A batch of 32 MNIST images goes into a dense layer with 128 units. What are the shapes, and what is the rule for when two matrices can be multiplied? (Maths for ML)
- A disease affects 1% of people. A test catches 90% of sick people (sensitivity) and wrongly flags 9% of healthy people (false positive rate). You test positive. What is the chance you are sick? (Maths for ML)
- Why do we not use mean squared error for classification with a sigmoid output? Compare the gradients numerically. (Maths for ML)
- Compute scaled dot-product attention by hand for two tokens with d_k = 2. (Maths for ML)
- A customer says: "We want AI for our support team." What do you do in the first meeting? (Forward Deployed Engineer)
- The customer's admin offers you a global admin account "to save time". What do you ask for instead, and why? (Forward Deployed Engineer)
- Your connector writes tickets into the customer's system. A request timed out, the connector retried, and now there are duplicate tickets. How do you make the writes safe to retry? (Forward Deployed Engineer)
- The customer wants RAG over 10 million documents, all behind their firewall. Size it for me on the whiteboard. (Forward Deployed Engineer)
- One deployment of your AI platform must serve 50 business units of the same enterprise. How do you design multi-tenancy? (Forward Deployed Engineer)