Skip to content

Implement a paged KV block manager with reference counts and copy-on-write

hardPremium edge casesimplement BlockManager

Before you code

Two requests share a partially filled KV block. Predict which state must change before one request can append independently.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

Write the block manager a paged KV cache runs on. Each sequence owns a block table of physical block ids; a block holds block_size tokens; forked sequences share blocks until one writes into a block that is still shared.

class OutOfBlocks(Exception)
class BlockManager(num_blocks, block_size)
  free     deque of free block ids, initially 0 .. num_blocks - 1 in order
  ref      list, reference count per block id
  tables   dict seq -> list of block ids
  lens     dict seq -> tokens written
  allocate(seq, n_tokens)  give a new sequence ceil(n / block_size) blocks, all or nothing
  append_slot(seq)         reserve the slot for one more token; returns (block, offset, cow_src)
                           cow_src is the block to copy from when a copy-on-write happened, else None
  fork(parent, child)      the child shares every parent block
  free_seq(seq)            drop the sequence; a block returns to free when nothing references it

append_slot allocates a fresh block when the last one is full. If the last block is partly full and shared, it copies on write: the writer gets a new block and the old one loses a reference. Any operation that runs out of blocks raises OutOfBlocks and leaves every structure unchanged.

The visible tests cover allocation, a block boundary and a fork. The full suite adds copy-on-write, a fork at a full block, exhaustion during allocate, append and copy-on-write, freeing a forked parent, and a random sequence of operations checked against the pool's invariants.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

Visible tests run for everyone; the edge cases and the solution are Premium.

The visible tests
def test_allocation_rounds_up_and_appends_at_the_boundary(ns):
    bm = ns["BlockManager"](8, 16)
    bm.allocate(1, 40)
    assert len(bm.tables[1]) == 3, f"40 tokens need 3 blocks of 16, got {len(bm.tables[1])}"
    for _ in range(8):                      # tokens 41..48 fill block 3
        bm.append_slot(1)
    assert len(bm.tables[1]) == 3, "48 tokens still fit in 3 blocks"
    block, offset, cow = bm.append_slot(1)  # token 49 crosses the boundary
    assert len(bm.tables[1]) == 4 and offset == 0 and cow is None, "token 49 needs a fourth block at offset 0"
    assert len(bm.free) == 4, f"8 blocks minus 4 held = 4 free, got {len(bm.free)}"


def test_fork_shares_blocks_instead_of_copying(ns):
    bm = ns["BlockManager"](8, 16)
    bm.allocate(1, 40)
    bm.fork(1, 2)
    assert bm.tables[2] == bm.tables[1], "the child must point at the parent's blocks"
    assert all(bm.ref[b] == 2 for b in bm.tables[1]), "every shared block has two references"
    assert len(bm.free) == 5, f"a fork allocates nothing; expected 5 free, got {len(bm.free)}"

The reasoning behind this component: the companion page.