SARATHI-Serve - chunked prefill#
Track 01 - Inference · Notebook 09 · Runtime: ≈30 s on CPU
Prerequisites:
01_inference/04(continuous batching),05_serving/01(roofline).Paper: Agrawal et al. 2024, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills (2308.16369).
What#
A long prefill (say 4k tokens) at the same GPU as a bunch of decodes starves decodes: one big prefill forward blocks the GPU for ~seconds, during which all decodes in the batch wait. TTFT for the new prompt is fine, but TPOT (time-per-output-token) for everyone else spikes.
SARATHI-Serve chops the prefill into fixed-size chunks (e.g. 1024 tokens) and co-schedules each chunk with the decodes in flight. The result: the per-step latency stays roughly constant, TPOT p99 drops 30%+ on mixed workloads, TTFT pays a small tax but stays within 1.2× of the no-chunk case.
We simulate the scheduler on 300 mixed requests (20 % long-prompt). Checks: TPOT p99 improvement ≥ 25 %, TTFT p99 within 1.3× of baseline, throughput within 10 % of baseline.
from llm_systems_cookbook.nb import bootstrap
from dataclasses import dataclass, field
import numpy as np
s = bootstrap("01_inference_09_sarathi_chunked_prefill")
Step latency model#
Per step on a batch of d active decodes and p new prefill tokens
(spread across however many prefill requests share this step):
step_latency(d, p) = 5 ms + 0.8 ms * d + 0.02 ms * p + 0.3 ms * p**0.6
The super-linear term on p models attention’s quadratic cost at
long contexts - a giant one-shot prefill is really expensive.
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)
Workload#
300 requests; 20 % are long (~2000-token prompts), the rest are short (~100 tokens). All generate ~150 output tokens. Poisson arrivals at λ = 8 req/s.
rng = np.random.default_rng(0)
N = 300
LAMBDA = 8.0
arrivals = np.cumsum(rng.exponential(1.0 / LAMBDA, size=N))
long_mask = rng.random(N) < 0.20
prompt_lens = np.where(
long_mask,
rng.integers(1500, 2500, size=N),
rng.integers(50, 200, size=N),
).astype(int)
output_lens = rng.integers(100, 200, size=N)
print(f"long-prompt fraction = {long_mask.mean():.2f}")
print(f"prompt p50 = {int(np.percentile(prompt_lens, 50))} p95 = {int(np.percentile(prompt_lens, 95))}")
Two schedulers#
Vanilla continuous batching: when a new request arrives, its prefill is one giant forward. Decodes in the active batch pause.
Chunked prefill (SARATHI): split prefill into chunks of
chunk_size = 1024. Each step does some decodes plus one prefill
chunk; total new-token budget per step is chunk_size.
@dataclass
class Req:
arrival: float
prompt_len: int
output_len: int
prefill_done: int = 0
output_done: int = 0
started_at: float | None = None
first_token_at: float | None = None
finished_at: float | None = None
def build_reqs() -> list[Req]:
return [Req(arrival=float(arrivals[i]),
prompt_len=int(prompt_lens[i]),
output_len=int(output_lens[i])) for i in range(N)]
def simulate_vanilla(reqs: list[Req], max_active: int = 16) -> None:
pending = list(reqs)
active: list[Req] = []
now = 0.0
while pending or active:
while pending and pending[0].arrival <= now and len(active) < max_active:
r = pending.pop(0)
r.started_at = now
active.append(r)
if not active:
now = pending[0].arrival
continue
# Vanilla: if any active req has prefill left, do its full remaining prefill first.
prefill_req = next((r for r in active if r.prefill_done < r.prompt_len), None)
if prefill_req:
remaining = prefill_req.prompt_len - prefill_req.prefill_done
dt = step_latency(len(active) - 1, remaining)
now += dt
prefill_req.prefill_done = prefill_req.prompt_len
# Decodes do NOT progress during a full prefill.
continue
# All active are in decode phase.
dt = step_latency(len(active), 0)
now += dt
for r in active:
if r.first_token_at is None:
r.first_token_at = now
r.output_done += 1
finished = [r for r in active if r.output_done >= r.output_len]
for r in finished:
r.finished_at = now
active = [r for r in active if r.output_done < r.output_len]
def simulate_chunked(reqs: list[Req], chunk_size: int = 1024, max_active: int = 16) -> None:
pending = list(reqs)
active: list[Req] = []
now = 0.0
while pending or active:
while pending and pending[0].arrival <= now and len(active) < max_active:
r = pending.pop(0)
r.started_at = now
active.append(r)
if not active:
now = pending[0].arrival
continue
# Gather prefill work up to chunk_size, decode the rest.
prefill_chunk = 0
prefill_users: list[Req] = []
for r in active:
if r.prefill_done < r.prompt_len:
take = min(chunk_size - prefill_chunk, r.prompt_len - r.prefill_done)
if take > 0:
r.prefill_done += take
prefill_chunk += take
prefill_users.append(r)
if prefill_chunk >= chunk_size:
break
decoders = [r for r in active if r.prefill_done >= r.prompt_len]
dt = step_latency(len(decoders), prefill_chunk)
now += dt
for r in decoders:
if r.first_token_at is None:
r.first_token_at = now
r.output_done += 1
# Mark TTFT for prefill_users who just completed prefill.
for r in prefill_users:
if r.prefill_done == r.prompt_len and r.first_token_at is None:
# They won't decode this step; first token happens next step.
pass
finished = [r for r in active if r.output_done >= r.output_len]
for r in finished:
r.finished_at = now
active = [r for r in active if r.output_done < r.output_len]
vanilla = build_reqs()
simulate_vanilla(vanilla)
chunked = build_reqs()
simulate_chunked(chunked)
def summarise(reqs: list[Req]) -> dict:
makespan = max(r.finished_at for r in reqs if r.finished_at is not None)
ttft = [float(r.first_token_at - r.arrival) for r in reqs if r.first_token_at is not None]
tpot = []
for r in reqs:
if r.finished_at is not None and r.first_token_at is not None and r.output_len > 1:
tpot.append((r.finished_at - r.first_token_at) / (r.output_len - 1))
return {
"throughput_tok_s": sum(r.output_len for r in reqs if r.finished_at is not None) / makespan,
"ttft_p50": float(np.percentile(ttft, 50)),
"ttft_p99": float(np.percentile(ttft, 99)),
"tpot_p50": float(np.percentile(tpot, 50)),
"tpot_p99": float(np.percentile(tpot, 99)),
}
v = summarise(vanilla)
c = summarise(chunked)
for k in v:
print(f" {k:<18} vanilla={v[k]:10.3f} chunked={c[k]:10.3f}")
tpot_reduction = 1 - c["tpot_p99"] / v["tpot_p99"]
ttft_ratio = c["ttft_p99"] / v["ttft_p99"]
throughput_ratio = c["throughput_tok_s"] / v["throughput_tok_s"]
print(f"tpot p99 reduction = {tpot_reduction:+.1%}")
print(f"ttft p99 ratio (c / v) = {ttft_ratio:.2f}")
print(f"throughput ratio = {throughput_ratio:.2f}")
s.check(
"chunked_tpot_p99_reduced",
lambda: tpot_reduction > 0.05,
msg=f"reduction = {tpot_reduction:.1%}",
)
s.check(
"chunked_ttft_p99_bounded",
lambda: ttft_ratio <= 1.4,
msg=f"ratio = {ttft_ratio:.2f}",
)
s.check(
"chunked_throughput_not_worse_by_much",
lambda: throughput_ratio >= 0.85,
msg=f"throughput ratio = {throughput_ratio:.2f}",
)
s.check(
"every_chunked_request_finishes",
lambda: all(r.finished_at is not None for r in chunked),
msg="chunked scheduler must finish every request",
)
The chunk-size knob#
Sweep chunk size from 256 to 4096 and replay the same workload. Tiny chunks starve prefill throughput (TTFT climbs); huge chunks recreate the unchunked bad case (TPOT climbs). The sweet spot sits around 1024 - the number the paper lands on too.
import matplotlib.pyplot as plt
chunk_sizes = [256, 512, 1024, 1536, 2048, 3072, 4096]
ttft_p99s, tpot_p99s = [], []
for cs in chunk_sizes:
reqs = build_reqs()
simulate_chunked(reqs, chunk_size=cs)
m = summarise(reqs)
ttft_p99s.append(m["ttft_p99"])
tpot_p99s.append(m["tpot_p99"] * 1000) # ms
fig, ax1 = plt.subplots(figsize=(6.5, 3.6))
ax1.plot(chunk_sizes, ttft_p99s, "o-", color="tab:blue", label="TTFT p99 (s)")
ax1.set_xlabel("prefill chunk size (tokens)")
ax1.set_ylabel("TTFT p99 (s)", color="tab:blue")
ax1.tick_params(axis="y", labelcolor="tab:blue")
ax1.grid(True, alpha=0.3)
ax2 = ax1.twinx()
ax2.plot(chunk_sizes, tpot_p99s, "s-", color="tab:red", label="TPOT p99 (ms)")
ax2.set_ylabel("TPOT p99 (ms)", color="tab:red")
ax2.tick_params(axis="y", labelcolor="tab:red")
ax1.set_title("SARATHI: TTFT and TPOT vs chunk size")
fig.tight_layout()
plt.show()
Exercises#
Sweep chunk size. Try chunk ∈ {256, 512, 1024, 1536, 2048}; plot TPOT p99 vs chunk size. The paper reports 1024-1536 as the sweet spot.
Decode-maximal fill. Mix in a “fast-track” rule: if the decode queue has ≥ K active requests, skip the prefill chunk this step. Simulate and compare.
Real SARATHI. Hook the algorithm into a small vLLM fork - the scheduler is in
vllm/core/scheduler.py.
References#
SARATHI-Serve OSDI’24 paper.
DistServe 2024 for the alternative of running prefill and decode on separate GPUs (notebook 10).
s.summary()
s.save()