E3.5 · LESSON 5 OF 6

Copy-on-write

Allocate logical token blocks to physical KV-cache blocks and quantify internal waste and shared-prefix behavior.

THE IDEA

Copy-on-write

Shared blocks are read-only in spirit. But sequences that share a prefix eventually diverge: one appends a token into a block that another sequence also uses. Writing in place would change the other's cache. Copy-on-write fixes this: each block has a reference count; if the block to be written has a count above 1, the writer takes a new block, copies the contents, lowers the old block's count and points its own table at the copy.

#0A ×1

#1A ×1

#2free

#3free

#4free

#5free

#6free

#7free

#8free

#9free

sequence	tokens	block table (logical → physical)
A	20	[0, 1]

free blocks: 8 of 10 in use: 2 (×n = reference count, striped = shared)

A allocated with 20 tokens

Editing a shared block in place would change the other sequence's cached state too. So the writer gets a private copy first; the reference count of the old block drops by one.

Press fork B, then B appends a token. Watch block 1 split into two blocks.

Predict first. Why avoid editing a block used by another sequence in place?

Chapter context

Contiguous reservation can waste memory when sequence lengths differ or grow unpredictably. Paging separates a sequence's logical block order from the physical blocks storing its cache.

TRY IT YOURSELF

Part 1. Do it hands-on

Predict the result of a copy-on-write step, then implement the block manager with reference counts, forking and copy-on-write.

#0A ×1

#1A ×1

#2free

#3free

#4free

#5free

#6free

#7free

#8free

#9free

sequence	tokens	block table (logical → physical)
A	20	[0, 1]

free blocks: 8 of 10 in use: 2 (×n = reference count, striped = shared)

A allocated with 20 tokens

Editing a shared block in place would change the other sequence's cached state too. So the writer gets a private copy first; the reference count of the old block drops by one.

Scenario for the exercise: A has 20 tokens, B is forked from A, then B appends one token. Use the buttons above to see it.

references to block 0 afterwards blocks in use B's table:

To do: Pool of 10 blocks of 16 tokens. Sequence A has 20 tokens (blocks 0 and 1). B is forked from A and shares both blocks. Then B appends token 21, which goes into the half-full block 1. Because block 1 is shared, B gets a private copy. How many sequences now reference block 0?
To do: How many physical blocks are in use now (0, 1 and the new copy)?
To do: Type B's block table now (two physical block numbers; the allocator hands out the lowest free number)

Part 2. Do it in code

The class BlockManager(num_blocks, block_size) keeps: free (sorted list of free block numbers), ref (block → reference count) and seqs (name → {"blocks": [...], "len": n}). Methods for you to write:
_take(): remove and return the LOWEST free block, set its reference count to 1. Raise MemoryError if none is free.
allocate(name, n_tokens): give a new sequence ceil(n_tokens / block_size) blocks (all taken from the free list). Raise MemoryError if there are not enough free blocks, and leave the manager unchanged in that case.
fork(parent, child): the child gets the same block list and length; every block's reference count goes up by 1 (no copying).
append(name): add one token. If the sequence length is a multiple of block_size, take a new block. Otherwise write into the last block: if its reference count is above 1, copy-on-write (take a new block, lower the old block's count by 1, and replace it in this sequence's table). Increase len by 1. Return True if a copy was made, else False.
used(): the number of blocks currently in use.

import math

class BlockManager:
    def __init__(self, num_blocks, block_size):
        self.block_size = block_size
        self.total = num_blocks
        self.free = list(range(num_blocks))
        self.ref = {}
        self.seqs = {}

    def _take(self):
        pass

    def allocate(self, name, n_tokens):
        pass

    def fork(self, parent, child):
        pass

    def append(self, name):
        pass

    def used(self):
        pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Why avoid editing a block used by another sequence in place?

AIt is slower than copying.BIt would change that sequence's cached state too, because both tables point to the same physical block.CBlocks cannot be edited at all.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For length T and block capacity P:

n_{blocks}=\lceil T/P\rceil,\quad slots=P\lceil T/P\rceil,\quad slack=slots-T.

Slack as a fraction of allocated slots is slack/slots. State the denominator when reporting waste percentages.

Worked example

A sequence has 33 tokens and blocks hold 16. It requires ceil(33/16)=3 blocks, allocating 48 token slots. Fifteen slots are unused, so internal slack is 15/48=31.25% of this allocation. It is not 31.25% of all GPU memory, and shared blocks change total ownership accounting.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import ceil,isclose
tokens,capacity=33,16
blocks=ceil(tokens/capacity)
slots=blocks*capacity
assert blocks == 3 and slots == 48
assert slots-tokens == 15
assert isclose((slots-tokens)/slots,.3125)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Paging improves allocation flexibility but does not make memory unlimited. Logical block order differs from physical layout. Internal slack and external fragmentation are different kinds of waste.