Simulate Chunked Prefill Scheduling

The problem, from first principles

A long prompt's prefill run in one shot can take far longer than a single decode step for other concurrently running requests — those requests effectively stall until the huge prefill finishes, spiking their inter-token latency. Simulate chunked prefill scheduling: a long prompt's prefill is split into fixed-size chunks and interleaved with ongoing decode steps for other requests, instead of running the entire prefill in one uninterrupted burst.

From theory to code

budget_left = token_budget - n_active_decode_requests    -- 1 token-unit per active decode request
chunk_size = min(budget_left, tokens_left_in_prefill)      -- spent on the next prefill request in FIFO order

Constraints

• Requests are given as (arrival_step, request_id, prefill_tokens) — after finishing prefill each produces exactly n_decode_tokens_each output tokens (a shared constant across all requests) before completing.
• Every step: 1 token-unit budget goes to each already-decoding request; remaining budget is spent, in FIFO arrival order among prefill-phase requests, up to chunk_size tokens each.
• A request whose prefill finishes this step starts decoding from the NEXT step onward.
• Return, per request, the step its prefill finished and the step it fully completed.

Hints