04 — Chunked Prefill & Stall-Free Scheduling (Sarathi-Serve)

Paper: Agrawal, Kedia, Panwar, Mohan, Kwatra, Gulavani, Tumanov, Ramjee — “Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve” (OSDI ‘24). arXiv: <phone_number_or_numberic_id_or_random_id_31>. PDF: https://www.usenix.org/system/files/osdi24-agrawal.pdf.

The problem Sarathi-Serve solves

Recall the picture from file 02, at step t=4:

batch = [R1(decode step 15), R2(decode step 30), R4(prefill of 1024 tokens)]

R4’s prefill is a 1024-token GEMM. On an 8B model that’s roughly 100-200 ms of compute. During that step R1 and R2’s decode step also takes 100-200 ms — because they’re all in one forward pass — even though their per-token decode work is measured in single-digit ms.

From R1’s user’s POV: they were getting nice 30ms tokens, and then suddenly there’s a 200ms gap. ITL p99 spikes. Repeated over a benchmark this destroys the tail latency, which destroys goodput.

Before Sarathi, engines fell into two bad choices:

  1. Prefill-prioritizing (e.g., vanilla vLLM before chunked prefill landed): run prefill of new requests to completion before continuing decode. Optimizes TTFT of new arrivals but starves ongoing decodes → catastrophic ITL.

  2. Decode-prioritizing (e.g., some TGI configurations): finish all decodes before admitting new prefills. Smooth ITL but terrible TTFT for new arrivals + underutilized compute during long decode-only phases.

Neither is Pareto-optimal. Both leave 30-50% of steady-state throughput on the table.

The insight: prefill is a knob, not a phase

A prefill of 1024 tokens is just a big GEMM. You can split it into chunks. Compute the KV for the first 256 tokens this step, the next 256 next step, etc. The math is identical (attention doesn’t care that we added the query tokens in chunks; the K/V for earlier tokens is already in the cache).

Given that, you can now co-schedule prefill chunks with decode steps in one hybrid batch, keeping the total tokens per step within a fixed token budget.

That token budget is the master knob. Say it’s 2048:

  • 30 decode sequences × 1 token each = 30 tokens

  • Remaining 2018 tokens = one prefill chunk of 2018 tokens (or two of 1009, or one of 1024 with room to spare)

Every step now takes approximately the same wall-clock time (~one prefill chunk’s worth of compute), so ITL is smooth. Prefill still advances — just in small bites interleaved with decodes.

This is what the paper calls stall-free scheduling: no decode step ever has to wait for a full prefill to finish before it can execute.

The token budget as an operating point

The token budget is the tunable in modern serving. Its tradeoff:

  • Large token budget (say 8192) → prefill chunks are big → TTFT is low (new requests finish prefill fast) → BUT decode steps that share a batch with a big prefill chunk are slower → ITL spikes still exist, just smaller.

  • Small token budget (say 512) → prefill chunks are tiny → decode steps are basically pure decode → ITL is beautifully smooth → BUT prefills take many more steps → TTFT balloons.

vLLM V1 (chunked prefill on by default since ~mid-2024, hardened in V1) exposes this as --max-num-batched-tokens. Default varies with model size; for an 8B on H100 it’s typically 8192.

Sarathi-Serve’s contribution beyond the mechanism is showing you can pick the token budget based on the model + hardware profile by profiling how prefill and decode compute scale with batch composition, then setting the budget so no step exceeds an ITL target.

Reported numbers

From the paper (numbers to remember):

  • Up to 5.6× serving-capacity improvement on Falcon-180B with pipeline parallelism.

  • 2.6× on Mistral-7B on a single A100 with tighter SLO.

  • 6.9× on Yi-34B on 2×A100 with TP=2.

“Serving capacity” here means goodput — requests/sec that meet a stated latency SLO. These are large multipliers even in the field’s noisy benchmark culture.

The scheduling algorithm, concretely

Essentially what vLLM V1 does now (Sarathi ideas absorbed):

def schedule(token_budget):
    scheduled = []

    # 1. Continue running decodes (cheap, priority)
    for req in running_decodes:
        scheduled.append((req, 1 decode token))
        token_budget -= 1

    # 2. Continue in-progress prefill chunks
    for req in running_prefills:
        chunk = min(req.remaining_prefill, token_budget)
        scheduled.append((req, chunk))
        token_budget -= chunk
        if token_budget == 0: break

    # 3. Admit new waiting requests (start their prefill in chunks)
    while waiting and token_budget > 0:
        req = waiting.popleft()
        cached = prefix_cache.lookup(req)
        req.remaining_prefill = req.prompt_len - cached
        chunk = min(req.remaining_prefill, token_budget)
        scheduled.append((req, chunk))
        token_budget -= chunk

    return scheduled

Note: decodes are scheduled first — they’re the cheap steady-state work that we don’t want to starve. Prefill (both continuation and new) fills the remaining budget.

Interaction with prefix caching

A nice side effect: when prefix caching (file 05) hits, the effective prefill length shrinks dramatically. A request with 4000 prompt tokens and a 3500-token cache hit only needs to compute 500 new tokens of prefill. Combined with chunked prefill, those 500 tokens might fit entirely inside one hybrid step — the request effectively starts generating within one iteration.

This is what makes agentic workloads (yours) so responsive on modern engines: cache hit + chunked prefill = ~10 ms from request arrival to first token. Numerically compare against a naive engine’s ~500ms and you see why the papers exist.

Anti-pattern: the un-tuned token budget

A very common production mistake: leave max-num-batched-tokens at its default, then complain about ITL spikes on long-prompt workloads. The resolution is almost always to lower the budget until the largest plausible prefill chunk stays inside a per-step compute budget derived from your ITL SLO.

Rule of thumb (derive it once for your hardware):

  1. Measure how long a decode-only step takes for the model at your expected batch size. Call this decode_step_ms.

  2. Measure a pure prefill’s ms/token. Call this prefill_ms_per_token.

  3. Set the budget so a hybrid step’s worst-case time is ≤ 2× decode_step_ms: budget ≈ decode_step_ms / prefill_ms_per_token.

On H100 + Llama-3-8B this usually lands around 4096-8192 tokens. On a tighter ITL SLO (say 20ms p99) it drops to 2048.

Reading exercise

Open vllm/v1/core/sched/scheduler.py. Find where token_budget (sometimes named num_new_tokens_budget or similar) is initialized in schedule(). Then trace:

  1. How decode vs prefill requests compete for the budget.

  2. When a prefill is chunked, where the remaining prefill state is tracked between iterations.

  3. What happens when a preempted request comes back — does it re-prefill from scratch, from a swap, or from cache?

When those three answers are clear you own the chunked-prefill implementation.

The intuition to internalize

Prefill is a variable-length compute pool that can be sliced; the token budget is how much of it you spend per step; the goal is to smooth the cost of every step so no ITL spike ever happens. Everything else — the SLO story, the goodput chart, the interaction with prefix caching — follows from that one mechanic.