[1/2] DistServe: Why Separate Prefill and Decode (Category #Research)

This post is about the DistServe paper by a team from Peking University, UC San Diego, and StepFun, published at OSDI 2024. I already discussed separating prefill and decode in my llm-d breakdown. DistServe does an especially good job of explaining the engineering side of this choice: what we gain by putting the stages on different GPUs, and what we pay for it. Alongside the paper, there is open-source code, so you can inspect not only the authors’ results but also the implementation if you wish.

Returning to the main point, an LLM request has two quite different parts:

• Prefill processes the input text, computes the KV cache, and produces the first response token. There are many positions and they can be processed in parallel; with a sufficiently long prompt, the stage is usually compute-bound.
• Decode emits subsequent tokens one at a time for each request. With a small batch there is relatively little computation, while the weights and cache must be read from memory again and again. Memory bandwidth is often the constraint here. Batching requests helps use the GPU more efficiently.

When both stages run together, a new heavy prefill delays generation already in progress. Give decode priority, and new users wait longer for the response to begin. Both stages also inherit one parallelism configuration even though their needs differ.

In DistServe, the process looks like this: Request → prefill queue → prompt processing and first token → KV-cache transfer → decode queue → generation of the remaining tokens.

Prefill and decode instances each have their own copy of the model weights. The KV cache holds intermediate attention keys and values; with it, another GPU can continue the computation without processing the prompt again. Prefill and decode can now be scaled and configured independently.

Then the familiar life of a distributed system begins: there is a network, extra memory, and one more place where a queue can grow :)

For OPT-66B, the authors estimate that the cache for a 512-token prompt is about 1.13 GB. At 10 requests/s, that is roughly 90 Gbit/s for the cache alone. Matching average demand to bandwidth is not enough: load spikes need headroom. If decode accepts work more slowly than prefill, completed caches also begin to wait and consume memory.

That is why DistServe accounts for topology. When the network between servers is slow, corresponding prefill and decode stages are placed on different GPUs within the same server so the cache can move over NVLink. In the authors’ measurements, more than 95% of requests transferred the cache in under 30 ms.

The authors’ target is goodput per GPU: the maximum request rate at which a specified share of latency requirements, or SLOs, is met. They evaluate TTFT (time to first token) and TPOT (time per output token), the average time for the next token. The target share is usually 90% of requests. An average TPOT, incidentally, still does not guarantee that generation contains no individual pauses.

They evaluated the approach on OPT-13B/66B/175B, a cluster of 32 A100 GPUs, and chat, code-completion, and summarization workloads. In the chat tests, they achieved 2–4.6 times more throughput per GPU than the version of vLLM available at the time; the maximum was 7.4 times versus DeepSpeed-MII. This is a gain for the complete configuration under the selected SLOs. Request arrival times were synthetic and the authors chose the latency thresholds themselves, so the multipliers cannot be transferred to an arbitrary service.

The most interesting question remains: how many GPUs should each stage receive, and how should the model be laid out across them? The second part of the post covers intra-op, inter-op, and queueing theory, which helps explain why the fastest individual request does not by itself determine the best service configuration.

#Research #AI #Architecture #Engineering #DistributedSystems