Continuous Batching and Preemption
Hard · Solved in PyTorch · Machine learning coding practice
Problem
Implement one step of a simplified nano-vLLM scheduler. The cache manager has already computed each request's block requirements; there is no GPU execution or prefix hashing here.
waiting is an int64 tensor (W, 3) with rows [id, remaining_prompt_tokens, new_blocks_needed] in FIFO order. running is (R, 3) with [id, new_blocks_needed_for_decode, blocks_already_held]. IDs are unique across both queues. Counts are nonnegative, remaining prompt tokens and held blocks are positive, and decode needs are 0 or 1. Each queue has at most 16 entries. Do not mutate either input.
Topics: inference, vllm, pytorch
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Continuous Batching and Preemption in the browser and run it against the tests.