Implement one step of a continuous-batching scheduler under a token budget

hardPremium edge casesimplement Scheduler

Before you code

A long prompt arrives while another request is decoding. Predict which work can fit in the next iteration under a fixed token budget.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

Write the two methods that drive a continuous-batching engine: schedule decides which tokens run in the next forward pass, and complete applies the step once it has run.

Req(rid, prompt_len, max_new, prefilled=0, generated=0)
Scheduler(token_budget, max_seqs, kv_tokens)
  waiting  deque of Req, first come first served;  running  list of Req, oldest first
  add(req)          provided; rejects a request that could never fit
  schedule()  -> list of (req, n_tokens) for this step, spending at most token_budget:
                 1. one token for every running request that has finished its prefill
                 2. the rest of any running request's prompt, as much as is left
                 3. admit waiting requests in order, while the budget, max_seqs and the KV
                    reservation (prompt_len + max_new per request) allow; never skip the head
  complete(plan) -> ids of requests that emitted a token in this step
                 a prefill chunk that finishes the prompt emits the first output token;
                 a request that has generated max_new tokens leaves and frees its reservation

The visible tests check the order and the chunk sizes when a long prompt arrives behind 64 running decodes. The full suite adds an in-progress prefill against a new arrival, a head-of-line request that does not fit, the sequence cap, the budget across a random workload, first-token emission, the release of the reservation and an impossible request.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

from collections import deque
from dataclasses import dataclass

@dataclass
class Req:
    rid: int
    prompt_len: int          # >= 1
    max_new: int             # >= 1; the first token comes out of the last prefill chunk
    prefilled: int = 0
    generated: int = 0

class Scheduler:
    def __init__(self, token_budget: int, max_seqs: int, kv_tokens: int):
        self.budget, self.max_seqs = token_budget, max_seqs
        self.kv_total = self.kv_free = kv_tokens  # reserve prompt + max_new on admit
        self.waiting: deque[Req] = deque()
        self.running: list[Req] = []              # admission order, oldest first

def add(self, r: Req) -> None:
        assert r.prompt_len >= 1 and r.max_new >= 1
        if r.prompt_len + r.max_new > self.kv_total:
            raise ValueError("request can never fit")    # else it blocks the queue
        self.waiting.append(r)

def schedule(self) -> list[tuple[Req, int]]:
        raise NotImplementedError

def complete(self, plan) -> list[int]:
        raise NotImplementedError

Visible tests run for everyone; the edge cases and the solution are Premium.

The visible tests

def _with_decoders(ns, budget, n=64):
    s = ns["Scheduler"](budget, 1024, 10**7)
    for i in range(n):
        s.add(ns["Req"](i, 1, 10_000))
    s.complete(s.schedule())               # one-token prompts: all are decoding now
    return s


def test_decodes_run_before_the_prompt(ns):
    s = _with_decoders(ns, 512)
    s.add(ns["Req"](999, 8192, 10))
    plan = s.schedule()
    kinds = [n for _, n in plan]
    assert kinds[:64] == [1] * 64, "the 64 decodes must be scheduled first, one token each"
    assert plan[-1][0].rid == 999 and plan[-1][1] == 448, f"512 - 64 = 448 prompt tokens fit; got {plan[-1][1]}"


def test_a_long_prompt_takes_19_chunks_ending_with_128(ns):
    s = _with_decoders(ns, 512)
    big = ns["Req"](999, 8192, 10)
    s.add(big)
    chunks = []
    for _ in range(100):
        if big.prefilled >= big.prompt_len:
            break
        plan = s.schedule()
        chunks += [n for r, n in plan if r.rid == 999]
        s.complete(plan)
    assert big.prefilled == big.prompt_len, f"after 100 steps the prompt has {big.prefilled} of 8,192 tokens prefilled"
    assert len(chunks) == 19 and chunks[-1] == 128, f"expected 19 chunks ending in 128, got {len(chunks)} ending in {chunks[-1]}"

The reasoning behind this component: the companion page.