Copy-on-write KV sharing

When two sequences start from the same prompt, their KV for that prompt is identical. Parallel sampling (asking for n answers to one prompt) and beam search both create this situation. With copy-on-write, the sequences point at the same physical blocks for the shared part. A block is copied only when a sequence needs to change it. The shared prompt is stored once.

The big idea

Think of a shared document with a copy for each reviewer. The reviewers share the same text until one of them edits a paragraph. Only that paragraph is split off. Everything else stays shared.

The picture breaks in one place. Copying a paragraph is cheap. Copying a KV block is a real cost on the GPU, and the reference counts must be exact. If they are wrong, a block can be freed while another sequence still reads it.

Walking through n = 4 samples

Take a 2,000-token prompt, 300 generated tokens per sample, and 16-token blocks. Numbers are illustrative.

Prompt blocks:        2,000 / 16 = 125 blocks
Generated blocks:     300 / 16 = 19 blocks per sample (rounded up)

Without sharing:      4 × (125 + 19) = 576 blocks
With sharing:         125 (shared, refcount 4) + 4 × 19 = 201 blocks

The prompt is about 4 times cheaper to keep, because it is stored once. Each sample still gets its own tail, since its generated tokens differ.

Here is how the blocks connect. Two sequences share their system prompt blocks, and each then has its own user blocks.

What copy-on-write does

1. Sequences start with the same block IDs in their tables, and each block's count goes up.
2. When a sequence needs to write into a shared block, the engine first copies that block to a new one and points the sequence at the copy.
3. The other sharers keep the original, and its count goes down.

Beam search depends on this. Each beam extends a shared history, and most history is never copied.

Common mistakes

• Turning on n = 8 without sharing. KV for the prompt is then stored eight times.
• Evicting shared blocks first. A block with a high count hurts every sequence that uses it. Preempt unique tails before shared prefixes.
• Assuming sharing is automatic. Confirm that your engine shares prompt blocks for parallel sampling and beam search.

Check yourself

1. Why does n = 4 sampling use less memory with shared prompt blocks? Answer: the prompt blocks are stored once and referenced four times, instead of stored four times.
2. Why should the scheduler evict a sequence's unique tail before a shared block? Answer: evicting a shared block hurts every sequence that uses it.

Key takeaways

• Sequences that share a prompt can share its KV blocks, using reference counts.
• Copy-on-write copies a block only at the point where a sequence diverges.
• Parallel sampling and beam search are the main beneficiaries.
• Preemption and prefix caching both depend on accurate reference counts.

Further learning

Next, read preemption: swap or recompute for what the scheduler does when memory runs out. For the same idea applied across requests, see prefix caching.