Implement paged cache allocation and copy-on-write

"""Contiguous KV cache allocator vs PagedAttention block tables with copy-on-write."""

from __future__ import annotations

from dataclasses import dataclass, field
import json
import math


@dataclass
class ContiguousSlot:
    seq_id: str
    reserved_tokens: int
    used_tokens: int


class ContiguousKvAllocator:
    """Reserves a contiguous chunk of max_seq_len tokens per active request."""

    def __init__(self, total_token_capacity: int, max_seq_len: int) -> None:
        if total_token_capacity <= 0 or max_seq_len <= 0:
            raise ValueError("Capacity and max_seq_len must be positive integers")
        self.total_token_capacity = total_token_capacity
        self.max_seq_len = max_seq_len
        self.max_concurrent_sequences = total_token_capacity // max_seq_len
        self.active: dict[str, ContiguousSlot] = {}

    def allocate_sequence(self, seq_id: str, initial_tokens: int) -> None:
        if seq_id in self.active:
            raise ValueError(f"Duplicate sequence ID: {seq_id}")
        if initial_tokens <= 0 or initial_tokens > self.max_seq_len:
            raise ValueError(
                f"initial_tokens={initial_tokens} outside 1..{self.max_seq_len}"
            )
        if len(self.active) >= self.max_concurrent_sequences:
            raise MemoryError(
                f"Contiguous KV allocator out of slots: {len(self.active)}/{self.max_concurrent_sequences} in use"
            )
        self.active[seq_id] = ContiguousSlot(
            seq_id=seq_id,
            reserved_tokens=self.max_seq_len,
            used_tokens=initial_tokens,
        )

    def append_tokens(self, seq_id: str, count: int = 1) -> None:
        slot = self.active[seq_id]
        if slot.used_tokens + count > slot.reserved_tokens:
            raise MemoryError(
                f"Sequence {seq_id} exceeded max_seq_len={self.max_seq_len}"
            )
        slot.used_tokens += count

    def free_sequence(self, seq_id: str) -> None:
        self.active.pop(seq_id, None)

    def metrics(self) -> dict[str, object]:
        reserved = sum(s.reserved_tokens for s in self.active.values())
        used = sum(s.used_tokens for s in self.active.values())
        wasted = reserved - used
        waste_ratio = (wasted / reserved) if reserved else 0.0
        return {
            "active_sequences": len(self.active),
            "reserved_tokens": reserved,
            "used_tokens": used,
            "wasted_tokens": wasted,
            "waste_ratio": round(waste_ratio, 4),
            "utilization_of_total_pool": round(used / self.total_token_capacity, 4),
        }


@dataclass
class PhysicalKvBlock:
    block_id: int
    block_size: int
    ref_count: int = 0
    tokens: list[int] = field(default_factory=list)


@dataclass
class BlockTable:
    seq_id: str
    block_size: int
    logical_blocks: list[int] = field(default_factory=list)
    num_tokens: int = 0


class PagedKvAllocator:
    """Manages fixed-size physical KV blocks with reference counting and copy-on-write."""

    def __init__(self, total_token_capacity: int, block_size: int = 16) -> None:
        if total_token_capacity <= 0 or block_size <= 0:
            raise ValueError("total_token_capacity and block_size must be positive")
        if total_token_capacity % block_size != 0:
            raise ValueError("total_token_capacity must be a multiple of block_size")
        self.total_token_capacity = total_token_capacity
        self.block_size = block_size
        self.num_blocks = total_token_capacity // block_size
        self.blocks: dict[int, PhysicalKvBlock] = {
            i: PhysicalKvBlock(block_id=i, block_size=block_size)
            for i in range(self.num_blocks)
        }
        self.free_block_ids: list[int] = list(range(self.num_blocks))
        self.tables: dict[str, BlockTable] = {}
        self.cow_copies_performed: int = 0

    def _acquire_block(self) -> PhysicalKvBlock:
        if not self.free_block_ids:
            raise MemoryError("PagedKvAllocator out of physical KV blocks")
        block_id = self.free_block_ids.pop(0)
        block = self.blocks[block_id]
        block.ref_count = 1
        block.tokens = []
        return block

    def _release_block(self, block_id: int) -> None:
        block = self.blocks[block_id]
        if block.ref_count <= 0:
            raise RuntimeError(f"Block {block_id} already has ref_count=0")
        block.ref_count -= 1
        if block.ref_count == 0:
            block.tokens = []
            self.free_block_ids.append(block_id)

    def allocate_sequence(self, seq_id: str, token_ids: list[int]) -> BlockTable:
        if seq_id in self.tables:
            raise ValueError(f"Duplicate sequence ID: {seq_id}")
        if not token_ids:
            raise ValueError("token_ids must be non-empty")
        needed_blocks = math.ceil(len(token_ids) / self.block_size)
        if needed_blocks > len(self.free_block_ids):
            raise MemoryError(
                f"Need {needed_blocks} blocks for {seq_id}, only {len(self.free_block_ids)} free"
            )
        table = BlockTable(seq_id=seq_id, block_size=self.block_size)
        self.tables[seq_id] = table
        for token in token_ids:
            self.append_token(seq_id, token)
        return table

    def fork_sequence(self, parent_seq_id: str, child_seq_id: str) -> BlockTable:
        """Share physical blocks with ref_count += 1 for beam search / parallel sampling."""
        if parent_seq_id not in self.tables:
            raise KeyError(f"Unknown parent sequence: {parent_seq_id}")
        if child_seq_id in self.tables:
            raise ValueError(f"Duplicate child sequence: {child_seq_id}")
        parent = self.tables[parent_seq_id]
        for block_id in parent.logical_blocks:
            self.blocks[block_id].ref_count += 1
        child = BlockTable(
            seq_id=child_seq_id,
            block_size=self.block_size,
            logical_blocks=list(parent.logical_blocks),
            num_tokens=parent.num_tokens,
        )
        self.tables[child_seq_id] = child
        return child

    def append_token(
        self, seq_id: str, token_id: int, enforce_cow: bool = True
    ) -> None:
        table = self.tables[seq_id]
        if (
            not table.logical_blocks
            or len(self.blocks[table.logical_blocks[-1]].tokens) == self.block_size
        ):
            new_block = self._acquire_block()
            new_block.tokens.append(token_id)
            table.logical_blocks.append(new_block.block_id)
            table.num_tokens += 1
            return

        last_block_id = table.logical_blocks[-1]
        last_block = self.blocks[last_block_id]
        if last_block.ref_count > 1 and enforce_cow:
            # Copy-on-write: allocate a private physical block, copy existing tokens, decrement old ref_count
            private_block = self._acquire_block()
            private_block.tokens = list(last_block.tokens)
            private_block.tokens.append(token_id)
            last_block.ref_count -= 1
            table.logical_blocks[-1] = private_block.block_id
            table.num_tokens += 1
            self.cow_copies_performed += 1
        else:
            last_block.tokens.append(token_id)
            table.num_tokens += 1

    def read_sequence_tokens(self, seq_id: str) -> list[int]:
        table = self.tables[seq_id]
        out: list[int] = []
        for block_id in table.logical_blocks:
            out.extend(self.blocks[block_id].tokens)
        return out

    def free_sequence(self, seq_id: str) -> None:
        table = self.tables.pop(seq_id)
        for block_id in table.logical_blocks:
            self._release_block(block_id)

    def metrics(self) -> dict[str, object]:
        allocated_blocks = [b for b in self.blocks.values() if b.ref_count > 0]
        reserved_tokens = len(allocated_blocks) * self.block_size
        physical_stored_tokens = sum(len(b.tokens) for b in allocated_blocks)
        logical_tokens = sum(t.num_tokens for t in self.tables.values())
        internal_frag_tokens = reserved_tokens - physical_stored_tokens
        frag_ratio = (
            (internal_frag_tokens / reserved_tokens) if reserved_tokens else 0.0
        )
        return {
            "active_sequences": len(self.tables),
            "allocated_blocks": len(allocated_blocks),
            "free_blocks": len(self.free_block_ids),
            "reserved_tokens": reserved_tokens,
            "physical_stored_tokens": physical_stored_tokens,
            "logical_tokens_served": logical_tokens,
            "internal_fragmentation_tokens": internal_frag_tokens,
            "waste_ratio": round(frag_ratio, 4),
            "cow_copies_performed": self.cow_copies_performed,
        }


def compare_contiguous_vs_paged() -> dict[str, object]:
    total_capacity = 4096
    max_seq_len = 512
    block_size = 16
    request_lengths = [
        68,
        95,
        114,
        52,
        140,
        83,
        77,
        129,
        61,
        105,
        88,
        73,
        119,
        91,
        64,
        110,
        79,
        102,
        58,
        134,
        85,
        70,
        99,
        122,
        66,
        108,
        74,
        115,
        81,
        93,
        62,
        125,
    ]

    contiguous = ContiguousKvAllocator(
        total_token_capacity=total_capacity, max_seq_len=max_seq_len
    )
    contiguous_admitted = 0
    contiguous_rejected = 0
    for idx, length in enumerate(request_lengths):
        try:
            contiguous.allocate_sequence(f"req-{idx}", length)
            contiguous_admitted += 1
        except MemoryError:
            contiguous_rejected += 1

    paged = PagedKvAllocator(total_token_capacity=total_capacity, block_size=block_size)
    paged_admitted = 0
    paged_rejected = 0
    for idx, length in enumerate(request_lengths):
        try:
            paged.allocate_sequence(f"req-{idx}", list(range(length)))
            paged_admitted += 1
        except MemoryError:
            paged_rejected += 1

    # Demonstrate shared prefix + copy-on-write parallel sampling
    cow_allocator = PagedKvAllocator(total_token_capacity=512, block_size=16)
    shared_prompt = list(
        range(40)
    )  # 2 full blocks (32 tokens) + 1 partial block (8 tokens)
    cow_allocator.allocate_sequence("beam-0", shared_prompt)
    cow_allocator.fork_sequence("beam-0", "beam-1")
    cow_allocator.fork_sequence("beam-0", "beam-2")
    blocks_before_decode = len(
        [b for b in cow_allocator.blocks.values() if b.ref_count > 0]
    )
    cow_allocator.append_token("beam-1", 9001)
    cow_allocator.append_token("beam-2", 9002)
    cow_allocator.append_token("beam-0", 9000)

    return {
        "pool_config": {
            "total_token_capacity": total_capacity,
            "max_seq_len": max_seq_len,
            "block_size": block_size,
            "submitted_requests": len(request_lengths),
        },
        "contiguous": {
            "admitted_requests": contiguous_admitted,
            "rejected_requests": contiguous_rejected,
            **contiguous.metrics(),
        },
        "paged": {
            "admitted_requests": paged_admitted,
            "rejected_requests": paged_rejected,
            **paged.metrics(),
        },
        "concurrency_multiplier": round(paged_admitted / contiguous_admitted, 2),
        "copy_on_write_demo": {
            "prompt_tokens": len(shared_prompt),
            "shared_beams": 3,
            "blocks_before_decode_append": blocks_before_decode,
            **cow_allocator.metrics(),
            "beam_tails": {
                "beam_0_last_token": cow_allocator.read_sequence_tokens("beam-0")[-1],
                "beam_1_last_token": cow_allocator.read_sequence_tokens("beam-1")[-1],
                "beam_2_last_token": cow_allocator.read_sequence_tokens("beam-2")[-1],
            },
        },
    }


def main() -> None:
    print(
        json.dumps(
            {
                "execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
                **compare_contiguous_vs_paged(),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()

"""Unit and failure-boundary tests for contiguous vs PagedAttention KV allocators."""

from __future__ import annotations

import unittest

from lab import ContiguousKvAllocator, PagedKvAllocator, compare_contiguous_vs_paged


class PagedKvCachingTest(unittest.TestCase):
    def test_paged_allocator_eliminates_reserved_slack_and_admits_4x_more_requests(
        self,
    ) -> None:
        report = compare_contiguous_vs_paged()
        self.assertEqual(report["contiguous"]["admitted_requests"], 8)
        self.assertEqual(report["paged"]["admitted_requests"], 32)
        self.assertGreater(report["contiguous"]["waste_ratio"], 0.80)
        self.assertLess(report["paged"]["waste_ratio"], 0.10)

    def test_contiguous_allocator_ooms_on_reserved_slack(self) -> None:
        """Contiguous allocation rejects a 9th 64-token request even though 87.5% of the pool is unused."""
        allocator = ContiguousKvAllocator(total_token_capacity=4096, max_seq_len=512)
        for i in range(8):
            allocator.allocate_sequence(f"req-{i}", initial_tokens=64)
        metrics = allocator.metrics()
        self.assertEqual(metrics["used_tokens"], 512)
        self.assertEqual(metrics["wasted_tokens"], 3584)
        with self.assertRaises(MemoryError):
            allocator.allocate_sequence("req-8", initial_tokens=64)

    def test_mutating_shared_block_without_cow_corrupts_sibling(self) -> None:
        """Appending to a shared tail block with enforce_cow=False overwrites/pollutes the sibling sequence."""
        safe = PagedKvAllocator(total_token_capacity=256, block_size=16)
        safe.allocate_sequence(
            "parent", list(range(20))
        )  # 1 full block + 4 tokens in tail block
        safe.fork_sequence("parent", "child")
        safe.append_token("child", 999, enforce_cow=True)
        self.assertEqual(safe.read_sequence_tokens("parent"), list(range(20)))
        self.assertEqual(safe.read_sequence_tokens("child"), list(range(20)) + [999])

        unsafe = PagedKvAllocator(total_token_capacity=256, block_size=16)
        unsafe.allocate_sequence("parent", list(range(20)))
        unsafe.fork_sequence("parent", "child")
        unsafe.append_token("child", 999, enforce_cow=False)
        # Without copy-on-write, parent's tail block now unexpectedly contains child's token 999!
        self.assertIn(999, unsafe.read_sequence_tokens("parent"))

    def test_freeing_forked_sequence_preserves_shared_prefix_until_last_reference(
        self,
    ) -> None:
        allocator = PagedKvAllocator(total_token_capacity=128, block_size=16)
        allocator.allocate_sequence("seq-a", list(range(32)))
        allocator.fork_sequence("seq-a", "seq-b")
        self.assertEqual(allocator.metrics()["allocated_blocks"], 2)
        allocator.free_sequence("seq-a")
        # Both blocks remain allocated because seq-b still holds ref_count=1
        self.assertEqual(allocator.metrics()["allocated_blocks"], 2)
        self.assertEqual(allocator.read_sequence_tokens("seq-b"), list(range(32)))
        allocator.free_sequence("seq-b")
        self.assertEqual(allocator.metrics()["allocated_blocks"], 0)


if __name__ == "__main__":
    unittest.main()


class GuardBoundaryTests(unittest.TestCase):
    def test_rejects_invalid_boundary_inputs_with_specific_reason(self):
        import lab
        cases = [('ContiguousKvAllocator(0,1)', 'Capacity and'), ('a=ContiguousKvAllocator(4,4)\na.allocate_sequence("x",1)\na.allocate_sequence("x",1)', 'Duplicate sequence'), ('ContiguousKvAllocator(4,4).allocate_sequence("x",0)', 'initial_tokens'), ('a=ContiguousKvAllocator(4,4)\na.allocate_sequence("x",4)\na.append_tokens("x")', 'exceeded max'), ('PagedKvAllocator(0,1)', 'total_token_capacity and'), ('PagedKvAllocator(5,2)', 'multiple of'), ('a=PagedKvAllocator(2,2)\na._acquire_block()\na._acquire_block()', 'out of physical'), ('PagedKvAllocator(2,2)._release_block(0)', 'already has ref_count'), ('a=PagedKvAllocator(4,2)\na.allocate_sequence("x",[1])\na.allocate_sequence("x",[1])', 'Duplicate sequence'), ('PagedKvAllocator(4,2).allocate_sequence("x",[])', 'non-empty'), ('PagedKvAllocator(2,2).allocate_sequence("x",[1,2,3])', 'Need 2 blocks'), ('PagedKvAllocator(2,2).fork_sequence("missing","x")', 'Unknown parent'), ('a=PagedKvAllocator(4,2)\na.allocate_sequence("x",[1])\na.fork_sequence("x","x")', 'Duplicate child')]
        for code, reason in cases:
            with self.subTest(code=code):
                with self.assertRaisesRegex((ValueError, MemoryError, KeyError, RuntimeError), reason):
                    exec(code, vars(lab).copy())