From deb2d7e1272d7ee3a258617ea37b10cb322b2f96 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 30 Jul 2026 17:18:07 +0800 Subject: [PATCH] refactor: rebuild KV cache with three-layer separation architecture - Replace CacheView/ContiguousCache/PageCache with SGLang-inspired design: KVStorage (flat token-level NHD buffers [n_layers, size, H, D]), ReqToTokenPool (index table [req_idx, pos] -> token_slot), Allocator + PrefixCache (slot allocation with LRU and prefix sharing) - Add KVCache as pure dataclass passed to model: k_buffer, v_buffer, req_to_token, req_pool_indices, seq_lens, out_cache_loc - PagePool orchestrates all three layers, supports contiguous mode (pre-allocated per-request blocks, default) and paged mode (page_size=1 or >1 with dynamic allocation and prefix caching) - Attention layers now do raw buffer indexing instead of opaque write/gather method calls on CacheView objects - Update executor.bind_tasks signature: seq_lens list + start_pos - Rename paged_cache -> kv_cache throughout model/ and inference/ --- astrai/inference/__init__.py | 18 +- astrai/inference/core/__init__.py | 18 +- astrai/inference/core/cache.py | 669 +++++++++++------------ astrai/inference/core/executor.py | 13 +- astrai/inference/core/scheduler.py | 20 +- astrai/inference/engine.py | 4 +- astrai/model/components/attention.py | 38 +- astrai/model/components/decoder_block.py | 6 +- astrai/model/transformer.py | 6 +- tests/inference/test_cache.py | 459 +++++++++------- 10 files changed, 644 insertions(+), 607 deletions(-) diff --git a/astrai/inference/__init__.py b/astrai/inference/__init__.py index d112b89..e0db8f4 100644 --- a/astrai/inference/__init__.py +++ b/astrai/inference/__init__.py @@ -30,21 +30,16 @@ from astrai.inference.api.openai import OpenAIResponseBuilder from astrai.inference.core import ( STOP, Allocator, - CacheView, - ContiguousCache, - ContiguousCacheView, Executor, InferenceScheduler, KVCache, - PageCache, - PageCacheView, + KVStorage, PagePool, PrefixCache, - Storage, + ReqToTokenPool, Task, TaskManager, TaskStatus, - TaskTable, page_hash, ) from astrai.inference.engine import GenerationRequest, InferenceEngine @@ -68,16 +63,11 @@ __all__ = [ "TaskManager", "TaskStatus", "Allocator", - "CacheView", "KVCache", - "ContiguousCache", - "ContiguousCacheView", - "PageCache", - "PageCacheView", + "KVStorage", "PagePool", "PrefixCache", - "Storage", - "TaskTable", + "ReqToTokenPool", "page_hash", "sample", "BaseSamplingStrategy", diff --git a/astrai/inference/core/__init__.py b/astrai/inference/core/__init__.py index a545e9b..f4d6de8 100644 --- a/astrai/inference/core/__init__.py +++ b/astrai/inference/core/__init__.py @@ -2,16 +2,11 @@ from astrai.inference.core.cache import ( Allocator, - CacheView, - ContiguousCache, - ContiguousCacheView, KVCache, - PageCache, - PageCacheView, + KVStorage, PagePool, PrefixCache, - Storage, - TaskTable, + ReqToTokenPool, page_hash, ) from astrai.inference.core.executor import Executor @@ -20,16 +15,11 @@ from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus __all__ = [ "Allocator", - "CacheView", "KVCache", - "ContiguousCache", - "ContiguousCacheView", - "PageCache", - "PageCacheView", + "KVStorage", "PagePool", "PrefixCache", - "Storage", - "TaskTable", + "ReqToTokenPool", "page_hash", "Executor", "InferenceScheduler", diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index e73db69..a710cfc 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -1,7 +1,21 @@ +"""KV cache architecture: three-layer separation (SGLang-inspired). + +Layer 1 — KVStorage: flat token-level K/V buffers [n_layers, size, H, D] +Layer 2 — ReqToTokenPool: index table [req_idx, pos] → physical token slot +Layer 3 — Allocator: slot/page allocation with ref-counting and LRU + +PagePool orchestrates all three plus PrefixCache (content addressing). +KVCache is a pure dataclass passed to the model for direct buffer access. + +Two modes: + - contiguous (default): pre-allocated per-request blocks, no dynamic alloc + - paged: shared pool with on-demand allocation, prefix caching support +""" + import threading -from abc import ABC, abstractmethod from collections import OrderedDict -from typing import Callable, Dict, List, Optional, Tuple +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional import torch from torch import Tensor @@ -108,418 +122,359 @@ class PrefixCache: self._hash_to_page[h] = page_idx -class PagePool: - """Orchestrates allocator (page management) and PrefixCache (content addressing).""" +class ReqToTokenPool: + """Maps [req_idx, pos] -> physical token slot in KV storage. - def __init__(self, allocator: Allocator, prefix: PrefixCache): - self._alloc = allocator - self._prefix = prefix - self._alloc.on_evict = prefix.evict + Each row is one request; each column is a sequence position. The value + at [req_idx, pos] is the flat index into the KV storage buffers. + """ - @property - def allocator(self) -> Allocator: - return self._alloc - - @property - def prefix(self) -> PrefixCache: - return self._prefix - - def alloc(self) -> int: - return self._alloc.alloc() - - def free(self, idx: int): - keep = self._prefix.has_page(idx) - self._alloc.free(idx, keep_cached=keep) - if not keep: - self._prefix.evict(idx) - - def inc_ref(self, idx: int): - self._alloc.inc_ref(idx) - - def lookup(self, token_ids: List[int]) -> List[int]: - hits = self._prefix.lookup(token_ids) - for p in hits: - self._alloc.touch(p) - return hits - - def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int): - self._prefix.record(page_idx, token_ids, logical_page_idx) - - -class TaskTable: - """Maps task_ids to page tables and cached token counts.""" - - def __init__(self, page_size: int): - self._page_size = page_size - self._pages: Dict[str, List[int]] = {} - self._cached: Dict[str, int] = {} + def __init__(self, size: int, max_context_len: int, device: torch.device): + self.size = size + self.max_context_len = max_context_len + self.req_to_token = torch.zeros( + (size, max_context_len), dtype=torch.long, device=device + ) + self.free_slots = list(range(size)) self._lock = threading.Lock() - def set(self, task_id: str, page_table: List[int], cached: int): + def alloc(self, num_reqs: int) -> Optional[List[int]]: with self._lock: - self._pages[task_id] = page_table - self._cached[task_id] = cached + if num_reqs > len(self.free_slots): + return None + slots = self.free_slots[:num_reqs] + self.free_slots = self.free_slots[num_reqs:] + return slots - def get(self, task_id: str) -> List[int]: + def free(self, req_indices: List[int]): with self._lock: - return self._pages.get(task_id, []) + self.free_slots.extend(req_indices) - def get_cached(self, task_id: str) -> int: - with self._lock: - return self._cached.get(task_id, 0) - - def pop(self, task_id: str) -> Tuple[List[int], int]: - with self._lock: - pages = self._pages.pop(task_id, []) - cached = self._cached.pop(task_id, 0) - return pages, cached - - def get_ref(self, task_id: str) -> List[int]: - with self._lock: - return self._pages.setdefault(task_id, []) - - def table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor: - with self._lock: - states = [self._pages.get(tid, []) for tid in task_ids] - max_pages = max((len(s) for s in states), default=0) - rows = [s + [-1] * (max_pages - len(s)) for s in states] - return torch.tensor(rows, dtype=torch.long, device=device) + def write(self, indices, values): + self.req_to_token[indices] = values -class Storage: - """KV-cache tensor storage with paged write/gather.""" +class KVStorage: + """Token-level flat KV cache storage with NHD layout. + + Buffers: [n_layers, size, n_kv_heads, head_dim]. Each token occupies + one contiguous row. Logical ordering is determined by ReqToTokenPool. + """ def __init__( self, + size: int, n_layers: int, - n_pages: int, - page_size: int, n_kv_heads: int, head_dim: int, device: torch.device, dtype: torch.dtype, ): - self.page_size = page_size - self.k_cache = torch.empty( - (n_layers, n_pages, page_size, n_kv_heads, head_dim), - device=device, - dtype=dtype, + self.size = size + self.k_buffer = torch.empty( + (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype ) - self.v_cache = torch.empty( - (n_layers, n_pages, page_size, n_kv_heads, head_dim), - device=device, - dtype=dtype, + self.v_buffer = torch.empty( + (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype ) - def write( - self, - layer_id: int, - page_table: Tensor, - start_pos: int, - k: Tensor, - v: Tensor, - ): - seq_len = k.size(1) - if seq_len == 0: - return - page_size = self.page_size - written = 0 - first_page = start_pos // page_size - last_page = (start_pos + seq_len - 1) // page_size - for pi in range(first_page, last_page + 1): - phys_pages = page_table[:, pi] - page_start = pi * page_size - write_start = max(page_start, start_pos) - write_end = min(page_start + page_size, start_pos + seq_len) - offset = write_start - page_start - chunk = write_end - write_start - valid = phys_pages >= 0 - if not valid.all(): - if valid.any(): - valid_pages = phys_pages[valid] - self.k_cache[layer_id, valid_pages, offset : offset + chunk] = k[ - valid, written : written + chunk - ] - self.v_cache[layer_id, valid_pages, offset : offset + chunk] = v[ - valid, written : written + chunk - ] - written += chunk - continue - self.k_cache[layer_id, phys_pages, offset : offset + chunk] = k[ - :, written : written + chunk - ] - self.v_cache[layer_id, phys_pages, offset : offset + chunk] = v[ - :, written : written + chunk - ] - written += chunk + def get_key_buffer(self, layer_id: int) -> Tensor: + return self.k_buffer[layer_id] - def gather( - self, layer_id: int, page_table: Tensor, total_len: int - ) -> Tuple[Tensor, Tensor]: - safe = page_table.clamp(min=0) - k = self.k_cache[layer_id, safe] - v = self.v_cache[layer_id, safe] - k = k.flatten(1, 2) - v = v.flatten(1, 2) - if (page_table < 0).any(): - invalid = ( - (page_table < 0) - .unsqueeze(-1) - .expand(-1, -1, self.page_size) - .flatten(1, 2) - ) - invalid = invalid[:, :, None, None].expand_as(k) - k = k.masked_fill(invalid, 0.0) - v = v.masked_fill(invalid, 0.0) - k = k[:, :total_len] - v = v[:, :total_len] - return k, v + def get_value_buffer(self, layer_id: int) -> Tensor: + return self.v_buffer[layer_id] + + def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None: + self.k_buffer[layer_id][loc] = k + self.v_buffer[layer_id][loc] = v -class CacheView(ABC): - """Abstract view passed to attention layers for KV-cache I/O.""" +@dataclass +class KVCache: + """Pure data struct passed to model for KV cache I/O. - @abstractmethod - def write(self, layer_id: int, k: Tensor, v: Tensor): ... + The attention layer does raw buffer indexing — no methods, no abstraction. - @abstractmethod - def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ... + Attributes: + k_buffer: [n_layers, size, n_kv_heads, head_dim] + v_buffer: [n_layers, size, n_kv_heads, head_dim] + req_to_token: [num_reqs, max_ctx_len] — index table + req_pool_indices: [batch_size] — row indices into req_to_token + seq_lens: [batch_size] — per-request total sequence lengths + out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices + """ + + k_buffer: Tensor + v_buffer: Tensor + req_to_token: Tensor + req_pool_indices: Tensor + seq_lens: Tensor + out_cache_loc: Tensor -class KVCache(ABC): - """Abstract KV-cache facade for scheduler/executor.""" +class PagePool: + """Top-level KV cache manager. - @abstractmethod - def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ... + Combines KVStorage + ReqToTokenPool + Allocator + PrefixCache. - @abstractmethod - def task_free(self, task_id: str): ... - - @abstractmethod - def task_extend(self, task_id: str, pos: int) -> bool: ... - - @abstractmethod - def bind_tasks( - self, - task_ids: List[str], - total_len: int, - device: torch.device, - write_positions: Optional[Tensor] = None, - ) -> CacheView: ... - - def task_cached(self, task_id: str) -> int: - return 0 - - def task_record_hashes( - self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0 - ): ... - - -class PageCacheView(CacheView): - """Bundles Storage + page_table + total_len for attention layers.""" - - def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0): - self._storage = storage - self._page_table = page_table - self._total_len = total_len - - def write(self, layer_id: int, k: Tensor, v: Tensor): - start_pos = self._total_len - k.size(1) - self._storage.write(layer_id, self._page_table, start_pos, k, v) - - def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: - return self._storage.gather(layer_id, self._page_table, self._total_len) - - -class PageCache(KVCache): - """Paged KV-cache with prefix sharing.""" + Args: + n_layers: Number of transformer layers. + n_kv_heads: Number of KV attention heads. + head_dim: Dimension per head. + max_batch_size: Maximum concurrent requests. + max_seq_len: Maximum sequence length per request. + device, dtype: Tensor device and dtype. + page_size: Page size for paged mode (1 = token-level). + n_tokens: Total token slots for paged mode. None = contiguous mode + (pre-allocates max_batch_size * max_seq_len). + """ def __init__( self, n_layers: int, - n_pages: int, - page_size: int, n_kv_heads: int, head_dim: int, - device: torch.device, - dtype: torch.dtype, - ): - self.page_size = page_size - self._pool = PagePool(Allocator(n_pages), PrefixCache(page_size)) - self._table = TaskTable(page_size) - self._storage = Storage( - n_layers, n_pages, page_size, n_kv_heads, head_dim, device, dtype - ) - - def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: - hits = self._pool.lookup(prompt_ids) - cached = len(hits) * self.page_size - for p in hits: - self._pool.inc_ref(p) - - remaining = len(prompt_ids) - cached - n_new = ( - (remaining + self.page_size - 1) // self.page_size if remaining > 0 else 0 - ) - new_pages: List[int] = [] - if n_new > 0: - for _ in range(n_new): - p = self._pool.alloc() - if p < 0: - for hp in hits: - self._pool.free(hp) - for np in new_pages: - self._pool.free(np) - return False - new_pages.append(p) - - self._table.set(task_id, hits + new_pages, cached) - return True - - def task_free(self, task_id: str): - page_table, _ = self._table.pop(task_id) - for idx in page_table: - self._pool.free(idx) - - def task_extend(self, task_id: str, pos: int) -> bool: - page_table = self._table.get(task_id) - needed = (pos + 1 + self.page_size - 1) // self.page_size - while len(page_table) < needed: - p = self._pool.alloc() - if p < 0: - return False - page_table.append(p) - return True - - def task_cached(self, task_id: str) -> int: - return self._table.get_cached(task_id) - - def task_record_hashes( - self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0 - ): - page_table = self._table.get(task_id) - full_pages = len(prompt_ids) // self.page_size - for i in range(start_logical_page, full_pages): - self._pool.record(page_table[i], prompt_ids, i) - - def bind_tasks( - self, - task_ids: List[str], - total_len: int, - device: torch.device, - write_positions: Optional[Tensor] = None, - ) -> PageCacheView: - page_table = self._table.table_tensor(task_ids, device) - return PageCacheView(self._storage, page_table, total_len) - - -class ContiguousCacheView(CacheView): - """Contiguous KV-cache view for attention layers.""" - - def __init__( - self, - cache: "ContiguousCache", - batch_indices: Tensor, - total_len: int = 0, - write_positions: Optional[Tensor] = None, - ): - self._cache = cache - self._batch_indices = batch_indices - self._total_len = total_len - self._write_positions = write_positions - - def write(self, layer_id: int, k: Tensor, v: Tensor): - seq_len = k.size(1) - indices = self._batch_indices - if self._write_positions is not None and seq_len == 1: - pos = self._write_positions - self._cache.k[layer_id, indices, pos] = k.squeeze(1) - self._cache.v[layer_id, indices, pos] = v.squeeze(1) - else: - start_pos = self._total_len - seq_len - self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k - self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v - - def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: - max_len = self._total_len - indices = self._batch_indices - k = self._cache.k[layer_id, indices, :max_len] - v = self._cache.v[layer_id, indices, :max_len] - return k, v - - -class ContiguousCache(KVCache): - """Contiguous per-slot KV cache (default implementation).""" - - def __init__( - self, - n_layers: int, max_batch_size: int, max_seq_len: int, - n_kv_heads: int, - head_dim: int, device: torch.device, dtype: torch.dtype, + page_size: int = 1, + n_tokens: Optional[int] = None, ): + self.page_size = page_size + self.max_batch_size = max_batch_size self.max_seq_len = max_seq_len - self.k = torch.zeros( - n_layers, - max_batch_size, - max_seq_len, - n_kv_heads, - head_dim, - device=device, - dtype=dtype, + self.device = device + self.dtype = dtype + self.n_layers = n_layers + self.n_kv_heads = n_kv_heads + self.head_dim = head_dim + + self.contiguous = n_tokens is None + if self.contiguous: + self.n_tokens = max_batch_size * max_seq_len + else: + self.n_tokens = n_tokens + + self._storage = KVStorage( + self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype ) - self.v = torch.zeros( - n_layers, - max_batch_size, - max_seq_len, - n_kv_heads, - head_dim, - device=device, - dtype=dtype, - ) - self._slot_len: Dict[int, int] = {} - self._task_slot: Dict[str, int] = {} - self._free_slots = list(range(max_batch_size)) - self._device = device + self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device) + + if self.contiguous: + for i in range(max_batch_size): + self._req_pool.req_to_token[i] = torch.arange( + i * max_seq_len, (i + 1) * max_seq_len, device=device + ) + self._alloc: Optional[Allocator] = None + self._prefix: Optional[PrefixCache] = None + else: + n_pages = self.n_tokens // page_size + self._alloc = Allocator(n_pages) + self._prefix = PrefixCache(page_size) if page_size > 1 else None + if self._prefix is not None: + self._alloc.on_evict = self._prefix.evict + + self._task_req: Dict[str, int] = {} + self._task_len: Dict[int, int] = {} + self._task_cached: Dict[str, int] = {} + self._task_slots: Dict[str, List[int]] = {} + self._task_pages: Dict[str, List[int]] = {} + self._lock = threading.Lock() + + # ---- task lifecycle ---- def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: - if not self._free_slots: + req_slots = self._req_pool.alloc(1) + if req_slots is None: return False - slot = self._free_slots.pop(0) - self._task_slot[task_id] = slot - self._slot_len[slot] = 0 + req_idx = req_slots[0] + self._task_req[task_id] = req_idx + + if self.contiguous: + self._task_len[req_idx] = len(prompt_ids) + self._task_cached[task_id] = 0 + return True + + n_tokens_needed = len(prompt_ids) + cached = 0 + + if self._prefix is not None: + hits = self._prefix.lookup(prompt_ids) + cached = len(hits) * self.page_size + for p in hits: + self._alloc.inc_ref(p) + self._task_pages[task_id] = list(hits) + self._task_slots[task_id] = [] + else: + self._task_pages[task_id] = [] + self._task_slots[task_id] = [] + + remaining = n_tokens_needed - cached + if remaining > 0: + if self.page_size == 1: + slots = self._alloc_tokens(remaining) + if slots is None: + for p in self._task_pages[task_id]: + self._alloc.free(p) + self._req_pool.free([req_idx]) + del self._task_req[task_id] + return False + self._task_slots[task_id] = slots + else: + n_new_pages = (remaining + self.page_size - 1) // self.page_size + new_pages = [] + for _ in range(n_new_pages): + p = self._alloc.alloc() + if p < 0: + for hp in self._task_pages[task_id]: + self._alloc.free(hp) + for np_ in new_pages: + self._alloc.free(np_) + self._req_pool.free([req_idx]) + del self._task_req[task_id] + return False + new_pages.append(p) + self._task_pages[task_id].extend(new_pages) + + self._write_req_to_token(task_id, prompt_ids, cached) + self._task_len[req_idx] = len(prompt_ids) + self._task_cached[task_id] = cached return True def task_free(self, task_id: str): - slot = self._task_slot.pop(task_id, None) - if slot is not None: - self._slot_len.pop(slot, None) - self._free_slots.append(slot) + req_idx = self._task_req.pop(task_id, None) + if req_idx is None: + return + self._task_len.pop(req_idx, None) + self._task_cached.pop(task_id, None) + + if not self.contiguous: + if self._prefix is not None: + for p in self._task_pages.get(task_id, []): + keep = self._prefix.has_page(p) + self._alloc.free(p, keep_cached=keep) + if not keep: + self._prefix.evict(p) + else: + for p in self._task_pages.get(task_id, []): + self._alloc.free(p) + self._task_pages.pop(task_id, None) + self._task_slots.pop(task_id, None) + + self._req_pool.free([req_idx]) def task_extend(self, task_id: str, pos: int) -> bool: - return pos < self.max_seq_len + req_idx = self._task_req.get(task_id) + if req_idx is None: + return False + + if self.contiguous: + return pos < self.max_seq_len + + if self.page_size == 1: + slots = self._alloc_tokens(1) + if slots is None: + return False + self._task_slots.setdefault(task_id, []).extend(slots) + self._req_pool.req_to_token[req_idx, pos] = slots[0] + else: + page_idx = pos // self.page_size + existing = self._task_pages.get(task_id, []) + if page_idx >= len(existing): + p = self._alloc.alloc() + if p < 0: + return False + existing.append(p) + self._task_pages[task_id] = existing + page_offset = pos % self.page_size + page = existing[page_idx] + token_slot = page * self.page_size + page_offset + self._req_pool.req_to_token[req_idx, pos] = token_slot + + self._task_len[req_idx] = pos + 1 + return True def task_cached(self, task_id: str) -> int: - slot = self._task_slot.get(task_id) - if slot is None: - return 0 - return self._slot_len.get(slot, 0) + return self._task_cached.get(task_id, 0) + + def task_record_hashes( + self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0 + ): + if self._prefix is None or self.contiguous: + return + pages = self._task_pages.get(task_id, []) + full_pages = len(prompt_ids) // self.page_size + for i in range(start_logical_page, min(full_pages, len(pages))): + self._prefix.record(pages[i], prompt_ids, i) + + # ---- bind for forward ---- def bind_tasks( self, task_ids: List[str], - total_len: int, + seq_lens: List[int], device: torch.device, - write_positions: Optional[Tensor] = None, - ) -> ContiguousCacheView: - slots = [self._task_slot[tid] for tid in task_ids] - batch_indices = torch.tensor(slots, dtype=torch.long, device=device) - for slot in slots: - if total_len > self._slot_len.get(slot, 0): - self._slot_len[slot] = total_len - return ContiguousCacheView( - self, batch_indices, total_len, write_positions=write_positions + start_pos: Optional[int] = None, + ) -> KVCache: + req_indices = [self._task_req[tid] for tid in task_ids] + req_pool_indices = torch.tensor(req_indices, dtype=torch.long, device=device) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.long, device=device) + + if start_pos is not None: + seq_len = seq_lens[0] + out_cache_loc = self._req_pool.req_to_token[ + req_pool_indices, start_pos:seq_len + ] + else: + write_pos = seq_lens_t - 1 + out_cache_loc = self._req_pool.req_to_token[ + req_pool_indices, write_pos + ].unsqueeze(-1) + + return KVCache( + k_buffer=self._storage.k_buffer, + v_buffer=self._storage.v_buffer, + req_to_token=self._req_pool.req_to_token, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens_t, + out_cache_loc=out_cache_loc, ) + + # ---- internals ---- + + def _alloc_tokens(self, n: int) -> Optional[List[int]]: + if self.page_size != 1: + raise RuntimeError("_alloc_tokens is for page_size=1 only") + slots = [] + for _ in range(n): + p = self._alloc.alloc() + if p < 0: + for s in slots: + self._alloc.free(s) + return None + slots.append(p) + return slots + + def _write_req_to_token(self, task_id: str, prompt_ids: List[int], cached: int): + req_idx = self._task_req[task_id] + total = len(prompt_ids) + + if self.contiguous: + return + + if self.page_size == 1: + slots = self._task_slots.get(task_id, []) + all_slots = slots[: total - cached] + if all_slots: + self._req_pool.req_to_token[req_idx, cached:total] = torch.tensor( + all_slots, dtype=torch.long, device=self.device + ) + else: + pages = self._task_pages.get(task_id, []) + for pos in range(cached, total): + page_idx = pos // self.page_size + page_offset = pos % self.page_size + if page_idx < len(pages): + token_slot = pages[page_idx] * self.page_size + page_offset + self._req_pool.req_to_token[req_idx, pos] = token_slot diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index 5c3d61a..8b13903 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -3,7 +3,7 @@ from typing import List, Optional import torch -from astrai.inference.core.cache import KVCache +from astrai.inference.core.cache import PagePool from astrai.inference.core.task import Task from astrai.inference.sample import sample from astrai.model.automodel import AutoModel @@ -19,7 +19,7 @@ class Executor: self, model: AutoModel, tokenizer: AutoTokenizer, - kv_cache: KVCache, + kv_cache: PagePool, device: Optional[str] = None, dtype: Optional[torch.dtype] = None, ): @@ -57,7 +57,9 @@ class Executor: input_ids, input_mask=input_mask, position_ids=position_ids, - paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device), + kv_cache=self.kv_cache.bind_tasks( + task_ids, [prompt_len] * batch_sz, self.device, start_pos=start_pos + ), ) def execute_decode( @@ -128,11 +130,10 @@ class Executor: outputs = self.model( input_ids.unsqueeze(1), input_mask=input_mask, - paged_cache=self.kv_cache.bind_tasks( + kv_cache=self.kv_cache.bind_tasks( task_ids, - total_len, + [t.next_pos + 1 for t in tasks], self.device, - write_positions=position_ids, ), position_ids=position_ids.unsqueeze(1), ) diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index 3b78b6e..dfeadd6 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional, Tuple import torch -from astrai.inference.core.cache import ContiguousCache, KVCache +from astrai.inference.core.cache import PagePool from astrai.inference.core.executor import Executor from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus from astrai.model.automodel import AutoModel @@ -25,7 +25,7 @@ class InferenceScheduler: max_seq_len: Optional[int] = None, device: Optional[str] = None, dtype: Optional[torch.dtype] = None, - cache: Optional[KVCache] = None, + cache: Optional[PagePool] = None, ): config = model.config @@ -46,14 +46,14 @@ class InferenceScheduler: if cache is not None: self._cache = cache else: - self._cache = ContiguousCache( - config.num_hidden_layers, - max_batch_size, - self.max_seq_len, - config.num_key_value_heads, - head_dim, - self.device, - self.dtype, + self._cache = PagePool( + n_layers=config.num_hidden_layers, + n_kv_heads=config.num_key_value_heads, + head_dim=head_dim, + max_batch_size=max_batch_size, + max_seq_len=self.max_seq_len, + device=self.device, + dtype=self.dtype, ) self._task_mgr = TaskManager( diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index 7751c0e..7b6f412 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -8,7 +8,7 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, import torch import torch.nn as nn -from astrai.inference.core.cache import KVCache +from astrai.inference.core.cache import PagePool from astrai.inference.core.scheduler import InferenceScheduler from astrai.inference.core.task import STOP from astrai.tokenize import AutoTokenizer @@ -111,7 +111,7 @@ class InferenceEngine: tokenizer: AutoTokenizer, max_batch_size: int = 1, max_seq_len: Optional[int] = None, - cache: Optional[KVCache] = None, + cache: Optional[PagePool] = None, ): self.model = model self.tokenizer = tokenizer diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py index 7467395..1fb0acd 100644 --- a/astrai/model/components/attention.py +++ b/astrai/model/components/attention.py @@ -6,7 +6,7 @@ import torch.nn.functional as F from torch import Tensor from astrai.factory import BaseFactory -from astrai.inference.core.cache import CacheView +from astrai.inference.core.cache import KVCache from astrai.model.components.linear import Linear from astrai.model.components.norm import RMSNorm from astrai.model.components.rope import apply_rotary_emb @@ -75,7 +75,7 @@ class GQA(nn.Module): x: Tensor, rotary_emb: Tensor, attn_mask: Tensor = None, - paged_cache: Optional[CacheView] = None, + kv_cache: Optional[KVCache] = None, is_causal: bool = False, ) -> Tensor: q = self._split_heads(self.q_proj(x), self.n_heads) @@ -86,9 +86,19 @@ class GQA(nn.Module): if self.use_qk_norm: q, k = self.q_norm(q), self.k_norm(k) - if paged_cache is not None: - paged_cache.write(self.layer_id, k, v) - k, v = paged_cache.gather(self.layer_id) + if kv_cache is not None: + kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k + kv_cache.v_buffer[self.layer_id][kv_cache.out_cache_loc] = v + + max_len = kv_cache.seq_lens.max() + indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len] + pos_mask = ( + torch.arange(max_len, device=x.device)[None, :] + < kv_cache.seq_lens[:, None] + ) + indices = torch.where(pos_mask, indices, torch.zeros_like(indices)) + k = kv_cache.k_buffer[self.layer_id][indices] + v = kv_cache.v_buffer[self.layer_id][indices] k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep) @@ -161,7 +171,7 @@ class MLA(nn.Module): x: Tensor, rotary_emb: Tensor, attn_mask: Tensor = None, - paged_cache: Optional[CacheView] = None, + kv_cache: Optional[KVCache] = None, is_causal: bool = False, ) -> Tensor: bsz, seq_len, _ = x.size() @@ -193,9 +203,19 @@ class MLA(nn.Module): q = self.q_norm(q) k = self.k_norm(k) - if paged_cache is not None: - paged_cache.write(self.layer_id, k, v) - k, v = paged_cache.gather(self.layer_id) + if kv_cache is not None: + kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k + kv_cache.v_buffer[self.layer_id][kv_cache.out_cache_loc] = v + + max_len = kv_cache.seq_lens.max() + indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len] + pos_mask = ( + torch.arange(max_len, device=x.device)[None, :] + < kv_cache.seq_lens[:, None] + ) + indices = torch.where(pos_mask, indices, torch.zeros_like(indices)) + k = kv_cache.k_buffer[self.layer_id][indices] + v = kv_cache.v_buffer[self.layer_id][indices] q = q.permute(0, 2, 1, 3) k = k.permute(0, 2, 1, 3) diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py index 7abc96f..2c1d880 100644 --- a/astrai/model/components/decoder_block.py +++ b/astrai/model/components/decoder_block.py @@ -4,7 +4,7 @@ from typing import Optional import torch.nn as nn from torch import Tensor -from astrai.inference.core.cache import CacheView +from astrai.inference.core.cache import KVCache from astrai.model.components.attention import AttnFactory from astrai.model.components.mlp import FFNFactory from astrai.model.components.norm import RMSNorm @@ -33,14 +33,14 @@ class DecoderBlock(nn.Module): x: Tensor, rotary_emb: Tensor, attention_mask: Optional[Tensor] = None, - paged_cache: Optional[CacheView] = None, + kv_cache: Optional[KVCache] = None, is_causal: bool = False, ) -> Tensor: attn_output = self.attention( self.input_norm(x), rotary_emb, attention_mask, - paged_cache, + kv_cache, is_causal, ) x = attn_output + x diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index 902d8e6..38b278b 100644 --- a/astrai/model/transformer.py +++ b/astrai/model/transformer.py @@ -5,7 +5,7 @@ import torch.nn as nn from torch import Tensor from astrai.config.model_config import AutoRegressiveLMConfig -from astrai.inference.core.cache import CacheView +from astrai.inference.core.cache import KVCache from astrai.model.automodel import AutoModel, ModelFactory from astrai.model.components.decoder_block import DecoderBlock from astrai.model.components.embedding import Embedding @@ -103,7 +103,7 @@ class AutoRegressiveLM(AutoModel): self, input_ids: Tensor, input_mask: Optional[Tensor] = None, - paged_cache: Optional[CacheView] = None, + kv_cache: Optional[KVCache] = None, position_ids: Optional[Tensor] = None, ) -> Dict[str, Tensor]: assert input_ids.ndim == 2 @@ -114,7 +114,7 @@ class AutoRegressiveLM(AutoModel): use_sdpa_causal_mask = attn_mask is None for layer in self.layers: - x = layer(x, rotary_emb, attn_mask, paged_cache, use_sdpa_causal_mask) + x = layer(x, rotary_emb, attn_mask, kv_cache, use_sdpa_causal_mask) hidden_states = self.norm(x) logits = self.lm_head(hidden_states) diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index e6f5dbb..4cd5820 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -4,17 +4,14 @@ import torch from astrai.inference import ( Allocator, - PageCache, + KVStorage, PagePool, PrefixCache, - Storage, - TaskTable, + ReqToTokenPool, page_hash, ) - -def make_pool(n_pages: int, page_size: int) -> PagePool: - return PagePool(Allocator(n_pages), PrefixCache(page_size)) +# ---- page_hash ---- def test_page_hash_full_page(): @@ -29,251 +26,335 @@ def test_page_hash_different_page_differs(): assert page_hash(token_ids, 0, 64) != page_hash(token_ids, 1, 64) -def test_page_pool_alloc_free_cycle(): - pool = make_pool(4, 64) - a = pool.alloc() - b = pool.alloc() +# ---- Allocator ---- + + +def test_allocator_alloc_free_cycle(): + alloc = Allocator(4) + a = alloc.alloc() + b = alloc.alloc() assert a != b - pool.free(a) - pool.free(b) - c = pool.alloc() + alloc.free(a) + alloc.free(b) + c = alloc.alloc() assert c in (a, b) -def test_page_pool_alloc_when_full(): - pool = make_pool(2, 64) - pool.alloc() - pool.alloc() - assert pool.alloc() == -1 +def test_allocator_alloc_when_full(): + alloc = Allocator(2) + alloc.alloc() + alloc.alloc() + assert alloc.alloc() == -1 -def test_page_pool_lru_eviction(): - pool = make_pool(2, 64) - p0 = pool.alloc() - p1 = pool.alloc() - pool.record(p0, list(range(64)), 0) - pool.record(p1, list(range(64, 128)), 0) - pool.free(p0) - pool.free(p1) - pool.alloc() - assert p0 in pool._alloc._lru or p1 in pool._alloc._lru +def test_allocator_lru_eviction(): + alloc = Allocator(2) + p0 = alloc.alloc() + p1 = alloc.alloc() + alloc.free(p0, keep_cached=True) + alloc.free(p1, keep_cached=True) + alloc.alloc() + assert p0 in alloc._lru or p1 in alloc._lru -def test_page_pool_inc_ref_and_free(): - pool = make_pool(2, 64) - p = pool.alloc() - pool.inc_ref(p) - assert pool._alloc._refs[p] == 2 - pool.free(p) - assert pool._alloc._refs[p] == 1 - pool.free(p) - assert pool._alloc._refs[p] == 0 +def test_allocator_inc_ref_and_free(): + alloc = Allocator(2) + p = alloc.alloc() + alloc.inc_ref(p) + assert alloc._refs[p] == 2 + alloc.free(p) + assert alloc._refs[p] == 1 + alloc.free(p) + assert alloc._refs[p] == 0 -def test_page_pool_keep_cached_realloc(): - """Free mask has priority over LRU; cached page returned only when no free pages.""" - pool = make_pool(3, 64) - p0 = pool.alloc() - p1 = pool.alloc() - p2 = pool.alloc() - for p in (p0, p1, p2): - pool.record(p, [p] * 64, 0) - pool.free(p0) - pool.free(p1) - pool.free(p2) - assert pool.alloc() == p0 +# ---- PrefixCache ---- def test_prefix_cache_lookup_returns_hits(): token_ids = list(range(256)) - pool = make_pool(16, 64) - pages = [pool.alloc() for _ in range(4)] + prefix = PrefixCache(64) + pages = [0, 1, 2, 3] for i, p in enumerate(pages): - pool.record(p, token_ids, i) - pool.free(p) - hits = pool.lookup(token_ids) + prefix.record(p, token_ids, i) + hits = prefix.lookup(token_ids) assert hits == pages def test_prefix_cache_lookup_stops_at_first_miss(): token_ids = list(range(256)) - pool = make_pool(16, 64) - p0 = pool.alloc() - pool.record(p0, token_ids, 0) - pool.free(p0) - p1 = pool.alloc() - pool.record(p1, [99] * 64, 1) - pool.free(p1) - hits = pool.lookup(token_ids) + prefix = PrefixCache(64) + prefix.record(0, token_ids, 0) + prefix.record(1, [99] * 64, 1) + hits = prefix.lookup(token_ids) assert len(hits) == 1 - assert hits[0] == p0 + assert hits[0] == 0 def test_prefix_cache_ignores_partial_last_page(): token_ids = list(range(100)) - pool = make_pool(16, 64) - p = pool.alloc() - pool.record(p, token_ids, 0) - pool.free(p) - hits = pool.lookup(token_ids) + prefix = PrefixCache(64) + prefix.record(0, token_ids, 0) + hits = prefix.lookup(token_ids) assert len(hits) == 1 def test_prefix_cache_on_evict_clears_mappings(): - pool = make_pool(4, 64) - p = pool.alloc() - pool.record(p, list(range(64)), 0) - pool.free(p) - assert p in pool._prefix._page_to_hash - pool._prefix.evict(p) - assert p not in pool._prefix._page_to_hash + prefix = PrefixCache(64) + prefix.record(0, list(range(64)), 0) + assert 0 in prefix._page_to_hash + prefix.evict(0) + assert 0 not in prefix._page_to_hash def test_prefix_cache_has_page(): - pool = make_pool(4, 64) - p = pool.alloc() - assert p not in pool._prefix._page_to_hash - pool.record(p, list(range(64)), 0) - pool.free(p) - assert p in pool._prefix._page_to_hash + prefix = PrefixCache(64) + assert not prefix.has_page(0) + prefix.record(0, list(range(64)), 0) + assert prefix.has_page(0) -def test_task_table_set_get(): - table = TaskTable(page_size=64) - table.set("task1", [0, 1, 2], 128) - assert table.get("task1") == [0, 1, 2] - assert table.get_cached("task1") == 128 +# ---- ReqToTokenPool ---- -def test_task_table_get_missing(): - table = TaskTable(page_size=64) - assert table.get("nonexistent") == [] - assert table.get_cached("nonexistent") == 0 +def test_req_to_token_pool_alloc_free(): + pool = ReqToTokenPool(4, 128, torch.device("cpu")) + slots = pool.alloc(2) + assert len(slots) == 2 + assert len(pool.free_slots) == 2 + pool.free(slots) + assert len(pool.free_slots) == 4 -def test_task_table_pop(): - table = TaskTable(page_size=64) - table.set("task1", [0, 1], 64) - pages, cached = table.pop("task1") - assert pages == [0, 1] - assert cached == 64 - assert table.get("task1") == [] +def test_req_to_token_pool_alloc_when_full(): + pool = ReqToTokenPool(2, 128, torch.device("cpu")) + pool.alloc(2) + assert pool.alloc(1) is None -def test_kv_cache_task_extend_allocates(): - cache = PageCache( - n_layers=1, - n_pages=8, - page_size=64, - n_kv_heads=2, - head_dim=8, - device=torch.device("cpu"), - dtype=torch.float32, - ) - cache._table.set("task1", [], 0) - ok = cache.task_extend("task1", 200) - assert ok - assert len(cache._table.get("task1")) == 4 +def test_req_to_token_pool_write(): + pool = ReqToTokenPool(4, 128, torch.device("cpu")) + slots = pool.alloc(1) + pool.write((slots[0], slice(0, 3)), torch.tensor([10, 20, 30])) + assert pool.req_to_token[slots[0], 0].item() == 10 + assert pool.req_to_token[slots[0], 2].item() == 30 -def test_kv_cache_task_extend_fails_when_pool_full(): - cache = PageCache( - n_layers=1, - n_pages=2, - page_size=64, - n_kv_heads=2, - head_dim=8, - device=torch.device("cpu"), - dtype=torch.float32, - ) - cache._table.set("task1", [0, 1], 0) - ok = cache.task_extend("task1", 300) - assert not ok +# ---- KVStorage ---- -def test_task_table_table_tensor(): - table = TaskTable(page_size=64) - table.set("a", [0, 1], 0) - table.set("b", [2, 3, 4], 0) - t = table.table_tensor(["a", "b"], torch.device("cpu")) - assert t.shape == (2, 3) - assert t[0].tolist() == [0, 1, -1] - assert t[1].tolist() == [2, 3, 4] - - -def test_task_table_table_tensor_empty_input(): - table = TaskTable(page_size=64) - t = table.table_tensor([], torch.device("cpu")) - assert t.numel() == 0 - - -def test_storage_write_gather_single_page(): - storage = Storage( +def test_kv_storage_set_and_get(): + storage = KVStorage( + size=16, n_layers=2, - n_pages=8, - page_size=4, - n_kv_heads=2, + n_kv_heads=4, head_dim=8, device=torch.device("cpu"), dtype=torch.float32, ) - page_table = torch.tensor([[0]], dtype=torch.long) - k = torch.randn(1, 2, 2, 8) - v = torch.randn(1, 2, 2, 8) - - storage.write(0, page_table, 0, k, v) - gk, gv = storage.gather(0, page_table, 2) - assert torch.allclose(gk, k) + loc = torch.tensor([[0, 1]], dtype=torch.long) + k = torch.randn(1, 2, 4, 8) + v = torch.randn(1, 2, 4, 8) + storage.set_kv_buffer(0, loc, k, v) + assert torch.allclose(storage.get_key_buffer(0)[loc], k) + assert torch.allclose(storage.get_value_buffer(0)[loc], v) -def test_storage_write_cross_page(): - storage = Storage( +def test_kv_storage_buffer_shape(): + storage = KVStorage( + size=32, + n_layers=3, + n_kv_heads=8, + head_dim=16, + device=torch.device("cpu"), + dtype=torch.float32, + ) + assert storage.k_buffer.shape == (3, 32, 8, 16) + assert storage.v_buffer.shape == (3, 32, 8, 16) + + +# ---- PagePool (contiguous mode) ---- + + +def _make_contiguous_pool(**kwargs): + defaults = dict( + n_layers=2, + n_kv_heads=4, + head_dim=8, + max_batch_size=4, + max_seq_len=64, + device=torch.device("cpu"), + dtype=torch.float32, + ) + defaults.update(kwargs) + return PagePool(**defaults) + + +def test_page_pool_contiguous_task_alloc_free(): + pool = _make_contiguous_pool() + assert pool.task_alloc("t1", [1, 2, 3]) + assert "t1" in pool._task_req + pool.task_free("t1") + assert "t1" not in pool._task_req + + +def test_page_pool_contiguous_task_extend(): + pool = _make_contiguous_pool() + pool.task_alloc("t1", [1, 2, 3]) + assert pool.task_extend("t1", 3) + assert pool.task_extend("t1", 63) + assert not pool.task_extend("t1", 64) + + +def test_page_pool_contiguous_task_cached(): + pool = _make_contiguous_pool() + pool.task_alloc("t1", [1, 2, 3]) + assert pool.task_cached("t1") == 0 + + +def test_page_pool_contiguous_bind_tasks_prefill(): + pool = _make_contiguous_pool() + pool.task_alloc("t1", list(range(10))) + pool.task_alloc("t2", list(range(10))) + kv = pool.bind_tasks(["t1", "t2"], [10, 10], torch.device("cpu"), start_pos=0) + assert kv.out_cache_loc.shape == (2, 10) + assert kv.seq_lens.tolist() == [10, 10] + assert kv.req_pool_indices.shape == (2,) + + +def test_page_pool_contiguous_bind_tasks_decode(): + pool = _make_contiguous_pool() + pool.task_alloc("t1", list(range(10))) + pool.task_alloc("t2", list(range(8))) + kv = pool.bind_tasks(["t1", "t2"], [11, 9], torch.device("cpu")) + assert kv.out_cache_loc.shape == (2, 1) + assert kv.seq_lens.tolist() == [11, 9] + + +def test_page_pool_contiguous_bind_roundtrip(): + """Write KV via bind_tasks, then gather via req_to_token indexing.""" + pool = _make_contiguous_pool(n_layers=1, n_kv_heads=2, head_dim=4) + pool.task_alloc("t1", list(range(4))) + + kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0) + k = torch.randn(1, 4, 2, 4) + v = torch.randn(1, 4, 2, 4) + kv.k_buffer[0][kv.out_cache_loc] = k + kv.v_buffer[0][kv.out_cache_loc] = v + + indices = kv.req_to_token[kv.req_pool_indices, :4] + gathered_k = kv.k_buffer[0][indices] + gathered_v = kv.v_buffer[0][indices] + assert torch.allclose(gathered_k, k) + assert torch.allclose(gathered_v, v) + + +# ---- PagePool (paged mode, page_size=1) ---- + + +def _make_paged_pool(**kwargs): + defaults = dict( n_layers=1, - n_pages=8, - page_size=4, n_kv_heads=2, - head_dim=8, + head_dim=4, + max_batch_size=4, + max_seq_len=64, device=torch.device("cpu"), dtype=torch.float32, + page_size=1, + n_tokens=128, ) - page_table = torch.tensor([[0, 1]], dtype=torch.long) - k = torch.randn(1, 8, 2, 8) - v = torch.randn(1, 8, 2, 8) - - storage.write(0, page_table, 0, k, v) - gk, gv = storage.gather(0, page_table, 8) - assert torch.allclose(gk, k) + defaults.update(kwargs) + return PagePool(**defaults) -def test_storage_gather_truncates_to_total_len(): - storage = Storage( +def test_page_pool_paged_task_alloc(): + pool = _make_paged_pool() + assert pool.task_alloc("t1", list(range(10))) + req_idx = pool._task_req["t1"] + slots = pool._task_slots["t1"] + assert len(slots) == 10 + assert pool._req_pool.req_to_token[req_idx, 0].item() == slots[0] + + +def test_page_pool_paged_task_extend(): + pool = _make_paged_pool() + pool.task_alloc("t1", list(range(4))) + assert pool.task_extend("t1", 4) + req_idx = pool._task_req["t1"] + slot = pool._req_pool.req_to_token[req_idx, 4].item() + assert slot >= 0 + + +def test_page_pool_paged_task_free_releases_slots(): + pool = _make_paged_pool(n_tokens=16) + pool.task_alloc("t1", list(range(8))) + pool.task_free("t1") + assert "t1" not in pool._task_req + assert len(pool._req_pool.free_slots) == 4 + + +def test_page_pool_paged_bind_roundtrip(): + pool = _make_paged_pool(n_layers=1, n_kv_heads=2, head_dim=4) + pool.task_alloc("t1", list(range(4))) + + kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0) + k = torch.randn(1, 4, 2, 4) + v = torch.randn(1, 4, 2, 4) + kv.k_buffer[0][kv.out_cache_loc] = k + kv.v_buffer[0][kv.out_cache_loc] = v + + indices = kv.req_to_token[kv.req_pool_indices, :4] + gathered_k = kv.k_buffer[0][indices] + assert torch.allclose(gathered_k, k) + + +# ---- PagePool (paged mode, page_size>1) ---- + + +def _make_paged_pool_ps64(**kwargs): + defaults = dict( n_layers=1, - n_pages=8, - page_size=4, n_kv_heads=2, - head_dim=8, + head_dim=4, + max_batch_size=4, + max_seq_len=256, device=torch.device("cpu"), dtype=torch.float32, + page_size=64, + n_tokens=512, ) - page_table = torch.tensor([[0, 1]], dtype=torch.long) - k = torch.randn(1, 6, 2, 8) - v = torch.randn(1, 6, 2, 8) - storage.write(0, page_table, 0, k, v) - - gk, gv = storage.gather(0, page_table, 5) - assert gk.shape == (1, 5, 2, 8) + defaults.update(kwargs) + return PagePool(**defaults) -def test_storage_gather_clamps_negative_padding(): - storage = Storage( - n_layers=1, - n_pages=8, - page_size=4, - n_kv_heads=2, - head_dim=8, - device=torch.device("cpu"), - dtype=torch.float32, - ) - page_table = torch.tensor([[0, -1]], dtype=torch.long) - gk, gv = storage.gather(0, page_table, 4) - assert gk.shape == (1, 4, 2, 8) +def test_page_pool_paged_ps64_task_alloc(): + pool = _make_paged_pool_ps64() + prompt = list(range(200)) + assert pool.task_alloc("t1", prompt) + assert pool.task_cached("t1") == 0 + n_pages = (200 + 63) // 64 + assert len(pool._task_pages["t1"]) == n_pages + + +def test_page_pool_paged_ps64_task_extend_crosses_page(): + pool = _make_paged_pool_ps64() + pool.task_alloc("t1", list(range(64))) + assert pool.task_extend("t1", 64) + assert len(pool._task_pages["t1"]) >= 2 + + +def test_page_pool_paged_ps64_bind_roundtrip(): + pool = _make_paged_pool_ps64(n_layers=1, n_kv_heads=2, head_dim=4) + prompt = list(range(128)) + pool.task_alloc("t1", prompt) + + kv = pool.bind_tasks(["t1"], [128], torch.device("cpu"), start_pos=0) + k = torch.randn(1, 128, 2, 4) + v = torch.randn(1, 128, 2, 4) + kv.k_buffer[0][kv.out_cache_loc] = k + kv.v_buffer[0][kv.out_cache_loc] = v + + indices = kv.req_to_token[kv.req_pool_indices, :128] + gathered_k = kv.k_buffer[0][indices] + assert torch.allclose(gathered_k, k)