Disaggregated prefill and decode
Compare kernel, parallelism and disaggregated-serving choices using memory, communication and latency constraints.
THE IDEA
Splitting prefill from decode
Prefill is compute-heavy and bursty; decode is memory-bandwidth-bound and steady. Disaggregated serving runs them on separate machines, each tuned for its job, so a big prompt does not stall requests that are decoding. The price: after prefill the request's KV state (512 MiB in the reference example) must be sent to the decode machine, and that transfer adds to the new request's time to first token. Over 32 GiB/s it is 15.6 ms; over 2 GiB/s it is 250 ms.
| prefill + decode on one machine | prefill machine → KV transfer → decode machine | |
|---|---|---|
| new request waits | prefill only | prefill + 15.63 ms transfer |
| running decodes | stalled about 200 ms by the prefill | not stalled |
Separating phases lets each machine specialise, but the KV state has to cross a network, and that needs coordination and bandwidth. It is not a win for free on a slow link.
Predict first. Is separating phases always beneficial regardless of network speed?
Chapter context
Scaling serving changes how work is placed and moved. An operation can be mathematically unchanged while its memory traffic changes dramatically. Device count alone does not describe the performance of the resulting system.
TRY IT YOURSELF
Part 1. Do it hands-on
Compute transfer times on two links, then write a model that decides whether splitting pays off.
| prefill + decode on one machine | prefill machine → KV transfer → decode machine | |
|---|---|---|
| new request waits | prefill only | prefill + 15.63 ms transfer |
| running decodes | stalled about 200 ms by the prefill | not stalled |
Separating phases lets each machine specialise, but the KV state has to cross a network, and that needs coordination and bandwidth. It is not a win for free on a slow link.
- To do: After prefill, 512 MiB of KV state must move from the prefill machine to the decode machine over a 32 GiB/s link. Milliseconds (3 decimals)? (512 MiB / 32 GiB = 1/64 s)
- To do: Over a slow 2 GiB/s link, the same transfer takes (ms, 1 decimal)?
- To do: Splitting phases avoids a 200 ms stall for running decodes but adds the transfer to the new request's time to first token. Over the slow 2 GiB/s link the transfer is 250 ms. Is the split a clear win for the NEW request's latency? Type yes or no
Part 2. Do it in code
- kv_transfer_ms(kv_bytes, bandwidth_bytes_per_s): kv_bytes / bandwidth · 1000. ValueError for a non-positive bandwidth.
- ttft_ms(queue_ms, prefill_ms, transfer_ms, first_decode_ms): the sum of the four parts (transfer_ms is 0 when prefill and decode run on the same machine).
- stall_avoided_ms(decodes, prompt_tokens, decode_cost=1.0, prefill_cost=0.1): the extra time running decodes would otherwise wait if the prompt were prefilled on their machine: prompt_tokens · prefill_cost (milliseconds).
- split_worthwhile(transfer_ms, stall_avoided, max_extra_ttft_ms): True if the transfer adds at most max_extra_ttft_ms to the new request AND stall_avoided is larger than transfer_ms.
- break_even_bandwidth(kv_bytes, max_transfer_ms): the smallest bandwidth (bytes/s) for which the transfer takes at most max_transfer_ms.
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Is separating phases always beneficial regardless of network speed?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
A weight-traffic lower-bound estimate is
This assumes the stated bytes really traverse the bottleneck bandwidth. KV transfer adds roughly bytes_to_move/network_bandwidth before protocol and synchronization overhead. Batching and caching can change the bytes per produced token.
Worked example
A fictional 14-GB dense weight scan across a sustained 700-GB/s memory path takes at least 14/700=0.02 seconds, giving a 50-step/s bandwidth ceiling under the one-scan-per-step assumption. Real execution includes other memory traffic and work. A 512-MiB KV transfer over 32 GiB/s has a payload-only lower bound of 1/64 second.
Reference code
The following example uses the Python standard library and does not call a remote service.
from math import isclose
assert isclose(14/700,.02)
assert isclose(1/(14/700),50)
transfer=(512*1024**2)/(32*1024**3)
assert isclose(transfer,1/64)
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Peak device bandwidth is not guaranteed sustained application bandwidth. Exact attention can be computed with different memory layouts. Additional GPUs can add communication bottlenecks.