Explain the KV Cache and Prefill-Decode Disaggregation in LLM Serving

This round for a machine learning engineer role focused on ML infrastructure rather than modeling. The interviewer asked conceptual questions about serving large language models, in particular the KV cache and prefill-decode (PD) disaggregation. Answer both parts as you would in a technical conversation: explain the mechanism, quantify it where you can, and discuss the trade-offs. ### Clarifying Questions - Which model shape should estimates assume: number of layers, number of attention heads and key-value heads (full multi-head attention or grouped-query attention), head dimension, and numeric precision? - What are typical prompt and output lengths, and how many requests are served concurrently? - Which latency targets matter most: time to first token, time per output token, or overall throughput? - What hardware and interconnect link the GPUs that would hold the prefill work and the decode work? ### Part 1 — The KV cache Explain what the KV cache stores during autoregressive generation with a decoder-only transformer, why it makes decoding faster, and how much memory it takes. Derive the cache size per token and per request for a model shape you state, and explain what that size means for batch size, context length and GPU memory management. ```hint Count the work at step t Write down what attention at generation step t needs from the earlier tokens, and which of those quantities would be recomputed at every step if nothing were saved. ``` ```hint Size it by its dimensions The cache holds tensors for every layer and every past token. Multiply out their dimensions and the bytes per element. ``` #### What This Part Should Cover - What is cached, per layer and per head, and what is deliberately not cached - The computation saved per decoding step and the memory paid for it - A per-token size formula with a worked number for a stated model - Consequences for batching and for managing GPU memory as requests grow ### Part 2 — Prefill-decode disaggregation Explain the two phases of LLM inference, prefill and decode, and how their resource profiles differ. Then explain PD disaggregation, which runs prefill and decode on separate groups of GPUs: why do it, what does it cost, and when is it not worth it? ```hint Compare the shapes of the work How many tokens per request does one forward pass process in prefill, and how many in decode? Ask which resource, compute or memory bandwidth, each phase runs out of first. ``` ```hint Something must move When prefill finishes on one machine and decoding continues on another, think about what state has to follow the request and how large it is. ``` #### What This Part Should Cover - Compute-bound prefill versus memory-bandwidth-bound decode, tied to the latency metric each phase drives - The interference that appears when both phases share the same GPUs - The cost of moving state between the two pools, and how to hide or reduce it - When keeping both phases on the same GPUs is the better choice ### What a Strong Answer Covers - States or asks for the model shape and workload before quantifying anything - Correct mechanics: what is cached and why, and what happens in each phase - Concrete arithmetic for the cache size and the transfer cost - A clear link from each mechanism to time to first token, time per output token and throughput - A balanced view of when disaggregation helps and when it only adds complexity ### Follow-up Questions - Requests have very different lengths and GPU memory keeps fragmenting. How would you allocate KV cache memory so that capacity is not wasted? - Many requests share the same long system prompt. How can the serving system avoid recomputing and re-storing its KV cache for every request? - The KV cache limits your batch size. Which techniques shrink it, and what does each cost in quality or compute? - How would you decide how many GPUs go to the prefill pool versus the decode pool, and adjust that split as traffic changes?