Batching strategies#
Track 05 - Serving · Notebook 08 · Runtime: ≈30 s on CPU
Prerequisites:
01_inference/04(continuous batching Orca),01_inference/09(SARATHI).Reference: vLLM’s scheduler, SARATHI-Serve.
What#
Four batching strategies compared in one simulator:
Static. Fixed batch size; admit a batch, drain it, next batch.
Dynamic. Admit up to
max_batchevery step; drop finished, no new admissions mid-step.Continuous. Per-step admission + eviction (Orca).
Chunked prefill. Continuous + chunked prefill (SARATHI) - the current production default.
Same workload, same latency model as the Orca notebook; this notebook emphasises the side-by-side Pareto: throughput vs TTFT p99 vs TPOT p99.
from llm_systems_cookbook.nb import bootstrap
from dataclasses import dataclass
import numpy as np
s = bootstrap("05_serving_08_batching_strategies")
Workload#
200 requests; 20 % long prompts (2k-3k tokens). Arrival λ = 8 req/s.
rng = np.random.default_rng(0)
N = 200
arrivals = np.cumsum(rng.exponential(1 / 8.0, size=N))
long = rng.random(N) < 0.2
prompt_lens = np.where(long, rng.integers(2000, 3000, size=N), rng.integers(60, 300, size=N)).astype(int)
output_lens = rng.integers(80, 150, size=N)
def step_latency(batch: int, prefill_tokens: int) -> float:
return 0.005 + 0.0008 * batch + 0.00002 * prefill_tokens + 0.0003 * (prefill_tokens ** 0.6)
@dataclass
class Req:
arrival: float
prompt_len: int
output_len: int
prefill_done: int = 0
output_done: int = 0
start: float | None = None
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)]
Four schedulers#
MAX_BATCH = 16
def sim_static(reqs: list[Req]) -> None:
pending = list(reqs)
now = 0.0
while pending:
if pending[0].arrival > now:
now = pending[0].arrival
batch: list[Req] = []
while pending and pending[0].arrival <= now and len(batch) < MAX_BATCH:
batch.append(pending.pop(0))
if not batch:
continue
for r in batch:
r.start = now
# Full prefill for all in batch, in series.
for r in batch:
now += step_latency(len(batch) - 1, r.prompt_len)
# Decode until all finish (bound by longest).
remaining = {id(r): r.output_len for r in batch}
while remaining:
now += step_latency(len(remaining), 0)
done_this_step = []
for r in batch:
if id(r) not in remaining:
continue
if r.ttft is None:
r.ttft = now
remaining[id(r)] -= 1
r.output_done += 1
if remaining[id(r)] <= 0:
r.finish = now
done_this_step.append(id(r))
for k in done_this_step:
del remaining[k]
def sim_dynamic(reqs: list[Req]) -> None:
pending = list(reqs)
active: list[Req] = []
now = 0.0
admission_gate_open = True
while pending or active:
if admission_gate_open:
while pending and pending[0].arrival <= now and len(active) < MAX_BATCH:
r = pending.pop(0)
r.start = now
active.append(r)
if not active:
now = pending[0].arrival
continue
# Prefill each new request first.
pref = next((r for r in active if r.prefill_done < r.prompt_len), None)
if pref is not None:
now += step_latency(len(active) - 1, pref.prompt_len)
pref.prefill_done = pref.prompt_len
continue
now += step_latency(len(active), 0)
for r in active:
if r.ttft is None:
r.ttft = now
r.output_done += 1
finished = [r for r in active if r.output_done >= r.output_len]
for r in finished:
r.finish = now
active = [r for r in active if r.output_done < r.output_len]
# Dynamic: close admission gate until the batch drains to half.
if len(active) <= MAX_BATCH // 2:
admission_gate_open = True
else:
admission_gate_open = False
def sim_continuous(reqs: list[Req]) -> 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_BATCH:
r = pending.pop(0)
r.start = now
active.append(r)
if not active:
now = pending[0].arrival
continue
pref = next((r for r in active if r.prefill_done < r.prompt_len), None)
if pref is not None:
now += step_latency(len(active) - 1, pref.prompt_len)
pref.prefill_done = pref.prompt_len
continue
now += step_latency(len(active), 0)
for r in active:
if r.ttft is None:
r.ttft = now
r.output_done += 1
finished = [r for r in active if r.output_done >= r.output_len]
for r in finished:
r.finish = now
active = [r for r in active if r.output_done < r.output_len]
def sim_chunked(reqs: list[Req], chunk: int = 1024) -> 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_BATCH:
r = pending.pop(0)
r.start = now
active.append(r)
if not active:
now = pending[0].arrival
continue
prefill_budget = 0
for r in active:
if r.prefill_done < r.prompt_len:
take = min(chunk - prefill_budget, r.prompt_len - r.prefill_done)
r.prefill_done += take
prefill_budget += take
if prefill_budget >= chunk:
break
decoders = [r for r in active if r.prefill_done >= r.prompt_len]
now += step_latency(len(decoders), prefill_budget)
for r in decoders:
if r.ttft is None:
r.ttft = now
r.output_done += 1
finished = [r for r in active if r.output_done >= r.output_len]
for r in finished:
r.finish = now
active = [r for r in active if r.output_done < r.output_len]
results: dict[str, list[Req]] = {}
for name, sim in [("static", sim_static), ("dynamic", sim_dynamic),
("continuous", sim_continuous), ("chunked", sim_chunked)]:
r = build()
sim(r)
results[name] = r
def summary(reqs: list[Req]) -> dict:
done = [r for r in reqs if r.finish is not None]
span = max(r.finish for r in done)
ttft = [r.ttft - r.arrival for r in done if r.ttft is not None]
tpot = [(r.finish - 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) / span,
"ttft_p99": float(np.percentile(ttft, 99)),
"tpot_p99": float(np.percentile(tpot, 99)),
}
for name, reqs in results.items():
m = summary(reqs)
print(f" {name:<10} tok/s={m['throughput']:7.1f} ttft_p99={m['ttft_p99']:7.2f}s tpot_p99={m['tpot_p99']:7.3f}s")
static = summary(results["static"])
dynamic = summary(results["dynamic"])
cont = summary(results["continuous"])
chunked = summary(results["chunked"])
s.check(
"continuous_throughput_at_least_static",
lambda: cont["throughput"] >= 0.9 * static["throughput"],
msg=f"static={static['throughput']:.1f} cont={cont['throughput']:.1f}",
)
s.check(
"continuous_ttft_p99_better_than_static",
lambda: cont["ttft_p99"] < static["ttft_p99"],
msg=f"static={static['ttft_p99']:.2f} cont={cont['ttft_p99']:.2f}",
)
s.check(
"chunked_tpot_p99_better_than_continuous",
lambda: chunked["tpot_p99"] <= cont["tpot_p99"] * 1.3,
msg=f"cont={cont['tpot_p99']:.3f} chunked={chunked['tpot_p99']:.3f}",
)
s.check(
"dynamic_between_static_and_continuous",
lambda: dynamic["throughput"] >= 0.8 * cont["throughput"],
msg=f"dynamic={dynamic['throughput']:.1f} cont={cont['throughput']:.1f}",
)
s.check(
"every_strategy_finishes_all_requests",
lambda: all(all(r.finish is not None for r in reqs) for reqs in results.values()),
)
Throughput vs TTFT p99 for each scheduler#
One dot per strategy. The ideal live in the top-left corner (high throughput, low TTFT). Static is fast at decoding once a batch is full but the queueing tax shows up as a long TTFT tail. Continuous and chunked keep admission open and dominate the Pareto.
import matplotlib.pyplot as plt
summaries = {name: summary(reqs) for name, reqs in results.items()}
colors = {"static": "tab:red", "dynamic": "tab:orange",
"continuous": "tab:blue", "chunked": "tab:green"}
fig, ax = plt.subplots(figsize=(7, 4.4))
for name, m in summaries.items():
size = 60 + m["tpot_p99"] * 3000
ax.scatter(m["ttft_p99"], m["throughput"], s=size,
color=colors[name], alpha=0.75, edgecolor="black", lw=0.7, zorder=5)
ax.annotate(f"{name}\n(tpot p99 {m['tpot_p99']*1000:.0f} ms)",
(m["ttft_p99"], m["throughput"]),
xytext=(8, 8), textcoords="offset points", fontsize=9)
ax.set_xlabel("TTFT p99 (s, lower is better)")
ax.set_ylabel("throughput (tokens/s, higher is better)")
ax.set_title("batching strategies on the throughput/TTFT Pareto (bubble size = TPOT p99)")
ax.grid(True, alpha=0.3)
ax.invert_xaxis()
plt.tight_layout()
plt.show()
Exercises#
Plot the Pareto. Scatter (TTFT p99, TPOT p99) for each strategy; size = throughput.
Admission priority. Give short prompts priority when admitting into continuous. Measure the effect on p99 TTFT for long requests.
Chunk-size sweep. For chunked, try
chunk ∈ {256, 512, 1024, 2048}. The sweet spot is usually 1024-1536.
References#
SARATHI-Serve paper.
vLLM’s
core/scheduler.pyfor the production continuous batching code.
s.summary()
s.save()