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]), ]