TL;DR: A free list, a refcount per physical block and a block table per sequence. A new block is allocated only when a sequence's length crosses a block boundary; a write into a partially filled block whose refcount is above 1 copies it first, and only that first writer copies, because the last holder can then write in place. The scheduler admits by watermark, decodes oldest first, and on exhaustion preempts the newest sequence and re-queues it for recompute.
How to approach it
State the invariants before writing code, because the grader's tests are those invariants: every block is either on the free list with refcount 0 or referenced exactly refcount times across all tables; a table holds exactly blocks; and no write lands in a block another live sequence can read. Ask what the fork is for (parallel sampling, beam search, prefix sharing), and whether preemption should swap to host memory or recompute.
A strong answer
from collections import deque
class OutOfBlocks(Exception):
pass
class BlockManager:
def __init__(self, num_blocks: int, block_size: int):
self.block_size = block_size
self.free = deque(range(num_blocks))
self.ref = [0] * num_blocks
self.tables: dict[int, list[int]] = {}
self.lens: dict[int, int] = {}
def _alloc(self) -> int:
if not self.free:
raise OutOfBlocks
b = self.free.popleft()
self.ref[b] = 1
return b
def _release(self, b: int) -> None:
self.ref[b] -= 1
assert self.ref[b] >= 0, "double free"
if self.ref[b] == 0:
self.free.append(b)
def allocate(self, seq: int, n_tokens: int) -> None:
need = -(-n_tokens // self.block_size) # ceil division
if need > len(self.free):
raise OutOfBlocks # all or nothing
self.tables[seq] = [self._alloc() for _ in range(need)]
self.lens[seq] = n_tokens
def append_slot(self, seq: int):
"""Reserve the slot for one new token: (block, offset, cow_src)."""
table, n = self.tables[seq], self.lens[seq]
offset, cow_src = n % self.block_size, None
if offset == 0: # last block full
table.append(self._alloc())
elif self.ref[table[-1]] > 1: # shared partial block
new = self._alloc() # allocate before releasing
cow_src = table[-1]
self._release(cow_src)
table[-1] = new # caller copies cow_src -> new
self.lens[seq] = n + 1
return table[-1], offset, cow_src
def fork(self, parent: int, child: int) -> None:
for b in self.tables[parent]:
self.ref[b] += 1
self.tables[child] = list(self.tables[parent])
self.lens[child] = self.lens[parent]
def free_seq(self, seq: int) -> None:
for b in self.tables.pop(seq):
self._release(b)
del self.lens[seq]
Three decisions carry the grade. append_slot allocates the new block before releasing the shared one, so an OutOfBlocks raised mid-call leaves every structure untouched and the caller can preempt and retry. Full blocks are shared freely, because no sequence ever writes into a full block again; only the trailing partial block needs copy-on-write. And after a fork, the first sequence to append copies and decrements the refcount to 1, so the second appends in place. Fork sequence 0 at 6 tokens with 4-token blocks: sequence 0 copies block 1 into block 2 and writes there, sequence 1 then owns block 1 alone and writes into it with no copy.
The scheduler step:
class Scheduler:
def __init__(self, bm: BlockManager, watermark_blocks: int = 1):
self.bm, self.wm = bm, watermark_blocks
self.waiting = deque() # (seq, tokens to prefill)
self.running: list[int] = []
def step(self) -> list[int]:
while self.waiting: # FCFS admission
seq, n = self.waiting[0]
need = -(-n // self.bm.block_size)
if len(self.bm.free) - need < self.wm:
break
self.waiting.popleft()
self.bm.allocate(seq, n)
self.running.append(seq)
decoded = []
for seq in list(self.running): # oldest first
while seq in self.running:
try:
self.bm.append_slot(seq)
decoded.append(seq)
break
except OutOfBlocks: # preempt the newest
victim = self.running.pop()
n = self.bm.lens[victim]
self.bm.free_seq(victim)
self.waiting.appendleft((victim, n)) # recompute prompt + output
return decoded
Both pieces were run before publishing: a randomized test of 120,000 allocate, append, fork and free operations checked the three invariants after every operation, and a scheduler simulation of 120 sequences on a 200-block pool ran to completion with 73 preemptions and zero leaked blocks.
The watermark is what stops thrashing. Admitting a prompt that leaves zero free blocks guarantees that the next decode step preempts someone, usually the sequence just admitted, which then re-prefills and preempts again. vLLM's original block manager reserved 1% of the pool this way; a fixed count of a few blocks works as well at small scale.
The reversal condition for recompute: swapping the victim's blocks to host memory and back is cheaper than re-running prefill only when the sequence is long and PCIe bandwidth is good. For short prompts on a fast GPU, recompute wins, and that is why engines default to it.
What interviewers probe next
"Add prefix sharing across unrelated requests." Hash each full block by its tokens and the hash of the block before it, so a block's identity covers its whole prefix. On admission, look up hashes and fork-style increment instead of allocating. Freed blocks with a live hash go to an evictable LRU instead of straight back to the free list.
"What block size would you pick?" Smaller blocks waste less, at most block size minus 1 tokens per sequence, but make tables longer and attention kernel gathers less contiguous. Sixteen tokens is the common default because it keeps waste under one percent at typical lengths without hurting the kernel.
"Beam search with width 4." Each beam forks from its parent at every step. Copy-on-write means beams that keep a shared history share every full block, and only trailing partial blocks get copied.
How it's scored
Weak: reserves a contiguous max_len buffer per sequence, or frees blocks without refcounts so a fork corrupts its sibling.
Hire: correct free list, refcounts and tables; allocates on boundary crossings; copy-on-write on shared partial blocks; FCFS admission with preemption.
Strong Hire: states the invariants up front and tests them, makes append_slot exception-safe, explains the watermark and the recompute-versus-swap decision with numbers, and extends to hash-based prefix sharing.
Common mistakes
- Copying on write for every shared block, including full ones, which turns a fork into a full cache copy.
- Releasing the shared block before allocating its replacement, so an out-of-blocks error leaves the table pointing at a block the sequence no longer owns.
- Preempting the sequence that hit the error rather than the newest one, which starves long-running requests.
- Forgetting that the re-queued length is prompt plus everything generated so far.