24. Serving many users: continuous batching

One user leaves the GPU idle

Decode at batch 1 is memory-bound: each step reads every weight once to produce one token (Chapter 10). A 1,024 × 4,096 BF16 linear layer does 2 FLOPs per weight read, for an arithmetic intensity of about 1 FLOP per byte. An H100 can do around 300 FLOPs in the time it reads one byte. At batch 1, more than 99% of its arithmetic capability idles.

Now run 64 users’ decode steps in one forward pass. The weights are read once and multiplied by 64 vectors: 64 times the work for the same weight traffic. PMPP §20.6 works the numbers for that layer: arithmetic intensity grows from 1 FLOP/B at batch 1 to 315 FLOP/B at batch 512. Until the batch is large enough to become compute-bound, each extra user’s decode is nearly free.

There’s an important exception. Each user has their own KV cache, so attention reads grow linearly with the batch:

$$ \text{bytes per step} \approx \underbrace{\text{weight bytes}}{\text{shared}} + \sum{\text{requests}} \underbrace{\text{KV bytes}(\text{context}i)}{\text{per request}} . $$

For Qwen3-0.6B (1.19 GB of weights, 112 KiB of KV per token), 64 users at 2,000 tokens each read about 15 GB of cache per step, twelve times the weights. Batching helps until the KV cache dominates, and the KV cache also limits how many users fit in memory. Those two facts drive the next two chapters: Chapter 25 packs caches without waste, and Chapter 26 gets more tokens out of each cache read.

Static batching and its waste

The simplest batching collects $B$ requests, pads their prompts to the same length, and generates until all of them finish:

step:         1  2  3  4  5  6  7  8  9  10
request A:    ■  ■  ■  ✓  ·  ·  ·  ·  ·  ·      finished at step 4, slot idles
request B:    ■  ■  ■  ■  ■  ■  ■  ■  ■  ✓
request C:    ■  ■  ✓  ·  ·  ·  ·  ·  ·  ·      new request D waits for the whole batch

Outputs vary from a few tokens to thousands, so most slots idle most of the time, and newly arrived requests wait for the slowest member of the current batch (head-of-line blocking).

Continuous batching (Orca, Yu et al., 2022) makes the decision at every iteration instead: after each step, finished requests leave and waiting requests join. The batch is a set of independent requests that happen to share a forward pass.

Prefill and decode in one step

Joining requests need their prompts prefilled, which is a very different workload from decode (Chapter 16): hundreds of tokens at once, compute-bound. If a 4,000-token prompt is prefilled in one step, every running user’s next token waits for it, and their time per output token jumps.

The fix is a token budget per step and chunked prefill (Sarathi-Serve, Agrawal et al., 2024):

1. Admit waiting requests while cache slots are free.
2. Decode first: every request that has finished its prefill gets one token in this step.
3. Spend the rest of the budget on prefill chunks of at most prefill_chunk tokens, oldest request first.

A long prompt is then spread over several steps, and decode latency stays steady. The budget is the knob: a large budget finishes prefills sooner (lower time to first token) but makes each step slower (higher time per output token).

class Scheduler:
    """Decode-first, token-budgeted, chunked-prefill planner.  (Your engine: Chapter 24)

    Each step: admit waiting requests while slots are free; give every running request that
    has finished prefill one decode token; spend what is left of the budget on prompt chunks
    of at most prefill_chunk tokens, oldest request first.
    """

    def __init__(self, max_slots=8, token_budget=64, prefill_chunk=32):
        if min(max_slots, token_budget, prefill_chunk) < 1:
            raise ValueError("Scheduler limits must be positive")
        self.max_slots, self.token_budget, self.prefill_chunk = max_slots, token_budget, prefill_chunk
        self.free_slots = deque(range(max_slots))
        self.waiting, self.running = deque(), []

    def submit(self, request):
        if not request.prompt or request.max_new_tokens < 1:
            raise ValueError("A request needs a prompt and at least one new token")
        self.waiting.append(request)

    def plan(self):
        """(Your engine: Chapter 24)"""
        while self.waiting and self.free_slots:
            request = self.waiting.popleft()
            request.slot = self.free_slots.popleft()
            self.running.append(request)
        budget = self.token_budget
        decode = [r for r in self.running if r.prefilled == len(r.prompt)][:budget]
        budget -= len(decode)
        prefill = []
        for request in self.running:
            remaining = len(request.prompt) - request.prefilled
            if remaining and budget:
                length = min(remaining, self.prefill_chunk, budget)
                prefill.append((request, request.prefilled, length))
                budget -= length
        return Plan(prefill, decode)

    def release(self, request):
        """Return a finished or cancelled request's slot."""
        self.running.remove(request)
        self.free_slots.append(request.slot)

    @property
    def idle(self):
        return not self.waiting and not self.running

The scheduler is pure policy: it moves no tensors. Keeping it separate from execution makes it testable with plain Python lists, and it’s where production engines differ most (priorities, preemption, fairness, prefix-aware ordering).

The engine loop

class ContinuousBatchingEngine:
    """Runs Scheduler plans on a model that follows the book's cache protocol.  (Your engine: Chapter 24)"""

    def __init__(self, model, max_slots=8, capacity=512, token_budget=64, prefill_chunk=32):
        self.model = model.eval()
        p = next(model.parameters())
        layers, kv_heads, head_dim = model.cache_spec()
        self.cache = StaticKVCache(layers, max_slots, kv_heads, capacity, head_dim, p.device, p.dtype)
        self.capacity, self.device = capacity, p.device
        self.scheduler = Scheduler(max_slots, token_budget, prefill_chunk)
        self.generators = {}
        self.steps = 0

    def submit(self, request):
        if len(request.prompt) + request.max_new_tokens > self.capacity:
            raise ValueError(f"{request.rid}: prompt plus output exceed slot capacity")
        request.arrival = time.perf_counter()
        self.generators[request.rid] = torch.Generator(device=self.device).manual_seed(request.seed)
        self.scheduler.submit(request)

    def _emit(self, request, logits_row):
        token = int(sample(logits_row[None], request.temperature, generator=self.generators[request.rid]))
        if request.first_token_time is None:
            request.first_token_time = time.perf_counter()
        request.output.append(token)
        if token in request.stop_ids:
            request.finish_reason = "stop"
        elif len(request.output) >= request.max_new_tokens:
            request.finish_reason = "length"

    @torch.inference_mode()
    def step(self):
        """One engine iteration. Returns the requests that finished during it.  (Your engine: Chapter 24)"""
        plan = self.scheduler.plan()
        for request, start, length in plan.prefill:            # prompt chunks, one request at a time
            ids = torch.tensor([request.prompt[start:start + length]], device=self.device)
            positions = torch.arange(start, start + length, device=self.device)[None]
            rows = torch.tensor([request.slot], device=self.device)
            logits = self.model(ids, self.cache, positions=positions, rows=rows)
            request.prefilled += length
            if request.prefilled == len(request.prompt):        # the last prompt position predicts token 1
                self._emit(request, logits[0, -1])
        if plan.decode:                                         # every decoding request, one batched forward
            ids = torch.tensor([[r.output[-1]] for r in plan.decode], device=self.device)
            positions = torch.tensor([[len(r.prompt) + len(r.output) - 1] for r in plan.decode], device=self.device)
            rows = torch.tensor([r.slot for r in plan.decode], device=self.device)
            logits = self.model(ids, self.cache, positions=positions, rows=rows)
            for i, request in enumerate(plan.decode):
                self._emit(request, logits[i, -1])
        finished = [r for r in self.scheduler.running if r.done]
        for request in finished:
            request.finish_time = time.perf_counter()
            self.scheduler.release(request)
            self.generators.pop(request.rid, None)
        self.steps += 1
        return finished

    def run(self, requests):
        """Submit everything, step until idle, return {rid: request}."""
        for request in requests:
            self.submit(request)
        results = {}
        while not self.scheduler.idle:
            for request in self.step():
                results[request.rid] = request
        return results

Each request owns one row of a StaticKVCache (Chapter 19), its slot. Rows hold different lengths; writes go to (row, position) pairs and the causal rule hides everything a row hasn’t written. That’s why one batched forward can decode requests at positions 7, 58 and 1,203 at once: the model receives positions of shape [B, 1] and rows of shape [B], and the cache protocol does the rest. No padding of prompts, no attention mask beyond key_pos <= query_pos.

Per-request state lives on the Request: prompt, output, how much has been prefilled, stop tokens, finish reason, and its own random generator, seeded per request so that a request’s samples don’t depend on who else shares the batch.

The correctness test

Batching is an optimization, so it must not change any request’s output. The milestone test runs four requests with different prompt and output lengths through 3 slots, with chunked prefill and a small budget, and checks that each one’s greedy output equals the output of generating it alone. This catches wrong positions, wrong rows, a stale cache slot from a previous request and off-by-one errors in the prefill-to-decode handoff. (With sampling, outputs match only if each request’s random stream is independent of the batch, which is why generators are per request.)

Note

Even with correct code, batched and solo outputs can differ in the last bits in BF16 on a GPU: a matmul over a batch of 8 may use a different kernel or reduction order than a batch of 1. Greedy tokens can then diverge at a near-tie. Production engines that need batch-invariant results use kernels designed for it (Thinking Machines, Defeating Nondeterminism in LLM Inference, 2025).

Run it

python run.py batch --requests 12 --slots 8

Twelve requests with prompts of 8-64 tokens and outputs of 8-48 tokens, through the 2-layer test model on a laptop CPU, first one at a time, then continuously batched:

{"requests": 12, "output_tokens": 322, "serial_tok_s": 398.4, "batched_tok_s": 1037.6, "engine_steps": 55}
{"rid": "5", "ttft_ms": 25.06, "tpot_ms": 8.15, "tokens": 9, "finish": "length"}
{"rid": "6", "ttft_ms": 29.77, "tpot_ms": 7.69, "tokens": 10, "finish": "length"}
{"rid": "2", "ttft_ms": 9.11, "tpot_ms": 7.53, "tokens": 21, "finish": "length"}

2.6× the throughput, in 55 engine steps instead of 322 sequential ones. With --slots 1 the engine does no batching and reached 549 tokens/s against the serial loop’s 640 on the same run: the scheduler’s bookkeeping costs a little, and batching is where the win comes from. On a GPU, where a batch of 8 costs about the same as a batch of 1, the gain approaches the slot count.

Serving metrics

Report percentiles (p50, p90, p99), not averages: a scheduler that’s great on average and terrible for 1% of users is a bad scheduler. And benchmark with a realistic arrival process (requests arriving over time, for example Poisson), not everything submitted at once, which hides queueing.

Production engines add more scheduling tools: preemption (pause a request and free its cache when memory runs out, then recompute or swap it back in), priorities, prefix-aware ordering (Chapter 25), and disaggregation, running prefill and decode on separate GPUs so the two workloads stop interfering (DistServe, Splitwise, 2024).

Batched prefill in production

This engine prefills one request’s chunk per forward call, which is simple and correct but launches many small forward passes. Production engines flatten all of a step’s tokens, prefill chunks and decode tokens of every request, into one long sequence of $N$ tokens, with per-token positions and request IDs. The linear layers see one [N, d] matrix, ideal for tensor cores. Attention uses a “varlen” kernel that receives the boundaries between requests (cu_seqlens in FlashAttention) and each request’s cache location. Your model’s positions-based interface already supports this; what’s missing is an attention kernel that takes ragged batches (Chapter 25’s paged decode kernel is a step in that direction).

Build it

pytest tests/test_ch24_batching.py
python run.py batch --impl engine

Stretch exercises

1. ★ Sweep --slots from 1 to 32 on a GPU with Qwen3-0.6B, and plot throughput and mean TPOT. Where does throughput stop growing? Where: create experiments/ch24.py from run.py’s cmd_batch, replacing its tiny model with engine.loaders.load_qwen3 and collecting per-request timings. Sweep max_slots in that script; run.py batch --slots N is the tiny-model baseline.
2. ★★ Add Poisson arrivals: submit requests at random times while the engine runs, and report p50 and p99 TTFT for token budgets 64, 256 and 1,024. Where: experiments/ch24.py (create it), adapting run.py’s cmd_batch and submitting requests to engine.scheduler.ContinuousBatchingEngine.
3. ★★ Serve your engine over HTTP with an OpenAI-compatible /v1/completions endpoint (FastAPI and server-sent events), running the engine loop in a background thread and streaming each request’s tokens as they’re emitted. Where: new experiments/batching_server.py, wrapping engine.scheduler.ContinuousBatchingEngine; Chapter 36 supplies the later unified server.
4. ★★★ Implement preemption: when a new high-priority request arrives and no slot is free, evict the request with the most remaining work, and recompute its cache from prompt + output when it’s readmitted. Verify outputs are still identical to solo runs. Where: Scheduler and ContinuousBatchingEngine in engine/scheduler.py.

Check your understanding

1. Why does batching speed up the weight reads of decode but not its KV-cache reads?
2. What does continuous batching fix that static batching can’t?
3. Why does the scheduler give decode tokens priority over prefill chunks?
4. How does one batched forward serve requests at different positions without padding?
5. Why should each request have its own random generator?

Going deeper

• PMPP §20.6 (KV cache arithmetic intensity and memory requirement, pp. 504-508): batching’s effect on intensity and the batch-size versus context-length trade-off.
• Yu et al., Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI 2022); Agrawal et al., Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve (OSDI 2024); Zhong et al., DistServe (2024).
• GPU Mode L35 (SGLang performance optimization), and the vLLM and SGLang schedulers (vllm/v1/core/sched/scheduler.py, sglang/srt/managers/scheduler.py), which follow the plan-then-execute structure used here.
• Anyscale’s blog post How continuous batching enables 23x throughput in LLM inference for measurements on real GPUs.