Paged-Attention KV Block Allocation

Background

vLLM's paged attention allocates the KV cache in fixed-size blocks — like virtual memory pages — instead of one large contiguous tensor per request. New tokens grab the next block; finished requests release their blocks back to a free list. Fragmentation drops; throughput rises.

Problem statement

Implement BlockAllocator(num_blocks, block_size) with:

• alloc(seq_len) — return a list of $\lceil \text{seq\_len} / \text{block\_size} \rceil$ block IDs from the free list, or None if not enough capacity.
• free(block_ids) — push block IDs back into the free list.
• fragmentation(allocations) — given a list of (block_ids, used_tokens) tuples, return total unused slots = $\sum (\text{capacity} - \text{used})$ for each allocation.

Input

• num_blocks — int >= 1, total pool size.
• block_size — int >= 1, tokens per block.
• alloc(seq_len) — seq_len >= 0.
• free(block_ids) — list of previously-allocated IDs.
• fragmentation(allocations) — list of (list[int], int).

Output

• alloc returns list[int] or None.
• free returns None.
• fragmentation returns int.

Examples

Example 1 — basic allocation

a = BlockAllocator(10, 16)
a.alloc(40) → 3 blocks  (ceil(40 / 16))

Example 2 — free returns blocks to pool

a = BlockAllocator(3, 16)
b1, b2, b3 = a.alloc(16), a.alloc(16), a.alloc(16)
a.alloc(16) → None       (no capacity)
a.free(b1)
a.alloc(16) → not None   (capacity restored)

Example 3 — fragmentation

a = BlockAllocator(10, 16)
blocks = a.alloc(17)
a.fragmentation([(blocks, 17)]) → 15   (2 blocks × 16 - 17 used)

Constraints

• alloc is atomic: either grant all needed blocks or return None (no partial allocation).
• free adds the released IDs back to the free list.
• fragmentation is informational only; doesn't change state.

Notes

• Why blocks, not contiguous. Contiguous KV requires the largest free chunk to be ≥ seq_len; under churn this fragments badly. Block allocation packs perfectly.
• Prefix sharing. Real vLLM also has copy-on-write: requests sharing a prefix share blocks until one writes a new token. This kernel doesn't model that.