Schedule Decode and Prefill Within a Token Budget
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Plan one decode-first mixed batch. This exercise isolates V1-style token budgeting: all requests are ready, KV space is available, and speculative decoding is disabled.
The int64 running: (R,) lists decode request IDs in priority order, each needing one token. The int64 waiting: (W, 2) lists [id, remaining_uncached_prompt_tokens] in FIFO order, with positive remainders. IDs are unique across both queues. First admit running requests in order, one token each. Then admit waiting requests in order, assigning each min(remainder, chunk_size, remaining_budget) tokens. Visit each request at most once; do not revisit a prefill to spend leftover budget.
Topics: inference, vllm, pytorch
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Schedule Decode and Prefill Within a Token Budget in the browser and run it against the tests.