S10Batching & Scheduling·442 lines

The Scheduler

At every step the scheduler decides which requests run. Fairness, latency SLOs and out-of-memory recovery all live here.

waiting/running queues
token budget
preemption
swap vs recompute
admission control

Motivation

The problem

S09 ended with while self.can_admit(), a function we never wrote. It has to answer, several hundred times per second:

How many requests should run this step?
Should we start a new 8,000-token prefill, or give the whole step to the forty requests already streaming?
We are out of KV blocks and a running request needs one more. Who loses their cache?
A request has been queued for nine seconds behind a long prompt. Is that acceptable?

No answer to these is universally correct. What you have instead is a policy, and that policy is your engine's observable behaviour.

Core idea

The solution

A scheduler with two queues, two budgets, and a preemption path. The structure is short enough to read in full:

the schedulerpython

def schedule(self) -> Batch:
    budget = self.max_num_batched_tokens
    scheduled = []

    # --- 1. running decodes first: a started request should finish ---
    for req in self.running:
        if budget < 1:
            break
        if not self.blocks.can_append(req):
            victim = self.pick_victim()      # last admitted, usually
            self.preempt(victim)
            if victim is req:
                continue
        self.blocks.append(req)
        scheduled.append(req)
        budget -= 1

    # --- 2. spend what is left admitting new prefills ---------------
    while self.waiting and budget > 0:
        req = self.waiting[0]
        if req.prompt_len > budget:
            break                            # cannot fit this step
        if not self.blocks.can_allocate(req, watermark=0.01):
            break                            # keep headroom for growth
        self.waiting.popleft()
        self.blocks.allocate(req)
        scheduled.append(req)
        budget -= req.prompt_len

    return Batch(scheduled)

Decode-first is not arbitrary

A request that has already been prefilled is holding KV memory. Every step it does not run, that memory sits occupied and produces nothing. Prioritising decodes minimises how long cache is held, and that raises the number of requests that fit. Decode-first is a memory policy as much as a latency one.

Mechanics

How it works

Diagramthe scheduler step

Two queues, a token budget, a memory watermark, and a preemption path. Almost every user-visible property of a serving system is decided here.

The two budgets are different constraints

Token budget (max_num_batched_tokens) bounds compute per step. It caps how long one forward pass takes, and so caps inter-token latency for everyone in the batch.
Memory watermark bounds KV blocks. Admitting up to 100% of the pool guarantees deadlock: every running request eventually needs one more block and none is available. Real engines keep a few percent free.

Raise the token budget and throughput improves while per-step latency worsens. Lower it and the reverse. There is no setting that is best for both; that is what an SLO is for.

Preemption: recompute or swap

When memory runs out, a running request must give up its cache. There are two ways to get it back, and the crossover between them is clean:

Recompute — throw the KV away and re-prefill the whole request when it is rescheduled. Cost is (prompt + generated) tokens of prefill: compute-bound, and reasonably fast.
Swap — copy the KV blocks to host memory and back. Cost is twice the cache size over PCIe. Many bytes, no arithmetic.

At short contexts, recompute wins easily. At very long contexts the prefill cost grows and swapping becomes competitive. vLLM defaults to recompute; prefill got much cheaper once chunking and prefix caching arrived.

Fairness, and the case against pure FCFS

Strict first-come-first-served is simple and starvation-free, and it has one bad failure mode: head-of-line blocking. A 32,000-token prompt at the front of the queue cannot be scheduled until the whole budget is free, and while it waits nothing behind it moves either.

Shortest-prompt-first fixes tail latency for small requests and starves large ones. Priority queues let you express business rules, and require you to have some. Scheduling policy belongs to the operator, so make it pluggable and measure P50 and P99. A policy that improves the mean while destroying the tail looks great on a dashboard and terrible to users.

Explore

Try it

Simulatorscheduler with preemptionshrink the KV pool until preemptions start

token budget / step512KV pool (blocks)180

queue policy

preemption

seed6

step 80 / 80

completed: 0/20
avg TTFT: 0.0 steps
p95 TTFT: 0 steps
preemptions: 0

waiting queue · 20

r11873 prompt
r12140 prompt
r171324 prompt
r681 prompt
r1464 prompt
r11114 prompt
r1669 prompt
r131109 prompt
r18194 prompt
r01461 prompt
r19146 prompt
r8182 prompt
r10929 prompt
r15219 prompt
r3136 prompt
r7156 prompt
r980 prompt
r244 prompt
r493 prompt
r5228 prompt

A prompt longer than the token budget can never be scheduled; it starves forever. Drop the budget to 128 and watch the long prompts pile up permanently. S11 fixes this.

running · 0

nothing running

Decodes are scheduled before prefills. A request that has started should finish, because until it does its KV cache stays occupied and everyone else is worse off.

per-step batch composition (prefill vs decode tokens)

blocks used: 0/180
prefill tok: 0
decode tok: 0
tokens recomputed: 0

Purple spikes are prefills. One long prompt can consume the entire step's budget, and every decode behind it waits: a visible latency spike for every user currently streaming, and the exact problem chunked prefill solves.

Three experiments, each reproducing a real production incident:

Set the token budget to 128. Long prompts can never be scheduled and queue forever. That is permanent starvation, and the engine looks healthy throughout. An unschedulable request at the head of an FCFS queue also blocks everything behind it, so a handful of oversized prompts can strand most of the traffic.
Shrink the KV pool to 60 blocks. Preemptions climb and the same tokens get recomputed over and over. This is thrashing: throughput collapses while GPU utilisation reads 100%.
Switch to shortest-prompt-first under sustained load. The median improves sharply and the long prompts do much worse. On a finite queue that drains, shortest-first wins outright; the starvation only appears when short requests keep arriving to overtake the long ones.

Build it

Implementation

code/s10_scheduler.py (excerpt)python

def preempt(self, req: Request):
    """Give a running request's KV memory back to the pool."""
    self.blocks.free(req)
    req.preemption_count += 1

    if self.preemption_mode == "recompute":
        # Cheapest for short contexts: forget everything, re-prefill later.
        req.computed = 0
        req.kv_blocks = []
    else:
        # Cheapest for long contexts: park the blocks in host memory.
        req.swapped_blocks = self.blocks.swap_out(req)

    self.waiting.appendleft(req)     # front of the queue — do not restart it last

Preempted requests go to the FRONT of the queue

Put a preempted request at the back and you can livelock: it is admitted, preempted and sent to the back again and again, making progress only during the few steps it runs before losing its cache again. Front-of-queue reinsertion guarantees forward progress. Picking the victim follows the same reasoning. The last-admitted request has generated the fewest tokens, so it loses the least work.

Run it locally

Runs a full scheduler against a synthetic arrival trace with configurable policies, then sweeps token budget and KV pool size to produce a throughput/latency table. Includes a thrashing detector that flags when recomputation exceeds useful work.

$ python code/s10_scheduler.py

Expect: A baseline that strands 46 of 60 requests behind just 5 oversized prompts, a chunked-prefill sweep that starves nobody at any budget, a policy comparison quoting TTFT per cohort, and a throughput/latency table in milliseconds, since a step is not a unit of time.

Only NumPy is required — setup instructions.

Production notes

In production

vLLM — Scheduler.schedule() with max_num_batched_tokens, max_num_seqs, a GPU watermark and recompute-by-default preemption. It emits a warning when the preemption rate is high; that warning is your thrashing signal.
SGLang — cache-aware scheduling that prefers requests whose prefix is already resident, turning the S07 cache into a scheduling input.
Kubernetes-scale deployments — a second scheduler above this one routes requests across replicas, ideally prefix-aware, so cache hits are not destroyed by round-robin load balancing.

Exercises

1
Implement a priority queue with aging: requests gain priority the longer they wait, so shortest-first cannot starve anyone. Demonstrate the starvation before and after.
2
Add an SLO-aware admission controller that rejects new requests outright when the queue implies a TTFT above your target. Rejecting early is usually better service than accepting and timing out.
3
Instrument preemption thrashing: log the ratio of recomputed tokens to newly generated ones, and trigger a backpressure signal above 20%.

Continue

What's next

The simulator showed the flaw: one long prefill eats an entire step and every streaming user feels the stall. S11 chops prefills into chunks and mixes them into decode batches. That removes the spike without giving up throughput.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

Why does the scheduler run all pending decodes before admitting any new prefill?

Score 0/6