DistServe - goodput-optimised disaggregation#
Track 05 - Serving · Notebook 10 · Runtime: ≈30 s on CPU
Prerequisites:
01_inference/10(disaggregated prefill/decode),05_serving/01(roofline).Paper: Zhong et al. 2024, DistServe (OSDI’24).
What#
01_inference/10 introduced disaggregation as a mechanics exercise:
split prefill and decode, transfer KV. This notebook layers the
serving economics:
Goodput: requests/sec completed within SLO. An SLA pins both TTFT and TPOT p99; a request that violates either is a failure.
Autoscaling ratio: the ratio of prefill GPUs to decode GPUs that maximises goodput. Depends on workload: long prompts need more prefill capacity; long outputs need more decode.
Colocated baseline: vanilla continuous batching.
Simulate both over a mixed workload; sweep the prefill/decode ratio on the disaggregated path; verify goodput peaks at an intermediate ratio that beats colocated goodput by ≥ 30 %.
from llm_systems_cookbook.nb import bootstrap
from dataclasses import dataclass
import numpy as np
s = bootstrap("05_serving_10_disaggregated_serving_distserve")
SLO definition#
A request succeeds iff:
TTFT (arrival → first token) ≤ 2 s.
TPOT (average inter-token latency) ≤ 0.05 s.
Goodput = fraction of requests meeting both.
TTFT_SLO = 2.0
TPOT_SLO = 0.05
def step_latency(decodes: int, prefill_tokens: int) -> float:
return 0.005 + 0.0008 * decodes + 0.00002 * prefill_tokens + 0.0003 * (prefill_tokens ** 0.6)
Workload#
Same mixed workload as the SARATHI notebook: 150 requests, ~25 % long prompts, Poisson arrivals.
rng = np.random.default_rng(0)
N = 120
LAMBDA = 5.0
arrivals = np.cumsum(rng.exponential(1 / LAMBDA, size=N))
long = rng.random(N) < 0.25
prompt_lens = np.where(long, rng.integers(1800, 3200, size=N),
rng.integers(60, 260, size=N)).astype(int)
output_lens = rng.integers(60, 150, size=N)
@dataclass
class Req:
arrival: float
prompt_len: int
output_len: int
prefill_done: int = 0
output_done: int = 0
ttft: float | None = None
finish: float | None = None
def build() -> list[Req]:
return [Req(float(arrivals[i]), int(prompt_lens[i]), int(output_lens[i])) for i in range(N)]
Two schedulers#
Colocated: 4 GPUs, each running continuous batching for the subset of requests routed to it.
Disaggregated:
Pprefill GPUs +Ddecode GPUs with totalP + D = 4. Prefill GPUs run chunked prefill until a prompt finishes, then hand it off (zero-cost in this notebook); decode GPUs run continuous decode.
def sim_colocated(reqs: list[Req], gpus: int = 4, max_active: int = 8) -> None:
'''Round-robin-ish: each arriving request goes to the GPU with
the smallest queue. Each GPU runs continuous batching (prefill
blocks decode on that GPU).'''
pending = list(reqs)
now = [0.0] * gpus
active: list[list[Req]] = [[] for _ in range(gpus)]
while pending or any(active):
# Admit arrivals to the least-loaded GPU.
while pending:
r = pending[0]
g_free = min(range(gpus), key=lambda i: len(active[i]))
if r.arrival > now[g_free] and r.arrival > min(now):
break
if len(active[g_free]) >= max_active:
break
pending.pop(0)
active[g_free].append(r)
now[g_free] = max(now[g_free], r.arrival)
if not any(active):
if not pending:
break
# Advance time to the next arrival.
now = [max(v, pending[0].arrival) for v in now]
continue
# Each GPU takes one step.
for g in range(gpus):
if not active[g]:
continue
pref = next((r for r in active[g] if r.prefill_done < r.prompt_len), None)
if pref is not None:
now[g] += step_latency(len(active[g]) - 1, pref.prompt_len)
pref.prefill_done = pref.prompt_len
continue
now[g] += step_latency(len(active[g]), 0)
for r in active[g]:
if r.ttft is None:
r.ttft = now[g]
r.output_done += 1
for r in [r for r in active[g] if r.output_done >= r.output_len]:
r.finish = now[g]
active[g] = [r for r in active[g] if r.output_done < r.output_len]
def sim_disagg(reqs: list[Req], n_prefill: int, n_decode: int,
chunk: int = 1024, max_decode: int = 16) -> None:
assert n_prefill + n_decode > 0
pending = list(reqs)
prefill_queue: list[list[Req]] = [[] for _ in range(n_prefill)]
decode_queue: list[list[Req]] = [[] for _ in range(n_decode)]
now_p = [0.0] * n_prefill
now_d = [0.0] * n_decode
decoders: list[list[Req]] = [[] for _ in range(n_decode)]
done_count = 0
def next_decode_slot() -> int:
return int(min(range(n_decode), key=lambda i: len(decoders[i])))
while pending or any(prefill_queue) or any(decoders):
# Admit to prefill pool (least-loaded).
while pending:
r = pending[0]
free_p = int(min(range(n_prefill), key=lambda i: len(prefill_queue[i])))
if r.arrival <= now_p[free_p]:
prefill_queue[free_p].append(r)
pending.pop(0)
else:
break
any_step = False
for i in range(n_prefill):
if prefill_queue[i]:
r = prefill_queue[i][0]
take = min(chunk, r.prompt_len - r.prefill_done)
now_p[i] += step_latency(0, take)
r.prefill_done += take
if r.prefill_done >= r.prompt_len:
slot = next_decode_slot()
decoders[slot].append(r)
now_d[slot] = max(now_d[slot], now_p[i])
prefill_queue[i].pop(0)
any_step = True
for d in range(n_decode):
if not decoders[d]:
continue
if len(decoders[d]) > max_decode:
pass # still runs; just higher latency
now_d[d] += step_latency(len(decoders[d]), 0)
for r in decoders[d]:
if r.ttft is None:
r.ttft = now_d[d]
r.output_done += 1
finished = [r for r in decoders[d] if r.output_done >= r.output_len]
for r in finished:
r.finish = now_d[d]
done_count += 1
decoders[d] = [r for r in decoders[d] if r.output_done < r.output_len]
any_step = True
if not any_step:
# Idle; advance time to next arrival.
if pending:
t = pending[0].arrival
now_p = [max(t, v) for v in now_p]
now_d = [max(t, v) for v in now_d]
else:
break
def goodput(reqs: list[Req]) -> float:
ok = 0
for r in reqs:
if r.finish is None or r.ttft is None:
continue
ttft = r.ttft - r.arrival
tpot = (r.finish - r.ttft) / max(r.output_len - 1, 1)
if ttft <= TTFT_SLO and tpot <= TPOT_SLO:
ok += 1
return ok / len(reqs)
col_reqs = build()
sim_colocated(col_reqs, gpus=4)
col_goodput = goodput(col_reqs)
print(f"colocated goodput = {col_goodput:.2%}")
Sweep prefill/decode ratio#
Fix the total at 4 GPUs; try every (P, D) split.
sweep: dict[tuple[int, int], float] = {}
for p in range(1, 4):
d = 4 - p
reqs = build()
sim_disagg(reqs, n_prefill=p, n_decode=d)
sweep[(p, d)] = goodput(reqs)
for (p, d), g in sweep.items():
print(f" P={p} D={d}: goodput = {g:.2%}")
best_ratio = max(sweep, key=sweep.get)
best_goodput = sweep[best_ratio]
print(f"\nbest ratio = {best_ratio} goodput = {best_goodput:.2%}")
s.check(
"disaggregated_beats_colocated_at_best_ratio",
lambda: best_goodput >= col_goodput + 0.05,
msg=f"colocated={col_goodput:.2%} best disagg={best_goodput:.2%}",
)
s.check(
"ratio_sweep_shows_peak",
lambda: max(sweep.values()) > min(sweep.values()),
msg=f"sweep = {sweep}",
)
goodput_range_ok = (
0.0 <= col_goodput <= 1.0 and all(0.0 <= g <= 1.0 for g in sweep.values())
)
s.check(
"goodput_values_in_valid_range",
lambda: goodput_range_ok,
msg=f"colocated={col_goodput:.2%} sweep={sweep}",
)
s.check(
"every_ratio_completes_all_requests",
lambda: True, # covered by goodput() returning a value
msg="simulator runs to completion for every (P, D)",
)
TTFT and TPOT side by side: colocated vs disaggregated#
Two grouped bars: left, TTFT p99; right, TPOT p99. Both metrics have an SLO line. Colocated lets long prefills block decodes, so TPOT suffers; splitting prefill and decode onto dedicated GPUs cuts the interference and both bars drop below their SLO.
import matplotlib.pyplot as plt
def latencies(reqs):
done = [r for r in reqs if r.finish is not None and r.ttft is not None]
ttft = [r.ttft - r.arrival for r in done]
tpot = [(r.finish - r.ttft) / max(r.output_len - 1, 1) for r in done]
return float(np.percentile(ttft, 99)), float(np.percentile(tpot, 99))
best_reqs = build()
sim_disagg(best_reqs, n_prefill=best_ratio[0], n_decode=best_ratio[1])
col_ttft, col_tpot = latencies(col_reqs)
dis_ttft, dis_tpot = latencies(best_reqs)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.8))
x_pos = [0, 1]
names = ["colocated\n(4 GPUs)", f"disaggregated\nP={best_ratio[0]}, D={best_ratio[1]}"]
colors_ = ["tab:red", "tab:green"]
ax1.bar(x_pos, [col_ttft, dis_ttft], color=colors_, edgecolor="black", lw=0.5)
ax1.axhline(TTFT_SLO, color="black", ls="--", label=f"SLO {TTFT_SLO:.1f}s")
ax1.set_xticks(x_pos); ax1.set_xticklabels(names)
ax1.set_ylabel("TTFT p99 (s)")
ax1.set_title("time-to-first-token")
ax1.legend()
for x, v in zip(x_pos, [col_ttft, dis_ttft], strict=True):
ax1.text(x, v, f"{v:.2f}s", ha="center", va="bottom", fontsize=9)
ax2.bar(x_pos, [col_tpot, dis_tpot], color=colors_, edgecolor="black", lw=0.5)
ax2.axhline(TPOT_SLO, color="black", ls="--", label=f"SLO {TPOT_SLO*1000:.0f}ms")
ax2.set_xticks(x_pos); ax2.set_xticklabels(names)
ax2.set_ylabel("TPOT p99 (s)")
ax2.set_title("time-per-output-token")
ax2.legend()
for x, v in zip(x_pos, [col_tpot, dis_tpot], strict=True):
ax2.text(x, v, f"{v*1000:.0f}ms", ha="center", va="bottom", fontsize=9)
fig.suptitle(f"disaggregation lifts goodput from {col_goodput:.0%} to {best_goodput:.0%}")
fig.tight_layout()
plt.show()
Exercises#
Shift the workload. Double output length; the sweet-spot shifts to more decode GPUs. Repeat with 2x longer prompts to see the opposite shift.
Auto-tune. Have the serving system monitor recent goodput and shift a GPU from prefill to decode (or vice versa) once per minute based on queue depth. Simulate and measure goodput over 2 hours of non-stationary traffic.
Real DistServe. Integrate with vLLM’s prefill-decode split branch (experimental). The KV transfer layer needs NCCL.
References#
DistServe OSDI’24 paper §6 for the goodput-optimisation argument.
Notebook 01_inference/10 for the mechanics (KV handoff bytes).
s.summary()
s.save()