Why AI servers waste memory, and how paging fixes it

"""The naive layout: one contiguous slab of slots per request, sized for the maximum length."""


class ContiguousPool:
    """A row of `n_slots` token slots. Each request gets one slab, first fit."""

    def __init__(self, n_slots: int):
        self.n_slots = n_slots
        self.slabs = {}   # request id -> (start, size)
        self.used = {}    # request id -> tokens actually stored

    def allocate(self, rid, size: int):
        """Reserve `size` adjacent slots. Return the start slot, or None if no gap is big enough."""
        for start, length in self.free_runs():
            if length >= size:
                self.slabs[rid] = (start, size)
                self.used[rid] = 0
                return start
        return None

    def release(self, rid) -> None:
        del self.slabs[rid]
        del self.used[rid]

    def free_runs(self):
        """Gaps between slabs as (start, length), left to right."""
        runs, cursor = [], 0
        for start, size in sorted(self.slabs.values()):
            if start > cursor:
                runs.append((cursor, start - cursor))
            cursor = start + size
        if cursor < self.n_slots:
            runs.append((cursor, self.n_slots - cursor))
        return runs

    def free_total(self) -> int:
        return sum(length for _, length in self.free_runs())

    def largest_free_run(self) -> int:
        return max((length for _, length in self.free_runs()), default=0)

    def internal_waste(self) -> int:
        """Reserved slots that hold no token."""
        return sum(size - self.used[rid] for rid, (_, size) in self.slabs.items())


def requests_that_fit_contiguous(pool_tokens: int, max_len: int) -> int:
    """Each request reserves max_len slots, whatever its real length turns out to be."""
    return pool_tokens // max_len


def requests_that_fit_paged(lengths, pool_tokens: int, block_size: int) -> dict:
    """Admit requests in order while their blocks fit. Returns counts and slot accounting."""
    free_blocks = pool_tokens // block_size
    admitted = used = 0
    for n in map(int, lengths):
        need = -(-n // block_size)            # ceil(n / block_size)
        if need > free_blocks:
            break
        free_blocks -= need
        admitted += 1
        used += n
    allocated = (pool_tokens // block_size - free_blocks) * block_size
    return {"admitted": admitted, "used": used, "allocated": allocated, "waste": allocated - used}

"""How big the KV cache is, and what it saves.

Standard library plus numpy. Run `python3 show_numbers.py` to print every
number the lesson quotes.
"""
from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class ModelConfig:
    name: str
    n_layers: int
    n_kv_heads: int      # key/value heads, not query heads
    head_dim: int
    bytes_per_elem: int  # 2 for 16-bit floats (fp16 or bf16)
    max_tokens: int      # longest sequence the model accepts


# Llama 3 8B: 32 layers, 8 key/value heads (grouped-query attention),
# hidden size 4096 / 32 query heads = head_dim 128, bf16, 8,192-token context.
LLAMA3_8B = ModelConfig("Llama 3 8B", 32, 8, 128, 2, 8192)
# OPT-13B: 40 layers, 40 heads of size 128 (hidden size 5120), fp16, 2,048 tokens.
OPT_13B = ModelConfig("OPT-13B", 40, 40, 128, 2, 2048)


def kv_bytes_per_token(cfg: ModelConfig) -> int:
    """2 (one key, one value) x layers x kv_heads x head_dim x bytes."""
    return 2 * cfg.n_layers * cfg.n_kv_heads * cfg.head_dim * cfg.bytes_per_elem


def kv_bytes(cfg: ModelConfig, tokens: int) -> int:
    return kv_bytes_per_token(cfg) * tokens


def pairs_without_cache(n_steps: int) -> int:
    """Step t must have keys and values for all t positions: 1 + 2 + ... + n."""
    return n_steps * (n_steps + 1) // 2


def pairs_with_cache(n_steps: int) -> int:
    """Each step computes the key and value of the one new position."""
    return n_steps


def _softmax(s):
    e = np.exp(s - s.max())
    return e / e.sum()


def decode_both_ways(n_steps: int, d: int = 8, seed: int = 0):
    """One attention head generating n_steps positions, with and without a cache.

    Returns (outputs_without, outputs_with, pairs_without, pairs_with).
    A "pair" is one key vector plus one value vector computed for one position.
    The input vectors are fixed random numbers; a real model feeds each output
    back in, but the counting is the same.
    """
    rng = np.random.default_rng(seed)
    w_q, w_k, w_v = rng.standard_normal((3, d, d)) / np.sqrt(d)
    x = rng.standard_normal((n_steps, d))

    out_a, pairs_a = [], 0
    for t in range(1, n_steps + 1):
        keys, values = x[:t] @ w_k, x[:t] @ w_v       # recompute the whole past
        pairs_a += t
        q = x[t - 1] @ w_q
        out_a.append(_softmax(keys @ q / np.sqrt(d)) @ values)

    out_b, pairs_b = [], 0
    k_cache, v_cache = [], []
    for t in range(1, n_steps + 1):
        k_cache.append(x[t - 1] @ w_k)                 # only the new position
        v_cache.append(x[t - 1] @ w_v)
        pairs_b += 1
        q = x[t - 1] @ w_q
        keys, values = np.array(k_cache), np.array(v_cache)
        out_b.append(_softmax(keys @ q / np.sqrt(d)) @ values)
    return np.array(out_a), np.array(out_b), pairs_a, pairs_b

"""A paged KV cache for one attention head: block pool, free list, block tables, copy-on-write."""
from collections import deque

import numpy as np


class OutOfBlocks(Exception):
    """A block is needed and the free list is empty."""


class BlockManager:
    def __init__(self, num_blocks: int, block_size: int, head_dim: int = 1):
        self.block_size = block_size
        # The pool: all physical memory, cut into fixed-size blocks, allocated once.
        self.k_pool = np.zeros((num_blocks, block_size, head_dim))
        self.v_pool = np.zeros((num_blocks, block_size, head_dim))
        self.free_list = deque(range(num_blocks))   # take from the left, give back on the right
        self.ref_count = [0] * num_blocks           # how many block tables point at each block
        self.tables = {}                            # request id -> physical block ids, in logical order
        self.lengths = {}                           # request id -> tokens stored

    def add_request(self, rid) -> None:
        self.tables[rid] = []
        self.lengths[rid] = 0                       # allocates nothing

    def _take_block(self) -> int:
        block = self.free_list.popleft()
        self.ref_count[block] = 1
        return block

    def append(self, rid, k: np.ndarray, v: np.ndarray) -> None:
        """Store the key and value of one new token at the end of request `rid`."""
        table, n = self.tables[rid], self.lengths[rid]
        offset = n % self.block_size
        last_is_full = offset == 0                  # also true when the table is empty
        last_is_shared = not last_is_full and self.ref_count[table[-1]] > 1
        if (last_is_full or last_is_shared) and not self.free_list:
            raise OutOfBlocks(rid)                  # nothing has changed yet
        if last_is_full:
            table.append(self._take_block())
        elif last_is_shared:                        # copy-on-write: give this request its own copy
            old, new = table[-1], self._take_block()
            self.k_pool[new], self.v_pool[new] = self.k_pool[old], self.v_pool[old]
            self.ref_count[old] -= 1
            table[-1] = new
        self.k_pool[table[-1], offset] = k
        self.v_pool[table[-1], offset] = v
        self.lengths[rid] = n + 1

    def fork(self, parent, child) -> None:
        """The child starts as the parent's twin. No data is copied: both tables name the same blocks."""
        self.tables[child] = list(self.tables[parent])
        self.lengths[child] = self.lengths[parent]
        for block in self.tables[child]:
            self.ref_count[block] += 1

    def free(self, rid) -> None:
        """Finish a request. A block returns to the free list when no table points at it any more."""
        for block in self.tables.pop(rid):
            self.ref_count[block] -= 1
            if self.ref_count[block] == 0:
                self.free_list.append(block)
        del self.lengths[rid]

    def locate(self, rid, position: int):
        """Logical token position -> (physical block, slot inside the block)."""
        return self.tables[rid][position // self.block_size], position % self.block_size

    def gather(self, rid):
        """All keys and values of `rid` in logical order, shape [tokens, head_dim]."""
        n, head_dim = self.lengths[rid], self.k_pool.shape[2]
        keys, values = np.zeros((n, head_dim)), np.zeros((n, head_dim))
        for pos in range(n):
            block, slot = self.locate(rid, pos)
            keys[pos], values[pos] = self.k_pool[block, slot], self.v_pool[block, slot]
        return keys, values


def paged_attention(q: np.ndarray, mgr: BlockManager, rid) -> np.ndarray:
    """One query vector attends over request `rid`, reading the cache one block at a time."""
    n, size, table = mgr.lengths[rid], mgr.block_size, mgr.tables[rid]
    scores = []
    for j, block in enumerate(table):
        valid = min(size, n - j * size)             # the last block may be partly filled
        scores.append(mgr.k_pool[block, :valid] @ q)
    scores = np.concatenate(scores) / np.sqrt(len(q))
    weights = np.exp(scores - scores.max())
    weights /= weights.sum()
    out = np.zeros_like(q, dtype=float)
    for j, block in enumerate(table):
        valid = min(size, n - j * size)
        out += weights[j * size : j * size + valid] @ mgr.v_pool[block, :valid]
    return out


def attention_contiguous(q: np.ndarray, keys: np.ndarray, values: np.ndarray) -> np.ndarray:
    """Reference: the same attention over keys and values already laid out in one array."""
    scores = keys @ q / np.sqrt(len(q))
    weights = np.exp(scores - scores.max())
    return (weights / weights.sum()) @ values


def worked_example():
    """The lesson's hand-worked example: block size 4, a pool of 8 blocks, requests A, B and C.

    Returns a list of (event, tables, free_list, lengths), one per step.
    """
    mgr = BlockManager(num_blocks=8, block_size=4)
    steps = []

    def grow(rid, tokens):
        for _ in range(tokens):
            n = mgr.lengths[rid]
            mgr.append(rid, np.array([n]), np.array([n]))

    def snap(event):
        steps.append((event, {r: list(t) for r, t in mgr.tables.items()}, list(mgr.free_list),
                      dict(mgr.lengths)))

    snap("start")
    for rid, prompt in (("A", 6), ("B", 3), ("C", 5)):
        mgr.add_request(rid)
        grow(rid, prompt)
        snap(f"{rid} arrives with a {prompt}-token prompt")
    grow("A", 2)
    snap("A generates 2 tokens")
    grow("A", 1)
    snap("A generates 1 more token")
    grow("B", 2)
    snap("B generates 2 tokens")
    mgr.free("B")
    snap("B finishes")
    grow("C", 4)
    snap("C generates 4 tokens")
    return steps


def prefix_sharing_example():
    """Block size 4, a 6-token prompt shared by two samples. Returns (event, tables, ref_counts, free_list)."""
    mgr = BlockManager(num_blocks=6, block_size=4)
    steps = []

    def snap(event):
        steps.append((event, {r: list(t) for r, t in mgr.tables.items()}, list(mgr.ref_count),
                      list(mgr.free_list)))

    def one(rid):
        n = mgr.lengths[rid]
        mgr.append(rid, np.array([n]), np.array([n]))

    mgr.add_request("S1")
    for _ in range(6):
        one("S1")
    snap("S1 holds the 6-token prompt")
    mgr.fork("S1", "S2")
    snap("S2 forks from S1")
    one("S1")
    snap("S1 writes token 6")
    one("S2")
    snap("S2 writes token 6")
    return steps

"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import numpy as np

from contiguous import ContiguousPool, requests_that_fit_contiguous, requests_that_fit_paged
from kv_sizes import (LLAMA3_8B, OPT_13B, decode_both_ways, kv_bytes, kv_bytes_per_token,
                      pairs_with_cache, pairs_without_cache)
from paged_kv import prefix_sharing_example, worked_example

GIB = 1024**3


def main():
    print("== steps with and without a cache")
    for n in (10, 100, 1000):
        print(f"n={n}: without={pairs_without_cache(n)}  with={pairs_with_cache(n)}  "
              f"ratio={pairs_without_cache(n) / pairs_with_cache(n):.1f}")
    out_a, out_b, pa, pb = decode_both_ways(1000)
    print(f"demo n=1000: counted without={pa} with={pb}  outputs equal={np.allclose(out_a, out_b)}")

    print("== sizes")
    for cfg in (LLAMA3_8B, OPT_13B):
        per = kv_bytes_per_token(cfg)
        print(f"{cfg.name}: {per} bytes/token = {per / 1024:.0f} KiB; "
              f"full context ({cfg.max_tokens}) = {kv_bytes(cfg, cfg.max_tokens)} bytes = "
              f"{kv_bytes(cfg, cfg.max_tokens) / GIB:.4f} GiB")
    print("Llama 3 8B, 1000 tokens:", kv_bytes(LLAMA3_8B, 1000) / 1024**2, "MiB")
    print("Llama 3 8B, 131072 tokens:", kv_bytes(LLAMA3_8B, 131072) / GIB, "GiB")

    print("== contiguous slabs (figure 1): 32 slots")
    pool = ContiguousPool(32)
    pool.allocate("R1", 12); pool.allocate("R2", 8); pool.allocate("R3", 8)
    pool.used.update(R1=5, R2=3, R3=6)
    print("before: internal waste", pool.internal_waste(), "runs", pool.free_runs())
    pool.release("R2")
    print("after R2: runs", pool.free_runs(), "free_total", pool.free_total(),
          "largest", pool.largest_free_run(), "internal waste", pool.internal_waste(),
          "D(12) ->", pool.allocate("D", 12))
    eq = ContiguousPool(32)
    print("12 reserved each:", [eq.allocate(r, 12) for r in "ABC"])

    print("== capacity: 16 GiB of cache, Llama 3 8B, max_len 8192, block size 16")
    pool_tokens = 16 * GIB // kv_bytes_per_token(LLAMA3_8B)
    lengths = np.random.default_rng(0).integers(100, 2001, size=500)
    fit_c = requests_that_fit_contiguous(pool_tokens, 8192)
    paged = requests_that_fit_paged(lengths, pool_tokens, 16)
    print("pool tokens", pool_tokens, "mean length", lengths.mean())
    print("contiguous admits", fit_c, "-> utilisation of reserved slots",
          f"{lengths[:fit_c].sum() / (fit_c * 8192):.3f}")
    print("paged", paged, f"waste fraction of allocated {paged['waste'] / paged['allocated']:.4f}",
          f"gain {paged['admitted'] / fit_c:.1f}x")

    print("== worked example")
    for event, tables, free, lengths_ in worked_example():
        used = sum(lengths_.values())
        allocated = 4 * sum(len(t) for t in tables.values())
        print(event, tables, free, lengths_, f"used={used} allocated={allocated} wasted={allocated - used}")
    print("== prefix sharing")
    for row in prefix_sharing_example():
        print(*row)


if __name__ == "__main__":
    main()

import numpy as np

from contiguous import ContiguousPool, requests_that_fit_contiguous, requests_that_fit_paged


def test_free_memory_in_two_gaps_cannot_hold_one_big_request():
    """External fragmentation: 12 slots are free, but never 12 in a row."""
    pool = ContiguousPool(32)
    for rid, size in (("R1", 12), ("R2", 8), ("R3", 8)):
        pool.allocate(rid, size)
    pool.used.update(R1=5, R2=3, R3=6)
    pool.release("R2")
    assert pool.free_runs() == [(12, 8), (28, 4)]
    assert pool.free_total() == 12 and pool.largest_free_run() == 8
    assert pool.allocate("D", 12) is None
    assert pool.internal_waste() == (12 - 5) + (8 - 6)     # reserved slots that hold no token


def test_equal_reservations_turn_the_third_request_away():
    """Reserving 12 slots each: the lesson's 32-slot pool fits two requests, the paged pool fits all three."""
    pool = ContiguousPool(32)
    assert [pool.allocate(r, 12) for r in "ABC"] == [0, 12, None]


def test_paged_waste_is_under_one_block_per_request():
    lengths = np.random.default_rng(1).integers(1, 500, size=300)
    result = requests_that_fit_paged(lengths, pool_tokens=20000, block_size=16)
    assert result["admitted"] > requests_that_fit_contiguous(20000, 2048)
    assert 0 <= result["waste"] < result["admitted"] * 16     # rounding down would make it negative


def test_lesson_capacity_example_matches_the_text():
    """16 GiB of Llama 3 8B cache is 131,072 slots; seed 0 draws 500 lengths from 100 to 2,000."""
    pool_tokens = 131072
    lengths = np.random.default_rng(0).integers(100, 2001, size=500)
    fit_contiguous = requests_that_fit_contiguous(pool_tokens, 8192)
    paged = requests_that_fit_paged(lengths, pool_tokens, block_size=16)
    assert fit_contiguous == 16
    assert round(lengths[:fit_contiguous].sum() / (fit_contiguous * 8192), 3) == 0.129
    assert (paged["admitted"], paged["allocated"], paged["waste"]) == (122, 130800, 913)

import numpy as np

from kv_sizes import (LLAMA3_8B, OPT_13B, decode_both_ways, kv_bytes, kv_bytes_per_token,
                      pairs_with_cache, pairs_without_cache)


def test_formula_matches_public_configurations():
    # Llama 3 8B: 2 x 32 layers x 8 kv heads x 128 x 2 bytes = 128 KiB per token
    assert kv_bytes_per_token(LLAMA3_8B) == 131072
    assert kv_bytes(LLAMA3_8B, 8192) == 1024**3          # exactly 1 GiB at the full 8,192-token context
    # The PagedAttention paper computes OPT-13B as 2 x 5120 x 40 x 2 bytes = 800 KB per token
    assert kv_bytes_per_token(OPT_13B) == 2 * 5120 * 40 * 2 == 819200


def test_cache_changes_the_cost_not_the_answer():
    """A cache that reads the wrong rows or is off by one would change the outputs."""
    n = 60
    without, with_cache, pairs_a, pairs_b = decode_both_ways(n)
    assert np.allclose(without, with_cache)
    assert (pairs_a, pairs_b) == (pairs_without_cache(n), pairs_with_cache(n)) == (n * (n + 1) // 2, n)

import math
import random

import numpy as np
import pytest

from paged_kv import (BlockManager, OutOfBlocks, attention_contiguous, paged_attention,
                      prefix_sharing_example, worked_example)

D = 3


def vec(rng):
    return rng.standard_normal(D)


class Reference:
    """No blocks: one Python list of (key, value) per request."""

    def __init__(self):
        self.seqs = {}


def check(mgr, ref, rid):
    keys, values = mgr.gather(rid)
    want_k = np.array([k for k, _ in ref.seqs[rid]]).reshape(-1, D)
    want_v = np.array([v for _, v in ref.seqs[rid]]).reshape(-1, D)
    assert np.allclose(keys, want_k) and np.allclose(values, want_v), f"request {rid} reads wrong data"


def test_blocks_in_use_is_ceil_of_length_over_block_size():
    rng = np.random.default_rng(0)
    mgr = BlockManager(num_blocks=10, block_size=4, head_dim=D)
    mgr.add_request("s")
    assert len(mgr.free_list) == 10                      # a new request holds no block
    for n in range(1, 18):
        mgr.append("s", vec(rng), vec(rng))
        assert len(mgr.tables["s"]) == math.ceil(n / 4)
        assert len(mgr.free_list) == 10 - math.ceil(n / 4)


def test_gather_follows_the_table_not_physical_order():
    """Two requests grow in turn, so their blocks interleave in the pool."""
    rng = np.random.default_rng(1)
    mgr, ref = BlockManager(8, 4, D), Reference()
    for rid in "ab":
        mgr.add_request(rid)
        ref.seqs[rid] = []
    for _ in range(9):
        for rid in "ab":
            k, v = vec(rng), vec(rng)
            mgr.append(rid, k, v)
            ref.seqs[rid].append((k, v))
    assert mgr.tables["a"] == [0, 2, 4] and mgr.tables["b"] == [1, 3, 5]   # interleaved in the pool
    check(mgr, ref, "a")
    check(mgr, ref, "b")


def test_out_of_blocks_changes_nothing_even_for_copy_on_write():
    rng = np.random.default_rng(2)
    mgr = BlockManager(num_blocks=2, block_size=4, head_dim=D)
    mgr.add_request("p")
    for _ in range(6):
        mgr.append("p", vec(rng), vec(rng))               # both blocks in use, last one partial
    mgr.fork("p", "c")
    before = (dict(mgr.tables), list(mgr.ref_count), list(mgr.free_list), dict(mgr.lengths))
    with pytest.raises(OutOfBlocks):
        mgr.append("c", vec(rng), vec(rng))               # needs a copy, no block is free
    assert (dict(mgr.tables), list(mgr.ref_count), list(mgr.free_list), dict(mgr.lengths)) == before


def test_copy_on_write_does_not_leak_into_the_sibling():
    rng = np.random.default_rng(3)
    mgr, ref = BlockManager(8, 4, D), Reference()
    mgr.add_request("p")
    ref.seqs["p"] = []
    for _ in range(6):
        k, v = vec(rng), vec(rng)
        mgr.append("p", k, v)
        ref.seqs["p"].append((k, v))
    mgr.fork("p", "c")
    ref.seqs["c"] = list(ref.seqs["p"])
    for rid in ("c", "p", "c"):                           # either side may write first
        k, v = vec(rng), vec(rng)
        mgr.append(rid, k, v)
        ref.seqs[rid].append((k, v))
    check(mgr, ref, "p")
    check(mgr, ref, "c")
    assert mgr.tables["p"][0] == mgr.tables["c"][0]       # the full block is still shared


def test_random_operations_match_the_reference_and_never_lose_a_block():
    """Guards double frees and leaked blocks across add, append, fork and free."""
    for seed in range(5):
        rnd, rng = random.Random(seed), np.random.default_rng(seed)
        mgr, ref, next_id = BlockManager(12, 4, D), Reference(), 0
        for _ in range(400):
            op = rnd.choice(["add", "append", "append", "append", "fork", "free"])
            live = list(ref.seqs)
            if op == "add" or not live:
                mgr.add_request(next_id)
                ref.seqs[next_id] = []
                next_id += 1
            elif op == "append":
                rid = rnd.choice(live)
                k, v = vec(rng), vec(rng)
                try:
                    mgr.append(rid, k, v)
                    ref.seqs[rid].append((k, v))
                except OutOfBlocks:
                    pass
            elif op == "fork":
                rid = rnd.choice(live)
                mgr.fork(rid, next_id)
                ref.seqs[next_id] = list(ref.seqs[rid])
                next_id += 1
            else:
                rid = rnd.choice(live)
                mgr.free(rid)
                del ref.seqs[rid]
            held = [b for t in mgr.tables.values() for b in t]
            assert sorted(set(held) | set(mgr.free_list)) == list(range(12)), "a block was lost"
            assert not set(held) & set(mgr.free_list), "a block is both in use and free"
            assert len(set(mgr.free_list)) == len(mgr.free_list), "a block is on the free list twice"
            assert all(mgr.ref_count[b] == held.count(b) for b in range(12))
            for rid in ref.seqs:
                check(mgr, ref, rid)


@pytest.mark.parametrize("block_size", [1, 4, 16])
def test_reading_through_the_table_equals_reading_one_array(block_size):
    """The point of the design: scattered blocks give the same attention output."""
    rng = np.random.default_rng(block_size)
    mgr = BlockManager(num_blocks=64, block_size=block_size, head_dim=4)
    for rid in ("x", "y"):                                # grown in turn below, so blocks are scattered
        mgr.add_request(rid)
    for i in range(13):
        for rid, n in (("x", 13), ("y", 7)):
            if i < n:
                mgr.append(rid, rng.standard_normal(4), rng.standard_normal(4))
    q = rng.standard_normal(4)
    for rid in ("x", "y"):
        keys, values = mgr.gather(rid)
        assert np.allclose(paged_attention(q, mgr, rid), attention_contiguous(q, keys, values))


def test_worked_example_matches_the_lesson_tables():
    """The tables printed in the lesson text. If the code changes, the lesson must change."""
    rows = [(tables, free) for _, tables, free, _ in worked_example()]
    assert rows == [
        ({}, [0, 1, 2, 3, 4, 5, 6, 7]),
        ({"A": [0, 1]}, [2, 3, 4, 5, 6, 7]),
        ({"A": [0, 1], "B": [2]}, [3, 4, 5, 6, 7]),
        ({"A": [0, 1], "B": [2], "C": [3, 4]}, [5, 6, 7]),
        ({"A": [0, 1], "B": [2], "C": [3, 4]}, [5, 6, 7]),
        ({"A": [0, 1, 5], "B": [2], "C": [3, 4]}, [6, 7]),
        ({"A": [0, 1, 5], "B": [2, 6], "C": [3, 4]}, [7]),
        ({"A": [0, 1, 5], "C": [3, 4]}, [7, 2, 6]),
        ({"A": [0, 1, 5], "C": [3, 4, 7]}, [2, 6]),
    ]
    steps = prefix_sharing_example()
    assert [(t, r[:3], f) for _, t, r, f in steps] == [
        ({"S1": [0, 1]}, [1, 1, 0], [2, 3, 4, 5]),
        ({"S1": [0, 1], "S2": [0, 1]}, [2, 2, 0], [2, 3, 4, 5]),
        ({"S1": [0, 2], "S2": [0, 1]}, [2, 1, 1], [3, 4, 5]),
        ({"S1": [0, 2], "S2": [0, 1]}, [2, 1, 1], [3, 4, 5]),
    ]