Prefill Disaggregation - Understanding from first principles

` In order to generate output tokens from an input prompt, LLM inference is split into two stages: prefill and decode. Prefill runs on the input tokens, populating KV caches, before entering the decode stage that generates tokens one-by-one. While a single decode step typically runs for tens of milliseconds, prefill takes substantially longer. If run on the same devices, mixing prefill with decode degrades decode performance. In this article we explore an established solution in the form of disaggregated prefill and decode, running them on separate devices to maximize both prefill throughput and decode latencies. `

Prefill and decode - simplified.

A modern LLM is what we generally call a decoder only neural network. The analogy of a question and answer makes this better. Lets say the question to an LLM is “How do I find good homes to buy within my budget”. The LLM has to process this entire question in one shot. The LLM then genrates the annswer word by word. Every word generated appends to the question.

Step 1: "How do I find good homes to buy within my budget -> start Step 2: "How do I find good homes to buy within my budget start -> by Step 3: "How do I find good homes to buy within my budget start by -> finding

The initial question has to be processed in every step. This is processed using scarce GPU compute. The next question is why not cache the processed question and all the intermediaries at each step.

That cache is known popularly as the kv cache because the cache contains key and value vectors ( no correlation to a key-value HashMap) Step 1 is called the prefill phase Step 2 - n is called the decode phase.

So prefill is a compute heavy phase and decode is a memory read heavy phase. One is constrained by GPU capacity the other is constrained by memory bandwidth.

Prefill - Decode separation

Let’s assume there are a pool of GPUs and this pool is allocated for the above steps. In reality traffic is coming in batches. Each batch has multiple user requests.

• batch 1 [user1, user2, user3] where each user = [token1, token2, token3] and each token = [1.0, 2,1, 2.3] an vector of floats.
• Similarly batch 2 , batch 3
• On the same GPU pool
  • timestaamp t1 , batch1 -> prefill
  • timestaamp t2 , batch1 -> first decode , batch2 prefill
  • timestaamp t3 , batch1 -> second decode , batch2 first decode, batch 3 prefill So the compute bound prefill ops are now interspersed with memory bound decode ops. This causes delays in the first token generated and delays in intermediate tokens as well. At scale, these delays become hard to measure and optimise.

It’s best to split them into different pools so that one can optimize compute bound parallel ops in Prefill separately and optimize the sequential decode separately. The split is the easy part. The difficulty lies in kv cache . It becomes the new bottleneck. It has to be shared across the prefill pool and the decode pool.

• For every step
  • Read kv-cache -> takes t time.
  • Load kv-cache to local SRAM. This is impossible. To load the data into local SRAM , the data has to be in HBM.
• So kv cache has to be transferred to the local HBM.

The transfer - not so pretty

The kv-cache cannot be shared as such. It has to be transferred from the Prefill pool to the Decode pool.Why transfer ? That needs a bit of GPU layout understanding.

A typical nVidia GPU rack from the bottom up:

• n GPUs in one node
• n nodes in one rack

• Hundreds or thousands of racks in one datacentre.

• Each GPU has a dedicated memory called the high bandwidth memory or HBM of a certain capacity.
• Each GPU within a node is interconnected by a slightly slower bus called NVLINK
• Each node in a rack is interconnected by a bus slower than NVLINK.
• Each rack is connected to the other by an even slower bus Infiniband.

The prefill kv-cache is bounded by the individual HBM capacity of each GPU. For a decode GPU that means, transferring the kv-cache to that respective GPU from the corresponding Prefill GPU. With millions of GPUs across Prefill and decode clusters, this transfer job gets increasingly complex.

Hence, for best utilization, kv-cache is distributed or split across individual GPUs. In modern LLMs this parallelism is pretty straight forwaard and mimics the Mixture of Experts (MOE) architecture.

The payoff

By managing prefill and decodde managed in different GPU pools or clusters, frontier labs and inference providers are able to optimize the times for TTFT -> the time to first token and also optmize the inter token latencies. The optimization is an ongoing process, because the GPU architectures, HBM capacities and rack organizations all change. A H100 rack is very different from a GB200 rack. So the transfer algorithms need to adapt for the new GPU architectures constantly.