TL;DR: Each step spends a fixed token budget in a fixed order: one token for every running decode, then the remainder of any partially prefilled prompt, then new requests admitted first come, first served, each taking at most what is left. On a TP=4 FP8 70B replica, a 2,048-token budget bounds a mixed step at about 36.5 ms of compute at peak FLOP/s, while an 8,192 budget lets one step reach 146 ms, which every running decode feels as TPOT.
How to approach it
Agree on the contract first: what one step is (one forward pass over a flat batch of tokens from many requests), what the budget counts (tokens in that pass, prefill and decode alike), and what admission needs (a sequence slot and KV memory). Then state the invariants you will test: the scheduled tokens never exceed the budget, every prompt is prefilled exactly once across its chunks, a request ready to decode gets a token every step while decodes fit the budget, and admission is in arrival order. Say early that this version reserves KV for the whole request at admission and never preempts; the preempting variant is a separate design.
A strong answer
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]]:
left, plan = self.budget, []
for r in self.running: # 1. decodes first: they hold TPOT
if r.prefilled == r.prompt_len and left > 0:
plan.append((r, 1)); left -= 1
for r in self.running: # 2. finish partial prefills
if r.prefilled < r.prompt_len and left > 0:
n = min(r.prompt_len - r.prefilled, left)
plan.append((r, n)); left -= n
while self.waiting and left > 0 and len(self.running) < self.max_seqs:
r = self.waiting[0] # 3. admit FCFS, no skipping ahead
need = r.prompt_len + r.max_new
if need > self.kv_free:
break
self.waiting.popleft(); self.kv_free -= need
self.running.append(r)
n = min(r.prompt_len, left)
plan.append((r, n)); left -= n
return plan
def complete(self, plan) -> list[int]:
"""Apply one executed step; return ids of requests that emitted a token."""
emitted = []
for r, n in plan:
if r.prefilled < r.prompt_len:
r.prefilled += n
if r.prefilled == r.prompt_len: # last chunk samples token 1
r.generated = 1; emitted.append(r.rid)
else:
r.generated += 1; emitted.append(r.rid)
for r in [r for r in self.running if r.generated == r.max_new]:
self.running.remove(r); self.kv_free += r.prompt_len + r.max_new
return emitted
The order of the three loops is the design. Decodes go first because each one is a user waiting on its next token and costs one token of budget; putting prefill first means a single 30,000-token prompt can take the whole step and stall every stream. Partial prefills go second so that a prompt started last step finishes before newer arrivals start theirs, which keeps time to first token roughly first come, first served. Admission goes last and stops at the first request that does not fit, instead of skipping ahead to a smaller one, because skipping lets a stream of short prompts starve a long one indefinitely.
The budget is the TPOT knob, and the chain is short:
Llama 3.3 70B, TP=4 H100 SXM, FP8 weights
prefill FLOPs per token = 2 x 70.55e9 = 1.41e11
peak = 4 x 1,979 TFLOP/s dense FP8 = 7.92e15 FLOP/s
budget 512: 512 x 1.41e11 / 7.92e15 = 9.1 ms
budget 2048: 36.5 ms
budget 8192: 146.0 ms
decode-only step floor (weights, FP8) = 70.55e9 B / 4 / 3.35e12 B/s = 5.3 ms
Those are floors at 100% of peak; real steps are slower. The point stands at any efficiency: a decode-only step costs about 5 ms, and a step carrying a full 8,192-token chunk costs 28 times that, so the budget is the worst TPOT a running stream sees while prompts are arriving. A smaller budget raises TTFT for long prompts instead: with 64 decodes running, a 2,048 budget leaves 1,984 tokens for prefill and a 32,768-token prompt needs 17 steps to reach its first token, against 5 steps at a budget of 8,192.
The last prefill chunk emits the first token, which is why complete sets generated = 1 when the prompt finishes rather than waiting one more step. A request with max_new = 1 leaves on the same step its prompt finishes.
What was run: 400 randomized simulations, 7,283 requests over 366,724 steps, with budgets from 1 to 256 tokens (budget 1 forces every prompt into one-token chunks and lets decodes crowd prefill out entirely), 1 to 11 sequence slots, prompts of 1 to 299 tokens, staggered arrivals and KV capacities from barely enough for the largest request to generous. After every step the test asserted the budget, the slot limit and nonnegative free KV; that the decodes scheduled were exactly the oldest ready ones up to the budget; and at the end that every prompt received exactly prompt_len prefill tokens and max_new - 1 decode steps, admission order matched arrival order and all KV was returned. A hand trace (budget 8, one decode running, a 20-token prompt arriving) produced chunks of 7, 7 and 6 tokens with the first token on the third step, and an empty scheduler returns an empty plan. A request larger than total KV capacity raises at add because without that check it would sit at the head of the queue forever and block everything behind it.
The reversal condition: reserving prompt_len + max_new at admission wastes memory when most requests stop well before max_new. If measured output lengths are a small fraction of the cap, switch to allocating KV as tokens arrive, admit by a watermark, and accept preemption, which turns this into a paged design with an eviction policy.
What interviewers probe next
"Decodes alone exceed the budget." Then some ready streams skip a step, and the oldest-first order means the newest ones skip. Either raise the budget above max_seqs so it cannot happen, or rotate the start index each step to spread the stall.
"Make prefill chunks a multiple of the KV block size." Round each chunk down to the block size unless it is the prompt's final chunk, which keeps KV writes block aligned at the cost of a few budget tokens per step.
"Add priorities." Keep the loop order but sort each class by priority then arrival, and add an age term so low priority requests still make progress.
How it's scored
Weak: schedules a whole prompt or nothing, so one long prompt stalls every decode, or lets the scheduled tokens exceed the budget.
Hire: decode-first ordering, chunks capped by the remaining budget, FCFS admission with slot and memory checks, and the first token emitted from the last chunk.
Strong Hire: converts the budget into milliseconds of TPOT and steps to first token, tests the invariants across randomized arrivals, rejects requests that can never fit, and states when reservation should give way to paged allocation with preemption.
Common mistakes
- Emitting the first token one step after the prompt completes, which adds a full step to every TTFT.
- Letting admission skip a blocked request for a smaller one behind it, which starves long prompts under load.
- Counting a decode as zero tokens of budget, so a step with 256 decodes and a full chunk overshoots by 256.
- Freeing the KV reservation when the prompt finishes instead of when the request finishes.