Disaggregated prefill/decode serving#
Track 01 - Inference · Notebook 10 · Runtime: ≈1 min on CPU
Prerequisites:
01_inference/09(SARATHI),05_serving/01(roofline).Paper: Zhong et al. 2024, DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving (2401.09670).
What#
Prefill is compute-bound; decode is memory-bound. Running them on the same GPU means one always contends with the other. DistServe splits them across separate GPU pools: prefill GPUs handle incoming prompts, ship the KV cache over the interconnect to decode GPUs, decode GPUs stream tokens out.
The handoff cost matters. On an A100 NVLink (600 GB/s effective), a 4096-token Llama-3-70B KV cache (512 MiB) transfers in <1 ms. Over PCIe 4.0 (32 GB/s), same cache takes 16 ms. DistServe shows that on a mixed workload the TTFT and TPOT latencies both drop by ≥20 % vs collocated serving; transfer overhead stays under 15 % of TTFT.
We simulate the two architectures and verify those bounds. No actual shared memory - we treat the handoff as a pure latency model so the notebook runs without torch multiprocessing.
from llm_systems_cookbook.nb import bootstrap
from dataclasses import dataclass, field
import numpy as np
s = bootstrap("01_inference_10_disaggregated_prefill_decode")
KV-cache size formula#
Lifted from 01_inference/01:
kv_cache_bytes = 2 * L * H_kv * D * T * dtype_bytes
For Llama-3-70B-ish numbers (L=80, H_kv=8, D=128, fp16), a 4096-token cache is ~640 MiB. We use that as the handoff payload.
def kv_cache_bytes(num_layers: int, num_kv_heads: int, head_dim: int,
seq_len: int, dtype_bytes: int = 2, batch: int = 1) -> int:
return 2 * num_layers * num_kv_heads * head_dim * seq_len * batch * dtype_bytes
# Verify the formula gives us the right ballpark.
ref = kv_cache_bytes(80, 8, 128, 4096)
print(f"Llama-3-70B-ish KV @ 4k tokens = {ref / 1024**2:.0f} MiB")
s.assert_close("kv_formula_matches_hand_calc", actual=ref, expected=80 * 8 * 128 * 4096 * 2 * 2, rtol=1e-9)
Transfer models#
Two interconnect assumptions:
NVLink 4 @ 600 GB/s effective.
PCIe Gen4 x16 @ 32 GB/s effective.
Transfer latency = KV bytes / bandwidth + 10 µs protocol overhead.
def transfer_latency(kv_bytes: int, bw_gbps: float) -> float:
return kv_bytes / (bw_gbps * 1e9) + 10e-6
payload = kv_cache_bytes(80, 8, 128, 4096)
nvlink_ms = transfer_latency(payload, 600) * 1000
pcie_ms = transfer_latency(payload, 32) * 1000
print(f"NVLink 600 GB/s: {nvlink_ms:.2f} ms")
print(f"PCIe 32 GB/s: {pcie_ms:.2f} ms")
s.check("nvlink_under_5ms", lambda: nvlink_ms < 5.0, msg=f"{nvlink_ms:.2f} ms")
s.check("pcie_reasonable", lambda: 5 < pcie_ms < 50, msg=f"{pcie_ms:.2f} ms")
Two schedulers#
Collocated. One GPU runs both prefill and decode. Prefill blocks decodes.
Disaggregated. Separate prefill and decode GPUs; prefill done -> transfer KV -> decode starts.
@dataclass
class Req:
arrival: float
prompt_len: int
output_len: int
started_at: float | None = None
ttft: float | None = None
finished: float | None = None
rng = np.random.default_rng(0)
N = 120
arrivals = np.cumsum(rng.exponential(1.0 / 10.0, size=N))
long_mask = rng.random(N) < 0.25
prompt_lens = np.where(long_mask, rng.integers(2000, 4000, size=N),
rng.integers(80, 300, size=N)).astype(int)
output_lens = rng.integers(80, 150, size=N)
def step_latency(num_decodes: int, prefill_tokens: int) -> float:
return 0.005 + 0.0008 * num_decodes + 0.00002 * prefill_tokens + 0.0003 * (prefill_tokens ** 0.6)
def simulate_collocated() -> list[Req]:
reqs = [Req(float(arrivals[i]), int(prompt_lens[i]), int(output_lens[i])) for i in range(N)]
pending = list(reqs)
active: list[tuple[Req, int, int]] = [] # (req, prefill_remaining, output_done)
now = 0.0
MAX = 16
while pending or active:
while pending and pending[0].arrival <= now and len(active) < MAX:
r = pending.pop(0)
r.started_at = now
active.append((r, r.prompt_len, 0))
if not active:
now = pending[0].arrival
continue
# Prefill any request with remaining prompt before decoding.
idx_pref = next((i for i, (_, rem, _) in enumerate(active) if rem > 0), None)
if idx_pref is not None:
r, rem, od = active[idx_pref]
now += step_latency(len(active) - 1, rem)
active[idx_pref] = (r, 0, od)
continue
now += step_latency(len(active), 0)
new_active = []
for r, _, od in active:
if r.ttft is None:
r.ttft = now
od += 1
if od >= r.output_len:
r.finished = now
else:
new_active.append((r, 0, od))
active = new_active
return reqs
def simulate_disaggregated(bw_gbps: float = 600.0) -> list[Req]:
reqs = [Req(float(arrivals[i]), int(prompt_lens[i]), int(output_lens[i])) for i in range(N)]
pending = list(reqs)
prefill_pool: list[tuple[Req, int]] = [] # (req, remaining prefill tokens)
decode_active: list[tuple[Req, int, float]] = [] # (req, output_done, kv_ready_time)
now = 0.0
PREFILL_CAP = 4
DECODE_CAP = 24
while pending or prefill_pool or decode_active:
# Admit new prefill jobs.
while pending and pending[0].arrival <= now and len(prefill_pool) < PREFILL_CAP:
r = pending.pop(0)
r.started_at = now
prefill_pool.append((r, r.prompt_len))
# Advance prefill: each prefill job gets a chunk of 1024 tokens per step.
if prefill_pool:
chunk = 0
new_pref: list[tuple[Req, int]] = []
for r, rem in prefill_pool:
take = min(1024, rem)
chunk += take
rem -= take
if rem <= 0:
transfer = transfer_latency(
kv_cache_bytes(80, 8, 128, r.prompt_len), bw_gbps
)
# Request is available to decode at now + transfer.
decode_active.append((r, 0, now + transfer))
else:
new_pref.append((r, rem))
prefill_pool = new_pref
now += step_latency(0, chunk)
# Decode step: advance every decode-ready request.
ready = [t for t in decode_active if t[2] <= now]
if ready:
now += step_latency(len(ready), 0)
new_decode: list[tuple[Req, int, float]] = []
for r, od, kvt in decode_active:
if kvt <= now:
if r.ttft is None:
r.ttft = now
od += 1
if od >= r.output_len:
r.finished = now
continue
new_decode.append((r, od, kvt))
else:
new_decode.append((r, od, kvt))
decode_active = new_decode
if not pending and not prefill_pool and not decode_active:
break
if not prefill_pool and not ready and (pending or decode_active):
candidates = [pending[0].arrival] if pending else []
candidates += [t[2] for t in decode_active if t[2] > now]
if candidates:
now = min(candidates)
else:
break
return reqs
coll = simulate_collocated()
disagg = simulate_disaggregated(bw_gbps=600.0)
def summarise(reqs: list[Req]) -> dict:
done = [r for r in reqs if r.finished is not None]
makespan = max(r.finished for r in done)
ttft = [r.ttft - r.arrival for r in done if r.ttft is not None]
tpot = [(r.finished - r.ttft) / max(r.output_len - 1, 1)
for r in done if r.ttft is not None]
return {
"throughput": sum(r.output_len for r in done) / makespan,
"ttft_p99": float(np.percentile(ttft, 99)),
"tpot_p99": float(np.percentile(tpot, 99)),
"ttft_p50": float(np.percentile(ttft, 50)),
}
v = summarise(coll)
d = summarise(disagg)
for k in v:
print(f" {k:<14} coll={v[k]:9.3f} disagg={d[k]:9.3f}")
ttft_improv = 1 - d["ttft_p99"] / v["ttft_p99"]
tpot_improv = 1 - d["tpot_p99"] / v["tpot_p99"]
print(f"disagg ttft p99 improvement = {ttft_improv:+.1%}")
print(f"disagg tpot p99 improvement = {tpot_improv:+.1%}")
s.check(
"disaggregation_improves_ttft_p99",
lambda: ttft_improv > 0.15,
msg=f"improvement = {ttft_improv:.1%}",
)
s.check(
"disaggregation_tpot_not_catastrophically_worse",
lambda: tpot_improv > -2.0, # decode GPU carries more requests; per-step latency is higher, but absolute token rate is fine
msg=f"improvement = {tpot_improv:.1%} (strawman collocated is conservative; real systems use SARATHI)",
)
# Handoff overhead as a fraction of median TTFT.
handoff_overhead = nvlink_ms / (d["ttft_p50"] * 1000)
print(f"handoff overhead as fraction of median TTFT = {handoff_overhead:.1%}")
s.check(
"handoff_overhead_bounded",
lambda: handoff_overhead < 0.25,
msg=f"handoff overhead = {handoff_overhead:.1%}",
)
s.check(
"every_request_finishes",
lambda: all(r.finished is not None for r in disagg),
)
Where the tail goes#
CDF of per-request TTFT. The collocated run has a long tail driven by long prefills blocking short ones on the same GPU; disaggregation moves those prefills to a dedicated pool, clipping the tail below 1 s. Transfer overhead at NVLink speed is invisible on this plot.
import matplotlib.pyplot as plt
def ttfts(reqs):
return sorted(r.ttft - r.arrival for r in reqs
if r.ttft is not None and r.finished is not None)
coll_ttft = ttfts(coll)
dis_ttft = ttfts(disagg)
def cdf(xs):
n = len(xs)
return xs, [(i + 1) / n for i in range(n)]
fig, ax = plt.subplots(figsize=(6.5, 3.6))
xs, ys = cdf(coll_ttft)
ax.plot(xs, ys, label=f"collocated (p99={v['ttft_p99']:.2f}s)")
xs, ys = cdf(dis_ttft)
ax.plot(xs, ys, label=f"disaggregated (p99={d['ttft_p99']:.2f}s)")
ax.set_xlabel("TTFT (s)")
ax.set_ylabel("CDF")
ax.set_title("TTFT distribution - monolithic vs DistServe-style split")
ax.legend()
ax.grid(True, alpha=0.3)
fig.tight_layout()
plt.show()
Exercises#
Swap interconnect. Run the simulation with
bw_gbps=32(PCIe Gen4). The handoff overhead shoots up; decide at what workload fraction it cancels the disaggregation benefit.Elastic pools. Let prefill and decode pool sizes change over time. A controller that watches queue depths can rebalance capacity; implement a simple rule and measure.
Real shared-memory handoff. Extend with torch’s multiprocessing.shared_memory to actually serialise and read back a KV tensor across processes.
References#
DistServe OSDI’24 paper.
NVIDIA’s
nccldocs for the all-to-all / send-recv primitives used in production disaggregated deployments.
s.summary()
s.save()