Pack Prefill and Decode Batches

Medium · Solved in PyTorch · Machine learning coding practice

Problem

Build a flat model input batch from scheduled slices. tokens is a list of B nonempty token-ID lists. The int64 tensors starts and counts have shape (B,); sequence i contributes tokens[i][starts[i]:starts[i]+counts[i]]. This covers both prompt chunks and one-token decode. Keep sequence order, including zero-length slices.

block_tables is int64 (B, L). Logical position p maps to physical slot block_tables[i, p // block_size] * block_size + p % block_size. Unused table entries are -1; every selected position has a valid mapping. Starts are absolute positions, not positions relative to the chunk.

Topics: inference, vllm, pytorch

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Pack Prefill and Decode Batches in the browser and run it against the tests.

Related problems

• Mask Logits with a Grammar State Easy
• Continuous Batching and Preemption Hard
• Measure Token Latency and SLO Goodput Medium
• Route Requests by Replica Load Easy
• Schedule Decode and Prefill Within a Token Budget Medium
• Coordinate Prefill-to-Decode KV Handoff Medium