Inference Engineering interview question · KV cache · mid level
Derive the KV cache size formula and work it out for Llama 3 70B.
Free to read in full. This is 1 of the 120 questions in the Inference Engineering interview pack, with 141 diagrams. Get the pack, ₹999 · $20
Why it is asked
Every capacity plan starts with this formula. The interviewer checks each factor and that you use KV heads, not query heads.
The answer
KV bytes per token = 2 x layers x KV heads x head dimension x bytes per value, which for Llama 3 70B in BF16 (2 bytes per number) is 2 x 80 x 8 x 128 x 2 = 327,680 bytes, about 0.33 MB per token.
Where each factor comes from:
- 2: one key vector and one value vector are stored.
- layers: every transformer layer has its own attention, so its own K and V. Llama 3 70B has 80.
- KV heads: the number of key-value heads, not query heads. Llama 3 70B has 64 query heads but uses grouped-query attention (GQA) with 8 KV heads, each shared by 8 query heads.
- head dimension: the length of each head's vector, 128 here. So each layer stores 8 x 128 = 1,024 values for K and 1,024 for V per token.
- bytes per value: 2 for BF16, 1 for FP8.
Then total KV = bytes per token x tokens per sequence x number of sequences.
For Llama 3 70B:
- 4,096 tokens: about 1.34 GB per sequence.
- 8,192 tokens: about 2.68 GB.
- 16,384 tokens (Llama 3.1 70B, same shape, longer context): about 5.37 GB.
With an FP8 KV cache the figure halves to 163,840 bytes per token. For comparison, Llama 3 8B (32 layers, 8 KV heads) needs 131,072 bytes per token.
Check: if you had used 64 heads instead of 8, you would get about 2.62 MB per token, 8 times too much. This is the most common slip.
What the interviewer listens for
- 2 x layers x KV heads x head_dim x bytes
- Use KV heads (8), not query heads (64)
- 327,680 bytes per token for Llama 3 70B in BF16
- Multiply by tokens and by concurrent sequences
Follow-ups the interviewer may ask
- How does tensor parallelism split this? Hint: KV heads are divided across GPUs, so with TP=8 each GPU holds one KV head per layer.
- Does the formula hold for DeepSeek-V3? Hint: no; its multi-head latent attention stores a compressed latent, so use its own formula.
Common mistakes
- Using the number of query heads
- Forgetting the factor 2 for K and V
- Mixing GB (10^9 bytes) and GiB (2^30 bytes) in the same calculation
Written by Roni Das, who builds distributed systems and AI systems for a living.
Learn it properly
This is taught step by step in the course lesson The Context Window: What Happens When a Prompt Does Not Fit.
One question is not enough for the interview.
This is 1 of the 120 questions in the Inference 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 (141 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
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)
- Why are LLMs bad at counting letters or reversing a word? (GenAI and LLM)
- Write out scaled dot-product attention with the shape of every tensor. (GenAI and LLM)
- How does speculative sampling keep exactly the target model's distribution? (GenAI and LLM)
- What is GRPO, and how do verifiable rewards train reasoning models? (GenAI and LLM)
- Design a system that answers questions from 300-page contracts accurately. (GenAI and LLM)
- 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)