| """Step 4.5: paged KV cache (vLLM-style blocks), a drop-in for SlotKVCache in run_continuous. |
| |
| The analogy is virtual memory: |
| block table = page table (one per sequence: logical block j -> physical block) |
| block = page (block_size tokens of K and V, every layer) |
| token t = (table[t // block_size], t % block_size) |
| |
| All layers share one block table: block b holds the same token positions in every layer's |
| pool. Blocks are allocated as a sequence grows, not reserved up front, and every block has |
| the same size, so a freed block fits any sequence: no external fragmentation, and at most |
| block_size - 1 wasted slots per sequence (its last, partly filled block). |
| |
| Left padding from the batched prefill is not stored: only real tokens go into blocks. |
| The pad count becomes a RoPE offset (`first`), so a token at storage index j keeps its |
| rotation position first + j and the next query lands at first + len, as in SlotKVCache. |
| |
| Attention reads the blocks by gathering each sequence's blocks into a contiguous |
| (S, H, n_blocks * block_size, D) view and calling SDPA as before. vLLM's PagedAttention |
| kernel reads the blocks in place instead; the gather is this toy's extra copy (measured |
| in docs/INFERENCE.md). |
| |
| When decode needs more blocks than are free, the newest sequence is preempted: its blocks |
| are freed and the engine requeues it to be recomputed (prompt + tokens so far), which |
| reproduces the same greedy tokens. This is vLLM's default recompute policy. |
| """ |
|
|
| import math |
|
|
| import torch |
| from torch import Tensor |
|
|
| from tinyllm.kv_cache import KVCache |
|
|
|
|
| class OutOfBlocks(RuntimeError): |
| pass |
|
|
|
|
| class BlockAllocator: |
| """Fixed-size blocks with a free list. Lowest ids are handed out first (reused LIFO).""" |
|
|
| def __init__(self, n_blocks: int, block_size: int): |
| self.n_blocks, self.block_size = n_blocks, block_size |
| self._free = list(range(n_blocks - 1, -1, -1)) |
| self._in_use = [False] * n_blocks |
| self.peak_used = 0 |
|
|
| @property |
| def num_free(self) -> int: |
| return len(self._free) |
|
|
| @property |
| def num_used(self) -> int: |
| return self.n_blocks - len(self._free) |
|
|
| def alloc(self) -> int: |
| if not self._free: |
| raise OutOfBlocks(f"all {self.n_blocks} KV blocks in use") |
| b = self._free.pop() |
| self._in_use[b] = True |
| self.peak_used = max(self.peak_used, self.num_used) |
| return b |
|
|
| def free(self, blocks: list[int]) -> None: |
| for b in blocks: |
| if not self._in_use[b]: |
| raise ValueError(f"block {b} freed twice (or never allocated)") |
| self._in_use[b] = False |
| self._free.append(b) |
|
|
|
|
| class PagedKVCache: |
| """S sequence slots over a shared pool of n_blocks KV blocks. |
| |
| Same model protocol as SlotKVCache (begin_step / layer / end_step, one token per slot) |
| and the same scheduling calls (free_slots / admit / release), plus can_admit() and |
| reserve_step() so the engine can respect the block budget. |
| """ |
|
|
| def __init__(self, n_layers: int, n_slots: int, n_kv_heads: int, head_dim: int, |
| n_blocks: int, block_size: int, max_len: int, dtype: torch.dtype, device): |
| self.alloc = BlockAllocator(n_blocks, block_size) |
| self.bs, self.n_slots, self.max_len = block_size, n_slots, max_len |
| self.max_blocks = math.ceil(max_len / block_size) |
| self.scratch = n_blocks |
| shape = (n_blocks + 1, n_kv_heads, block_size, head_dim) |
| self.k = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)] |
| self.v = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)] |
| |
| self.tables: list[list[int]] = [[] for _ in range(n_slots)] |
| self.lens = [0] * n_slots |
| self.first = [0] * n_slots |
| self.active = [False] * n_slots |
| self._admit_no = [0] * n_slots |
| self._admits = 0 |
| self._table_dev = torch.full((n_slots, self.max_blocks), self.scratch, |
| dtype=torch.long, device=device) |
| self._lens_dev = torch.zeros(n_slots, dtype=torch.long, device=device) |
| self._first_dev = torch.zeros(n_slots, dtype=torch.long, device=device) |
| self._active_dev = torch.zeros(n_slots, dtype=torch.long, device=device) |
| self._rows = torch.arange(n_slots, device=device) |
| self.attn_kwargs: dict = {} |
| self._n_read_blocks = 0 |
| self._blk = self._off = None |
|
|
| |
| def free_slots(self) -> list[int]: |
| return [s for s in range(self.n_slots) if not self.active[s]] |
|
|
| def can_admit(self, lengths: list[int]) -> bool: |
| """Enough free blocks for these prompts, plus room for each one's first decode token.""" |
| return sum(math.ceil((n + 1) / self.bs) for n in lengths) <= self.alloc.num_free |
|
|
| def admit(self, slots: list[int], prefill: KVCache, pad_lens: list[int]) -> None: |
| """Copy row j of a padded prefill (real tokens only) into fresh blocks for slots[j].""" |
| L = prefill.pos |
| dev = self._lens_dev.device |
| rows, src, blk, off = [], [], [], [] |
| for j, (s, p) in enumerate(zip(slots, pad_lens)): |
| n = L - p |
| if n >= self.max_len: |
| raise ValueError(f"prompt length {n} >= max_len {self.max_len}") |
| blocks = [self.alloc.alloc() for _ in range(math.ceil(n / self.bs))] |
| self.tables[s] = blocks |
| self.lens[s], self.first[s], self.active[s] = n, p, True |
| self._admits += 1 |
| self._admit_no[s] = self._admits |
| self._table_dev[s].fill_(self.scratch) |
| self._table_dev[s, :len(blocks)] = torch.tensor(blocks, device=dev) |
| t = torch.arange(n) |
| rows.append(torch.full((n,), j)) |
| src.append(t + p) |
| blk.append(torch.tensor(blocks, dtype=torch.long)[t // self.bs]) |
| off.append(t % self.bs) |
| rows, src, blk, off = (torch.cat(x).to(dev) for x in (rows, src, blk, off)) |
| for i in range(len(self.k)): |
| self.k[i][blk, :, off] = prefill.k[i][rows, :, src] |
| self.v[i][blk, :, off] = prefill.v[i][rows, :, src] |
| self._sync() |
|
|
| def release(self, slot: int) -> None: |
| self.alloc.free(self.tables[slot]) |
| self.tables[slot] = [] |
| self.active[slot] = False |
| self.lens[slot] = self.first[slot] = 0 |
| self._table_dev[slot].fill_(self.scratch) |
| self._sync() |
|
|
| def reserve_step(self) -> list[int]: |
| """Make sure the next decode step's new blocks can be allocated. Preempts the newest |
| sequences (frees their blocks) until they can; returns the preempted slots, newest first.""" |
| victims = [] |
| while True: |
| need = sum(1 for s in range(self.n_slots) |
| if self.active[s] and len(self.tables[s]) * self.bs <= self.lens[s]) |
| if need <= self.alloc.num_free or not any(self.active): |
| return victims |
| v = max((s for s in range(self.n_slots) if self.active[s]), key=self._admit_no.__getitem__) |
| self.release(v) |
| victims.append(v) |
|
|
| def _sync(self) -> None: |
| self._lens_dev.copy_(torch.tensor(self.lens)) |
| self._first_dev.copy_(torch.tensor(self.first)) |
| self._active_dev.copy_(torch.tensor(self.active, dtype=torch.long)) |
|
|
| def used_tokens(self) -> int: |
| return sum(n for n, a in zip(self.lens, self.active) if a) |
|
|
| def allocated_tokens(self) -> int: |
| return self.alloc.num_used * self.bs |
|
|
| def reserved_tokens(self) -> int: |
| """Pool capacity: the KV memory this cache holds.""" |
| return self.alloc.n_blocks * self.bs |
|
|
| def nbytes(self) -> int: |
| return sum(t.numel() * t.element_size() for t in self.k + self.v) |
|
|
| |
| def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]: |
| if T != 1: |
| raise ValueError("PagedKVCache decodes one token per slot; prefill goes through KVCache") |
| if max(self.lens) >= self.max_len: |
| raise ValueError(f"a sequence is full (max_len {self.max_len})") |
| new_rows, new_cols, new_blks = [], [], [] |
| for s in range(self.n_slots): |
| if self.active[s] and len(self.tables[s]) * self.bs <= self.lens[s]: |
| b = self.alloc.alloc() |
| new_rows.append(s) |
| new_cols.append(len(self.tables[s])) |
| new_blks.append(b) |
| self.tables[s].append(b) |
| if new_rows: |
| dev = self._table_dev.device |
| self._table_dev[torch.tensor(new_rows, device=dev), |
| torch.tensor(new_cols, device=dev)] = torch.tensor(new_blks, device=dev) |
| self._n_read_blocks = math.ceil((max(self.lens) + 1) / self.bs) |
| self._blk = self._table_dev[self._rows, self._lens_dev // self.bs] |
| self._off = self._lens_dev % self.bs |
| k_pos = torch.arange(self._n_read_blocks * self.bs, device=self._lens_dev.device) |
| self.attn_kwargs = {"attn_mask": (k_pos <= self._lens_dev[:, None])[:, None, None, :]} |
| return rope.at((self._first_dev + self._lens_dev)[:, None]) |
|
|
| def end_step(self, T: int) -> None: |
| for s in range(self.n_slots): |
| if self.active[s]: |
| self.lens[s] += 1 |
| self._lens_dev += self._active_dev |
|
|
| def layer(self, i: int) -> "PagedLayerKV": |
| return PagedLayerKV(self, i) |
|
|
|
|
| class PagedLayerKV: |
| def __init__(self, cache: PagedKVCache, layer_idx: int): |
| self.cache = cache |
| self.i = layer_idx |
|
|
| @property |
| def attn_kwargs(self) -> dict: |
| return self.cache.attn_kwargs |
|
|
| def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: |
| """k, v (S, H, 1, D): write each slot's token into its block, then gather its blocks.""" |
| c, i = self.cache, self.i |
| c.k[i][c._blk, :, c._off] = k[:, :, 0] |
| c.v[i][c._blk, :, c._off] = v[:, :, 0] |
| table = c._table_dev[:, :c._n_read_blocks] |
| S, nb = table.shape |
| _, H, bs, D = c.k[i].shape |
|
|
| def gather(pool: Tensor) -> Tensor: |
| return pool[table].transpose(1, 2).reshape(S, H, nb * bs, D) |
|
|
| return gather(c.k[i]), gather(c.v[i]) |
|
|
