Build a continuous-batching scheduler with paged KV memory

Driftwood (fictional) serves a chat model and wants a simulator of its scheduler, to try policies before they touch GPUs. Time moves in whole steps. Python 3 and the standard library. No AI tools. Each part should keep the earlier parts' tests passing.

Write simulate(requests, max_running, block_size=None, num_blocks=None, token_budget=None, log=None). It returns {rid: {"first_token": step, "finish": step, "preemptions": n}}.

@dataclass(frozen=True)
class Request:
    rid: str
    arrival: int       # first step at which it may be scheduled
    prompt_len: int    # >= 1
    output_len: int    # >= 1 tokens to generate

Part 1 (15 to 20 minutes): continuous batching. In each step s:

1. Admit waiting requests in order of (arrival, rid) while fewer than max_running are running. A request can be admitted at its arrival step.
2. Every running request does one unit of work. A new request processes its whole prompt and emits its first token in this step. A request that already has tokens emits one more.
3. A request that has emitted output_len tokens finishes at s and leaves after the step.

Part 2 (20 minutes): paged KV memory. Memory is num_blocks blocks of block_size tokens.

• Before a request can emit its next token, the cache must hold prompt_len + emitted tokens, where emitted counts the tokens it has already sent. So in that step it needs ceil((prompt_len + emitted) / block_size) blocks.
• Order inside a step:
  • (a) Grow: each decoding request, oldest first, takes one more block if it needs one. If none is free, preempt the most recently admitted running request (possibly the one asking): free all its blocks and put it at the front of the waiting queue. It keeps the tokens it already sent and rebuilds its cache later.
  • (b) Admit: the head of the queue is admitted only if its blocks fit. If they don't, stop: strict first come, first served, and nobody skips ahead. A preempted request needs blocks for its prompt plus its sent tokens. In the step it comes back, it reprocesses all of them and emits its next token.
  • (c) Run. (d) Finish and free blocks.
• Raise ValueError if a request could never fit, even alone.

Part 3 (15 minutes): chunked prefill. token_budget caps the tokens processed in one step. Decoding requests go first, 1 token each. The rest of the budget goes to prompt processing in admission order. A prompt can be split across steps, and it emits its token in the step its last chunk runs. Raise ValueError if token_budget < max_running.

If log is a list, append one dict per step with the keys step, running, emitted, tokens and blocks_used.

Examples, as (first_token, finish):

• Part 1, max_running = 2: A(arrival 0, prompt 4, output 3), B(0, 2, 1), C(1, 3, 2) give A (0, 2), B (0, 0), C (1, 2).
• Part 2, max_running = 3, block_size = 4, num_blocks = 4: A(0, 6, 4), B(0, 4, 5), C(1, 3, 2) give A (0, 3), B (0, 5) with 1 preemption, and C (4, 5).
• Part 3, max_running = 2, token_budget = 4: A(0, 6, 3), B(0, 3, 2) give A (1, 3), B (2, 3).

Talk only, last 5 minutes: what changes if the scheduler process crashes, or if it runs on 100 GPUs?

Write the order of operations inside one step, and what happens when a running sequence needs a new block and none is free.