Skip to main content
Inference optimization

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

nano-vllm Scheduler Continuous batching Scheduling Preemption Source reading

📑 Contents


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:

Sequence state flow: a new request enters the waiting queue; an unfinished chunked prefill stays at the waiting head; after the last prefill chunk is scheduled it enters the running queue; in running it decodes step by step, finishes and frees KV on EOS or max_tokens; when GPU memory is short it is preempted back to waiting

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_allocate returns -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 uses min(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 continue drops it. Only after the last chunk finishes is the sampled token the real first generated token.
  • Finish conditions: hit EOS (and ignore_eos is not set), or the generated count reaches max_tokens. Either one marks FINISHED, frees KV blocks, and removes the sequence from running; unfinished sequences stay in running for the next decode round. A request with max_tokens = 1 finishes on this prefill step; by then schedule has already moved it into running, so remove does not fail.
  • ignore_eos / max_tokens are instance attributes on Sequence (flattened from SamplingParams at 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:

Scheduling loop: schedule() first checks whether the prefill branch picked sequences; if yes, run Prefill under the Token Budget (first sequence may be chunked); if no (waiting empty or head does not fit), walk running for Decode; when can_append fails, preempt the running tail back to waiting for recomputation; when it succeeds, may_append and put into the batch; then model forward plus sampling; postprocess registers hashes, writes back, decides finish, frees KV, and the next round begins

📝 Summary

  • The scheduler only decides; it does not run matrix math: it keeps waiting / running deques, and schedule() returns (seqs, is_prefill).
  • Prefill first, no mixed batches: if the prefill branch picks sequences, do prefill; otherwise decode. When the head of waiting does 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 of waiting to continue, and middle-chunk samples are dropped.
  • When decode runs out of free blocks it preempts the running tail, using recomputation rather than swap: send back to the head of waiting, reset is_prefill=True, free KV. Freed blocks keep their hashes, so re-prefill may hit the prefix cache.
  • postprocess first registers hashes for newly filled blocks and updates num_cached_tokens, then appends the token; on EOS or max_tokens it marks FINISHED and 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_tokens and max_num_seqs mainly 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_allocate returning -1 leads to an immediate break, and the head-of-line blocking that causes
  • Can retell the preemption flow when decode is short of free blocks, including where while ... else goes 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 continue drops in postprocess, and the two conditions that finish a sequence
  • Can contrast nano-vllm and the vLLM V1 scheduler on mixed Prefill/Decode batches

📚 References