Scheduling & Prefix Caching · AI infra interview lesson

Why It Matters

A needs ten more tokens; B needs one. Can C join as soon as B finishes? Draw active requests at each iteration before comparing batching, prefill and cache budgets.

A concrete number: four requests generate [20, 200, 50, 30] tokens with a batch limit of 2. Static batching takes 200 + 50 = 250 steps for 300 useful tokens, a utilization of 300 / (2*250) = 60%; continuous batching takes 200 steps, 75% (the Step 1 code computes both). Prefix caching: 100 requests sharing a 2000-token system prompt skip 99 * 2000 = 198,000 prefill tokens; at LLaMA-2-7B's 512 KiB of KV per token, one copy of that prefix is about 1000 MiB, and sharing keeps a single copy in memory.

Walk the Shapes First

nano-vllm's Scheduler with block_size=4, 5 physical blocks and max_num_batched_tokens=8; three requests A (6 tokens), B (5), C (3), each generating at most 4 tokens. Each row is one schedule() call ("length" is the sequence length at scheduling time):

Rules read off the table: whenever some prefill can be scheduled, the step is prefill only; only the first prefill may be chunked; before decode, a sequence needs one new block when len % bs == 1, and if none is available the tail of the running queue is preempted (or the sequence itself, if nobody is left); a preempted sequence drops all of its KV, goes back to the front of the waiting queue, and is prefilled again later, including the tokens it already generated.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: static vs continuous batching

What is new: quantifying "a finished request gives up its slot immediately". Continuous batching is "whichever slot frees up first takes the next request", which a min-heap of slot free times computes without simulating every step.

Formula: static batch k takes max(out_lens[batch]) steps and the total is the sum over batches; under continuous batching request i finishes at finish[i] = heappop(slots) + out_lens[i]. Lower bound: total >= max(max(out_lens), ceil(sum(out_lens) / B)).

import heapq, random

def simulate_static(out_lens, B):
    t, finish = 0, [0] * len(out_lens)
    for s in range(0, len(out_lens), B):
        batch = range(s, min(s + B, len(out_lens)))
        for i in batch:
            finish[i] = t + out_lens[i]
        t += max(out_lens[i] for i in batch)          # the whole batch waits for its longest request
    return t, finish

def simulate_continuous(out_lens, B):
    finish, slots = [0] * len(out_lens), [0] * min(B, len(out_lens))
    for i, L in enumerate(out_lens):                  # in arrival order, the earliest free slot takes it
        finish[i] = heapq.heappop(slots) + L
        heapq.heappush(slots, finish[i])
    return (max(finish) if finish else 0), finish

def naive_continuous(out_lens, B):                    # step-by-step simulation, used as the reference
    t, finish, nxt, running = 0, [0] * len(out_lens), 0, {}
    while nxt < len(out_lens) or running:
        while len(running) < B and nxt < len(out_lens):
            running[nxt] = out_lens[nxt]; nxt += 1
        t += 1
        for i in list(running):
            running[i] -= 1
            if running[i] == 0:
                finish[i] = t; del running[i]
    return t, finish

lens = [20, 200, 50, 30]
assert simulate_static(lens, 2) == (250, [20, 200, 250, 230])
assert simulate_continuous(lens, 2) == (200, [20, 200, 70, 100])
assert sum(lens) / (2 * 250) == 0.6 and sum(lens) / (2 * 200) == 0.75
assert simulate_static([5, 1, 1, 1], 2) == (6, [5, 1, 6, 6])
assert simulate_continuous([5, 1, 1, 1], 2) == (5, [5, 1, 2, 3])
rng = random.Random(0)
for _ in range(300):
    lens = [rng.randint(1, 30) for _ in range(rng.randint(1, 15))]
    B = rng.randint(1, 6)
    got = simulate_continuous(lens, B)
    assert got == naive_continuous(lens, B)
    assert max(max(lens), -(-sum(lens) // B)) <= got[0] <= simulate_static(lens, B)[0]

Step 2: nano-vllm Scheduler: prefill first, preempt the tail, recompute

What is new: a real scheduler also manages memory. nano-vllm's schedule() first schedules as much prefill as it can (bounded by max_num_seqs, max_num_batched_tokens and can_allocate; only the first sequence may be cut into a chunk) and returns right away if it scheduled any; only otherwise does it run decode, preempting whenever can_append fails. postprocess advances num_cached_tokens, discards samples taken in the middle of a chunk, and frees finished sequences.

Key state: num_cached_tokens counts tokens whose KV is already in the cache, num_scheduled_tokens counts what this step computes, and the prompt is done when they add up to num_tokens. Preemption is recompute: just deallocate, and waiting.appendleft makes it the first to resume. The BlockManager below drops prefix caching and keeps only the interface the scheduler uses.

from collections import deque
from enum import Enum, auto

class Status(Enum):
    WAITING = auto(); RUNNING = auto(); FINISHED = auto()

class Seq:
    block_size = 4
    def __init__(self, name, token_ids, max_tokens, ignore_eos=True):
        self.name, self.token_ids, self.max_tokens, self.ignore_eos = name, list(token_ids), max_tokens, ignore_eos
        self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens = len(token_ids), 0, 0
        self.block_table, self.status, self.is_prefill = [], Status.WAITING, True
    @property
    def num_tokens(self): return len(self.token_ids)
    def __len__(self): return len(self.token_ids)
    @property
    def num_blocks(self): return (len(self) + self.block_size - 1) // self.block_size
    @property
    def num_completion_tokens(self): return len(self) - self.num_prompt_tokens
    def append_token(self, t): self.token_ids.append(t)

class BlockManager:                                   # trimmed: no prefix caching
    def __init__(self, num_blocks, block_size):
        self.block_size, self.free_block_ids = block_size, deque(range(num_blocks))
    def can_allocate(self, seq): return 0 if len(self.free_block_ids) >= seq.num_blocks else -1
    def allocate(self, seq, num_cached_blocks):
        seq.block_table = [self.free_block_ids.popleft() for _ in range(seq.num_blocks)]
        seq.num_cached_tokens = 0
    def deallocate(self, seq):
        self.free_block_ids.extend(reversed(seq.block_table))
        seq.block_table, seq.num_cached_tokens = [], 0
    def can_append(self, seq): return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)
    def may_append(self, seq):
        if len(seq) % self.block_size == 1:
            seq.block_table.append(self.free_block_ids.popleft())
    def hash_blocks(self, seq): pass

Choose requests for the next iteration

Inspect waiting and running queues, then choose prefill or decode. Check cache capacity before admission. This teaching model uses preemption and recomputation under pressure. Trace one iteration with two requests before reading the full class.

class Scheduler:
    def __init__(self, num_blocks, block_size, max_num_batched_tokens, max_num_seqs=8, eos=-1):
        self.block_manager = BlockManager(num_blocks, block_size)
        self.block_size, self.max_num_batched_tokens, self.max_num_seqs, self.eos = block_size, max_num_batched_tokens, max_num_seqs, eos
        self.waiting, self.running = deque(), deque()

    def is_finished(self): return not self.waiting and not self.running
    def add(self, seq): self.waiting.append(seq)

    def schedule(self):
        scheduled, num_batched = [], 0
        while self.waiting and len(scheduled) < self.max_num_seqs:          # prefill first
            seq = self.waiting[0]
            remaining = self.max_num_batched_tokens - num_batched
            if remaining == 0:
                break
            if not seq.block_table:
                num_cached_blocks = self.block_manager.can_allocate(seq)
                if num_cached_blocks == -1:
                    break
                num_tokens = seq.num_tokens - num_cached_blocks * self.block_size
            else:                                                           # only one chunk was computed last step
                num_tokens = seq.num_tokens - seq.num_cached_tokens
            if remaining < num_tokens and scheduled:                        # only the first one may be chunked
                break
            if not seq.block_table:
                self.block_manager.allocate(seq, num_cached_blocks)
            seq.num_scheduled_tokens = min(num_tokens, remaining)
            num_batched += seq.num_scheduled_tokens
            if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:
                seq.status = Status.RUNNING
                self.running.append(self.waiting.popleft())
            scheduled.append(seq)
        if scheduled:
            return scheduled, True
        while self.running and len(scheduled) < self.max_num_seqs:          # decode
            seq = self.running.popleft()
            while not self.block_manager.can_append(seq):
                if self.running:
                    self.preempt(self.running.pop())                        # preempt the tail (latest admitted)
                else:
                    self.preempt(seq)                                       # nobody else left
                    break
            else:
                seq.num_scheduled_tokens, seq.is_prefill = 1, False
                self.block_manager.may_append(seq)
                scheduled.append(seq)
        assert scheduled
        self.running.extendleft(reversed(scheduled))
        return scheduled, False

    def preempt(self, seq):
        seq.status, seq.is_prefill = Status.WAITING, True
        self.block_manager.deallocate(seq)                                  # recompute: just drop the KV
        self.waiting.appendleft(seq)

    def postprocess(self, seqs, token_ids, is_prefill):
        for seq, tok in zip(seqs, token_ids):
            self.block_manager.hash_blocks(seq)
            seq.num_cached_tokens += seq.num_scheduled_tokens
            seq.num_scheduled_tokens = 0
            if is_prefill and seq.num_cached_tokens < seq.num_tokens:
                continue                                                    # mid-chunk: no sampling
            seq.append_token(tok)
            if (not seq.ignore_eos and tok == self.eos) or seq.num_completion_tokens == seq.max_tokens:
                seq.status = Status.FINISHED
                self.block_manager.deallocate(seq)
                self.running.remove(seq)

def fake_model(seqs):           # next token depends only on all computed tokens, so recompute is invisible
    return [(sum(s.token_ids[:s.num_cached_tokens + s.num_scheduled_tokens]) * 31 + 7) % 1000 for s in seqs]

def expected(prompt, n):
    toks = list(prompt)
    for _ in range(n):
        toks.append((sum(toks) * 31 + 7) % 1000)
    return toks[len(prompt):]

# the trace from the shape table
sch = Scheduler(num_blocks=5, block_size=4, max_num_batched_tokens=8)
seqs_all = [Seq(n, [ord(n) * 100 + j for j in range(L)], max_tokens=4) for n, L in [("A", 6), ("B", 5), ("C", 3)]]
for s in seqs_all:
    sch.add(s)
trace = []
while not sch.is_finished():
    seqs, is_prefill = sch.schedule()
    trace.append(("P" if is_prefill else "D", [(s.name, s.num_scheduled_tokens) for s in seqs], len(sch.block_manager.free_block_ids)))
    sch.postprocess(seqs, fake_model(seqs), is_prefill)
assert trace == [("P", [("A", 6)], 3), ("P", [("B", 5), ("C", 3)], 0), ("D", [("A", 1), ("B", 1), ("C", 1)], 0),
                 ("D", [("A", 1), ("B", 1)], 1), ("D", [("A", 1), ("B", 1)], 0), ("P", [("C", 5)], 3), ("D", [("C", 1)], 3)]
assert all(s.token_ids[s.num_prompt_tokens:] == expected(s.token_ids[:s.num_prompt_tokens], 4) for s in seqs_all)
assert len(sch.block_manager.free_block_ids) == 5

# tight memory + chunked prefill: outputs must match generating each request alone
for bs in [4, 16, 256]:
    Seq.block_size = bs
    prompts = [list(range(i * 10000, i * 10000 + n)) for i, n in enumerate([3 * bs, 1, bs + 1, 2 * bs - 1, bs])]
    sch = Scheduler(num_blocks=5, block_size=bs, max_num_batched_tokens=bs + bs // 2)
    seqs_all = [Seq(str(i), p, max_tokens=bs + 2) for i, p in enumerate(prompts)]
    for s in seqs_all:
        sch.add(s)
    n_prefill_steps = 0
    while not sch.is_finished():
        seqs, is_prefill = sch.schedule()
        assert not is_prefill or sum(s.num_scheduled_tokens for s in seqs) <= bs + bs // 2
        n_prefill_steps += is_prefill
        sch.postprocess(seqs, fake_model(seqs), is_prefill)
    assert all(s.token_ids[len(p):] == expected(p, bs + 2) for s, p in zip(seqs_all, prompts))
    assert len(sch.block_manager.free_block_ids) == 5 and n_prefill_steps > len(prompts)   # chunking or recompute happened

Step 3: chunked prefill and the token budget

What is new: prefilling one very long prompt in a single step makes that step slow, every concurrent decode waits for it, and TPOT spikes. (Pro lesson)

Allocate a token budget

Equal request counts can hide very different prefill work. Chunk prompts so prefill and decode share a token budget. (Pro lesson)

Step 4: RadixTree: longest-prefix match and edge splitting

What is new: nano-vllm's prefix cache works at block granularity through chained hashes, so it only hits whole blocks. (Pro lesson)

Step 5: lock_ref and LRU eviction

What is new: when the cache is full, something has to go. (Pro lesson)

In Real Systems

nano-vllm nanovllm/engine/scheduler.py : schedule , preempt and postprocess match Step 2 line by line; (Pro lesson)

Common Pitfalls

Appending the sampled token after a mid-chunk prefill → a wrong token enters the sequence and everything after it is wrong → in postprocess… (Pro lesson)

Say It in the Interview

Continuous batching schedules at the iteration level: finished requests leave the batch after every step and waiting ones join, so short req… (Pro lesson)

Hands-on Notebook

The first half of the notebook follows nano-vllm's Scheduler interface (the setup cell provides a full BlockManager ), and the second half i… (Pro lesson)