Scaling Out: Parallelism, Chunked Prefill & Disaggregation

Once a model is too big for one GPU, or traffic is too big for one server, you must spread the work. How you spread it matters, because the two phases of inference want different things. This chapter covers the main ways, and one problem that surprises people: a single long prompt can freeze everyone else's stream.

A highway with one lane works until a very long truck pulls on. Every car behind it slows down. You could let the truck enter in sections between cars, or give trucks their own lane. Those are chunked prefill and disaggregation. When a single lane is not enough to carry the weight, you split the load across several lanes, which is parallelism.

• -Tensor parallelism splits each layer's big matrices across GPUs, and the GPUs combine their partial results at every layer. Each token gets faster, since the total memory bandwidth is shared, but it needs a very fast connection between GPUs, so it is usually kept inside one server.
• -Pipeline parallelism puts different layers on different GPUs, like stages on an assembly line. Communication is lighter, but there are idle gaps called bubbles unless many micro-batches are in flight.
• -Data parallelism copies the whole model to several groups of GPUs and spreads requests across them. It raises throughput, not the size of a model that can be served.
• -Mixture-of-experts models add expert parallelism, placing different experts on different GPUs.

The earlier code showed why you need it: a 70B model in 16-bit needs about 141 GB for weights, more than one 80 GB GPU. After weights, you also need room for the KV cache, which pushes the minimum higher. A common choice is tensor parallelism inside a server for the model, and data parallel replicas of that server for traffic.

In a standard batch, prefill for a long new prompt runs in the same step as everyone else's decoding. That step now takes far longer, and every user in the batch waits for it to finish before their next token. To them it looks like a freeze. The SARATHI paper by Agrawal and colleagues proposed chunked prefill: split the long prompt into chunks and mix one chunk into each decoding step. They reported up to 10 times higher decode throughput for a LLaMA-13B model on one GPU type, and smoother latency.

The code models the trade-off. A busy server runs decoding steps of 25 ms for 30 users. A new user pastes an 8,000-token document, and each prompt token adds 0.02 ms to the step it is processed in. Press Run and compare the longest pause other users feel with how long the new user waits.

With the prompt processed all at once, everyone else stalls for 185 ms and the new user gets a first token after 210 ms. At 512 tokens per chunk, other users never pause more than 35 ms, close to the normal 25 ms, but the new user waits 585 ms. At 256 tokens it is 30 ms against 985 ms. Chunk size is a dial between the two. Pick it from your latency goals, and notice that it moves latency between users and does not make it vanish.

Prefill is compute-bound and decode is memory-bound, so they want different hardware settings, and mixing them on the same GPUs makes each interfere with the other. Disaggregation runs prefill on one pool of GPUs and decode on another, and sends the finished KV cache between them. The DistServe paper by Zhong and colleagues reports that this serves 7.4 times more requests, or meets a 12.6 times tighter latency target, while staying within the targets for over 90% of requests, compared with the baselines it tested.

• -Benefit: each pool is sized and tuned for its own phase, and a long prefill can no longer stall decoding.
• -Cost: the KV cache has to move between pools over a fast network, there are more moving parts, and the two pools must be balanced for your traffic mix.
• -Best fit: large deployments with strict targets for both time to first token and time per token.

State your latency goals in two numbers: a target for time to first token, and a target for time per output token. A useful system-level measure is goodput, the highest request rate that still meets both targets. Optimizing raw tokens per second can raise throughput while breaking those targets, which is why the techniques in this module are judged against goodput.