3.5 Scheduler and continuous batching
The scheduling core of nano-vllm: Waiting/Running queues, separate Prefill/Decode scheduling, Token Budget, chunked prefill, and preemption with recomputation when GPU memory runs short
📑 Contents
- 1. Three questions the scheduler answers
- 2. Two queues and one schedule
- 3. Prefill scheduling: spend the Token Budget
- 4. Decode scheduling: preempt to make room
- 5. postprocess: write back and decide finish
- Summary
- Self-check checklist
- References
1. Three questions the scheduler answers
Every engine step (step) starts by asking the scheduler: who runs this round? Prefill or decode? Does their KV fit? scheduler.py is only about 90 lines; its job is to answer those three questions. It never touches matrix multiplication. It only decides how to allocate resources and assemble a batch.
Think of the scheduler as a restaurant host: guests at the door are waiting, guests already seated and ordering are running. Each round the host first checks whether another table can be seated at the door; only if that fails does it look after the seated guests. When tables run short, the table seated last is asked to stand up and rejoin the line at the door.
nano-vllm follows the same default as vLLM V0: a round is either all prefill or all decode, never mixed. The two phases have different input shapes: prefill is several variable-length full prompts and goes through flash_attn_varlen_func; decode is one token per sequence, goes through flash_attn_with_kvcache, and can also replay a CUDA Graph (see 3.3, 3.6). With separate scheduling, each step only prepares one kind of input and takes one execution path.
2. Two queues and one schedule
The scheduler keeps two deques and a few budget constraints:
class Scheduler:
def __init__(self, config):
self.max_num_seqs = config.max_num_seqs # max sequences per round
self.max_num_batched_tokens = config.max_num_batched_tokens # max tokens per round
self.eos = config.eos
self.block_size = config.kvcache_block_size
self.block_manager = BlockManager(config.num_kvcache_blocks,
config.kvcache_block_size)
self.waiting = deque() # sequences that have not finished prefill (new requests, unfinished chunks, preempted ones)
self.running = deque() # prefill done; decoding token by token
How a sequence moves is below. Yellow is the waiting queue, blue is the running queue, and the arc underneath is preemption:
schedule() returns (seqs, is_prefill): the chosen batch and whether this round is prefill or decode. The logic is prefill first: as long as the prefill branch picks at least one sequence, this round does prefill; only when it picks none (waiting is empty, or the head sequence does not have enough blocks) does it go to running for decode.
3. Prefill scheduling: spend the Token Budget
The prefill branch takes sequences from the head of waiting one by one and packs them into a batch until it hits any budget limit:
def schedule(self):
scheduled_seqs = []
num_batched_tokens = 0
while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
seq = self.waiting[0]
remaining = self.max_num_batched_tokens - num_batched_tokens # remaining token budget
if remaining == 0:
break
if not seq.block_table: # first schedule: check prefix cache, confirm blocks fit
num_cached_blocks = self.block_manager.can_allocate(seq)
if num_cached_blocks == -1: # not enough blocks
break
num_tokens = seq.num_tokens - num_cached_blocks * self.block_size
else: # resume chunked prefill: blocks already allocated
num_tokens = seq.num_tokens - seq.num_cached_tokens
if remaining < num_tokens and scheduled_seqs:
break # budget short and batch already has sequences: stop
if not seq.block_table:
self.block_manager.allocate(seq, num_cached_blocks)
# allow chunked prefill only for "the first sequence this round": compute only what fits the budget
seq.num_scheduled_tokens = min(num_tokens, remaining)
num_batched_tokens += seq.num_scheduled_tokens
if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:
seq.status = SequenceStatus.RUNNING # move to running only when the full prompt is done
self.waiting.popleft(); self.running.append(seq)
scheduled_seqs.append(seq)
if scheduled_seqs:
return scheduled_seqs, True # is_prefill = True
# ...otherwise enter the decode branch (see section 4)
Two budgets together make the Token Budget:
- 🔑
max_num_batched_tokens: how many tokens one round may process, mainly a prefill constraint (prompts can be long). Over it, stop adding to the batch. - 🔑
max_num_seqs: how many sequences one round may hold, mainly a decode constraint (each sequence is one token, but the count cannot grow without bound).
The loop has three exits, each with a different meaning:
can_allocatereturns-1→break: when the head sequence does not fit, shorter requests behind it are not scheduled either, even if they would fit. That is strict first-come-first-served, and the cost is head-of-line blocking.remaining < num_tokens and scheduled_seqs→break: only when the batch already has sequences does a short budget stop packing. If the head is a very long prompt and the batch is still empty, the scheduler usesmin(num_tokens, remaining)and schedules only the in-budget chunk. That is chunked prefill. The source comment says “only allow chunked prefill for the first seq”.remaining == 0→break: after the first sequence is truncated and the budget is exactly used up, the next loop iteration exits immediately, so a truncated sequence is always the only one in that round.
A truncated sequence does not leave waiting: it stays at the head, block_table was fully allocated on the first schedule, status is still WAITING, and is_prefill is still True. The next round takes the else branch, computes the remainder as num_tokens - num_cached_tokens, and keeps prefilling until the last chunk is scheduled and it moves into running.
💡 Tip: prefill-first means that as long as the head of waiting can be scheduled, sequences in running do not decode this round. When a long prompt is split into several chunked-prefill pieces, requests that are already generating pause for several steps, and the gap between output tokens grows. The vLLM V1 scheduler drops the prefill/decode phase split: each step assigns how many tokens to compute per request under a fixed token budget (the schedule result is {request_id: num_tokens}), so prefill chunks and decode tokens can sit in the same batch.
4. Decode scheduling: preempt to make room
When the prefill branch picks nothing, schedule() turns to decode the sequences in running. Decode adds only +1 token per sequence per step; the risk is that this token needs a new block and free blocks are already gone. nano-vllm’s answer is preemption:
# decode branch
while self.running and len(scheduled_seqs) < self.max_num_seqs:
seq = self.running.popleft()
while not self.block_manager.can_append(seq):
if self.running:
self.preempt(self.running.pop()) # sacrifice the sequence at the running tail
else:
self.preempt(seq); break # nobody else to sacrifice; sacrifice yourself
else:
seq.num_scheduled_tokens = 1
seq.is_prefill = False # set False the first time it enters a decode batch
self.block_manager.may_append(seq) # take a new block if a boundary is really crossed
scheduled_seqs.append(seq)
assert scheduled_seqs
self.running.extendleft(reversed(scheduled_seqs)) # put chosen sequences back at the head in original order
return scheduled_seqs, False # is_prefill = False
The inner while may preempt several sequences in a row until the current one can append. After the loop, extendleft(reversed(...)) puts this round’s chosen sequences back at the head of running in their original order; sequences that max_num_seqs never reached stay in the queue behind them.
What preemption does:
def preempt(self, seq):
seq.status = SequenceStatus.WAITING # send back to the waiting queue
seq.is_prefill = True # must prefill again next time
self.block_manager.deallocate(seq) # free every KV block it held; num_cached_tokens cleared
self.waiting.appendleft(seq) # insert at the head; reschedule with priority next round
📌 Key point: nano-vllm’s preemption strategy is recomputation, not swap. A preempted sequence frees its KV and goes back to waiting; next time the prompt plus the tokens already generated are treated as one full stretch and prefilled again. That is much simpler than moving KV to CPU memory and back; the cost is recomputation. The recomputed amount is usually smaller than the whole stretch: deallocate only returns blocks to the free queue and does not clear hashes (see 3.4), so as long as those blocks have not been overwritten by other requests, the re-prefill can hit the prefix cache.
What gets sacrificed is the tail of running, i.e. the sequence that finished prefill most recently. Earlier requests can keep moving, so overall behavior is close to first-come-first-served. That sequence usually also has the fewest generated tokens, but the KV it holds still depends on prompt length and is not necessarily the smallest.
⚠️ Note: the else of while ... else runs only when the loop condition becomes false (can_append succeeds). A break out of the loop skips it. The “sacrifice yourself” branch ends with break, so that sequence does not enter this round’s batch; preempt already put it back at the head of waiting.
5. postprocess: write back and decide finish
After the model runs and tokens are sampled, the engine calls scheduler.postprocess to write results back into sequences and decide whether they are done:
def postprocess(self, seqs, token_ids, is_prefill):
for seq, token_id in zip(seqs, token_ids):
self.block_manager.hash_blocks(seq) # register hashes for blocks newly filled this round (see 3.4)
seq.num_cached_tokens += seq.num_scheduled_tokens # tokens computed this round count as cached
seq.num_scheduled_tokens = 0
if is_prefill and seq.num_cached_tokens < seq.num_tokens:
continue # chunked prefill not finished: drop this round's sample
seq.append_token(token_id) # append the newly sampled token
if (not seq.ignore_eos and token_id == self.eos) \
or seq.num_completion_tokens == seq.max_tokens:
seq.status = SequenceStatus.FINISHED
self.block_manager.deallocate(seq) # free KV on finish
self.running.remove(seq)
A few points:
- Middle chunks of chunked prefill still sample: the model still computes logits at the last position of this chunk and samples a token, but that position is not the end of the prompt, so the result is meaningless and
continuedrops it. Only after the last chunk finishes is the sampled token the real first generated token. - Finish conditions: hit EOS (and
ignore_eosis not set), or the generated count reachesmax_tokens. Either one marksFINISHED, frees KV blocks, and removes the sequence fromrunning; unfinished sequences stay inrunningfor the next decode round. A request withmax_tokens = 1finishes on this prefill step; by thenschedulehas already moved it intorunning, soremovedoes not fail. ignore_eos/max_tokensare instance attributes onSequence(flattened fromSamplingParamsat construction; see 3.4).
Putting this section together, the scheduling loop each step looks like the figure below. Left is the prefill path (yellow), right is the decode path (blue), and the red box is preemption when decode runs out of KV:
📝 Summary
- The scheduler only decides; it does not run matrix math: it keeps
waiting/runningdeques, andschedule()returns(seqs, is_prefill). - Prefill first, no mixed batches: if the prefill branch picks sequences, do prefill; otherwise decode. When the head of
waitingdoes not fit, stop immediately; strict first-come-first-served. - Token Budget =
max_num_batched_tokens(mainly prefill) +max_num_seqs(mainly decode). Only the first sequence of a round can be truncated by chunked prefill; after truncation it stays at the head ofwaitingto continue, and middle-chunk samples are dropped. - When decode runs out of free blocks it preempts the
runningtail, using recomputation rather than swap: send back to the head ofwaiting, resetis_prefill=True, free KV. Freed blocks keep their hashes, so re-prefill may hit the prefix cache. postprocessfirst registers hashes for newly filled blocks and updatesnum_cached_tokens, then appends the token; on EOS ormax_tokensit marksFINISHEDand frees KV.
🎯 Self-check checklist
- Can say why nano-vllm schedules Prefill and Decode separately, and what cost prefill-first imposes on requests that are already decoding
- Can explain which phase
max_num_batched_tokensandmax_num_seqsmainly constrain - Can say why chunked prefill only applies to the first sequence of a round, and how a truncated sequence continues next round
- Can explain why
can_allocatereturning-1leads to an immediatebreak, and the head-of-line blocking that causes - Can retell the preemption flow when decode is short of free blocks, including where
while ... elsegoes on “sacrifice yourself” - Can explain why nano-vllm chose recomputation over swap, and how the prefix cache shrinks the recompute amount
- Can name which sample
continuedrops inpostprocess, and the two conditions that finish a sequence - Can contrast nano-vllm and the vLLM V1 scheduler on mixed Prefill/Decode batches
📚 References
- nano-vllm GitHub repository
- Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)
- Orca: A Distributed Serving System for Transformer-Based Generative Models (Yu et al., 2022)
- vLLM docs - Architecture overview
- vLLM V1: A Major Upgrade to vLLM’s Core Architecture (vLLM blog, 2025-01-27)