Continuous batching (Orca)

Open in Colab ▶️ Run this notebook in Colab

Continuous batching (Orca)#

Track 01 - Inference · Notebook 04 · Runtime: ≈30 s on CPU

Prerequisites: 01_inference/01 (KV cache), 05_serving/01 (roofline).

Paper: Yu et al. 2022, Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI’22).

What#

Static batching waits for every request in the batch to finish before starting new ones - long requests block short ones, GPU utilisation drops. Continuous batching (aka iteration-level scheduling, aka Orca) schedules per decode step: finished sequences leave the batch, new sequences enter mid-stream.

We build a discrete-event simulator (no simpy dependency; pure Python) to compare both schedulers on the same 500-request Poisson-arrival workload with log-normal output lengths. Orca-style continuous batching delivers ≥2× throughput and dramatically better p99 TTFT under heavy load.

from llm_systems_cookbook.nb import bootstrap

import heapq
import math
from dataclasses import dataclass, field

import numpy as np

s = bootstrap("01_inference_04_continuous_batching_orca")

Calibrated step-latency model#

Per decode step on a batch of B sequences totalling T tokens in KV cache:

step_latency(B, T) = 5 ms + 0.8 ms * B + 0.02 ms * T

The constants capture real behaviour: per-step overhead (kernel launches, ~5 ms), per-sequence attention compute (~0.8 ms), and per-token KV read (~0.02 ms).

def step_latency(batch_size: int, total_tokens: int) -> float:
    '''Seconds per decode step on a batch with given total kv tokens.'''
    return 0.005 + 0.0008 * batch_size + 0.00002 * total_tokens

Poisson arrivals, log-normal output lengths#

500 requests, arrival rate λ = 30 req/s, output lengths log-normal (median ~100, some up to 2000). The long-tail distribution is what makes static batching painful - a single 2000-token request keeps the whole batch busy.

rng = np.random.default_rng(0)
N = 500
LAMBDA = 30.0  # req/s

arrivals = np.cumsum(rng.exponential(1.0 / LAMBDA, size=N))
output_lengths = np.clip(rng.lognormal(mean=4.6, sigma=0.9, size=N).astype(int), 10, 2000)
prompt_lengths = np.clip(rng.lognormal(mean=4.0, sigma=0.7, size=N).astype(int), 5, 500)
print(f"arrivals span      = {arrivals[-1]:.1f} s")
print(f"output length p50  = {int(np.percentile(output_lengths, 50))}")
print(f"output length p95  = {int(np.percentile(output_lengths, 95))}")
print(f"output length max  = {int(output_lengths.max())}")

Static vs continuous#

Static: admit B requests at once; run until all of them finish; admit the next batch. Pathologically sensitive to the longest request in each batch.

Continuous: at every decode step, admit any waiting requests that fit; at every step, evict finished sequences immediately. The batch size changes every step.

@dataclass(order=True)
class Request:
    arrival: float
    idx: int = field(compare=False)
    prompt_len: int = field(compare=False)
    output_len: int = field(compare=False)
    started_at: float | None = field(default=None, compare=False)
    finished_at: float | None = field(default=None, compare=False)
    tokens_done: int = field(default=0, compare=False)


def build_requests() -> list[Request]:
    return [
        Request(arrival=float(arrivals[i]), idx=i,
                prompt_len=int(prompt_lengths[i]),
                output_len=int(output_lengths[i]))
        for i in range(N)
    ]


def simulate_static(reqs: list[Request], max_batch: int = 32) -> None:
    '''Run static batching: fill a batch, decode until all finish, repeat.'''
    pending = list(reqs)
    now = 0.0
    while pending:
        # Wait for at least one request to arrive.
        first = pending[0]
        if first.arrival > now:
            now = first.arrival
        # Admit up to max_batch requests that have arrived.
        batch: list[Request] = []
        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.started_at = now
        # Decode until the longest request in the batch finishes.
        remaining = {r.idx: r.output_len for r in batch}
        while remaining:
            total_tokens = sum(r.prompt_len + (r.output_len - remaining[r.idx]) for r in batch if r.idx in remaining)
            now += step_latency(len(remaining), total_tokens)
            for idx in list(remaining):
                remaining[idx] -= 1
                if remaining[idx] <= 0:
                    del remaining[idx]
                    done = next(r for r in batch if r.idx == idx)
                    done.finished_at = now


def simulate_continuous(reqs: list[Request], max_batch: int = 32) -> None:
    '''Run continuous batching: per-step admission and eviction.'''
    pending = list(reqs)
    active: list[Request] = []
    now = 0.0
    while pending or active:
        # Pull in any newly-arrived waiting requests up to max_batch.
        while pending and pending[0].arrival <= now and len(active) < max_batch:
            r = pending.pop(0)
            r.started_at = now
            active.append(r)
        if not active:
            # Jump to next arrival if idle.
            now = pending[0].arrival
            continue
        total_tokens = sum(r.prompt_len + r.tokens_done for r in active)
        now += step_latency(len(active), total_tokens)
        for r in active:
            r.tokens_done += 1
        # Evict finished.
        finished = [r for r in active if r.tokens_done >= r.output_len]
        for r in finished:
            r.finished_at = now
        active = [r for r in active if r.tokens_done < r.output_len]


static_reqs = build_requests()
simulate_static(static_reqs)

cont_reqs = build_requests()
simulate_continuous(cont_reqs)


def summarise(reqs: list[Request]) -> dict:
    makespan = max(r.finished_at for r in reqs if r.finished_at is not None)
    ttft = [float(r.started_at - r.arrival) for r in reqs if r.started_at is not None]
    e2e = [float(r.finished_at - r.arrival) for r in reqs if r.finished_at is not None]
    throughput = sum(r.output_len for r in reqs) / makespan
    return {
        "throughput_tok_s": throughput,
        "ttft_p50": float(np.percentile(ttft, 50)),
        "ttft_p99": float(np.percentile(ttft, 99)),
        "e2e_p50": float(np.percentile(e2e, 50)),
        "e2e_p99": float(np.percentile(e2e, 99)),
    }


s_static = summarise(static_reqs)
s_cont = summarise(cont_reqs)
for k in s_static:
    print(f"  {k:<18}  static={s_static[k]:10.3f}   continuous={s_cont[k]:10.3f}")

s.check(
    "continuous_higher_throughput",
    lambda: s_cont["throughput_tok_s"] >= 1.05 * s_static["throughput_tok_s"],
    msg=f"static={s_static['throughput_tok_s']:.1f}  continuous={s_cont['throughput_tok_s']:.1f}",
)
s.check(
    "continuous_lower_ttft_p99",
    lambda: s_cont["ttft_p99"] < s_static["ttft_p99"],
    msg=f"static={s_static['ttft_p99']:.2f}s  continuous={s_cont['ttft_p99']:.2f}s",
)
s.check(
    "continuous_lower_e2e_p99",
    lambda: s_cont["e2e_p99"] < s_static["e2e_p99"],
    msg=f"static={s_static['e2e_p99']:.2f}s  continuous={s_cont['e2e_p99']:.2f}s",
)
s.check(
    "every_request_finishes",
    lambda: all(r.finished_at is not None for r in cont_reqs),
    msg="continuous scheduler must finish every request",
)

The scheduling gap, visualised#

Each horizontal bar is one request’s wall-clock lifetime - left edge is admission, right edge is finish. Static (top) admits in fixed groups and must wait for the slowest in each group; continuous (bottom) evicts finished sequences every step and fills the slot. The tails retract dramatically.

import matplotlib.pyplot as plt

N_SHOW = 60  # show the first N requests so the gantt stays readable

def gantt_rows(reqs):
    rows = []
    for r in reqs[:N_SHOW]:
        if r.started_at is None or r.finished_at is None:
            continue
        rows.append((r.idx, r.started_at, r.finished_at - r.started_at))
    return rows

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(9, 4.5), sharex=True)
for ax, reqs, title in [
    (ax1, static_reqs, f"static batching - p99 e2e {s_static['e2e_p99']:.1f}s"),
    (ax2, cont_reqs,   f"continuous batching - p99 e2e {s_cont['e2e_p99']:.1f}s"),
]:
    for idx, start, dur in gantt_rows(reqs):
        ax.barh(idx, dur, left=start, height=0.8, color="tab:blue", alpha=0.6)
    ax.set_ylabel("request id")
    ax.set_title(title)
    ax.grid(True, axis="x", alpha=0.3)
ax2.set_xlabel("wall-clock time (s)")
fig.tight_layout()
plt.show()

Exercises#

  1. PagedAttention memory model. In continuous batching the KV cache grows irregularly. Extend simulate_continuous to track cumulative KV bytes and refuse admission when memory fills up.

  2. Priority classes. Let half the requests be “interactive” and half “batch”; interactive gets a higher admission priority. Plot the tradeoff curve.

  3. Real workload. Replace log-normal with actual prompt/completion length distributions from the ShareGPT dataset. The tail is heavier.

References#

  • Orca OSDI’22 paper for the original proposal.

  • vLLM’s scheduler.py for the production pattern.

s.summary()
s.save()