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

@@ -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,
+    KVStorage,
-    PageCacheView,
     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",
+    "KVStorage",
-    "ContiguousCacheView",
-    "PageCache",
-    "PageCacheView",
     "PagePool",
     "PrefixCache",
-    "Storage",
+    "ReqToTokenPool",
-    "TaskTable",
     "page_hash",
     "sample",
     "BaseSamplingStrategy",

astrai/inference/core/__init__.py

@@ -2,16 +2,11 @@
 
 from astrai.inference.core.cache import (
     Allocator,
-    CacheView,
-    ContiguousCache,
-    ContiguousCacheView,
     KVCache,
-    PageCache,
+    KVStorage,
-    PageCacheView,
     PagePool,
     PrefixCache,
-    Storage,
+    ReqToTokenPool,
-    TaskTable,
     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",
+    "KVStorage",
-    "ContiguousCacheView",
-    "PageCache",
-    "PageCacheView",
     "PagePool",
     "PrefixCache",
-    "Storage",
+    "ReqToTokenPool",
-    "TaskTable",
     "page_hash",
     "Executor",
     "InferenceScheduler",

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:
+class ReqToTokenPool:
-    """Orchestrates allocator (page management) and PrefixCache (content addressing)."""
+    """Maps [req_idx, pos] -> physical token slot in KV storage.
 
-    def __init__(self, allocator: Allocator, prefix: PrefixCache):
+    Each row is one request; each column is a sequence position. The value
-        self._alloc = allocator
+    at [req_idx, pos] is the flat index into the KV storage buffers.
-        self._prefix = prefix
+    """
-        self._alloc.on_evict = prefix.evict
 
-    @property
+    def __init__(self, size: int, max_context_len: int, device: torch.device):
-    def allocator(self) -> Allocator:
+        self.size = size
-        return self._alloc
+        self.max_context_len = max_context_len
-
+        self.req_to_token = torch.zeros(
-    @property
+            (size, max_context_len), dtype=torch.long, device=device
-    def prefix(self) -> PrefixCache:
+        )
-        return self._prefix
+        self.free_slots = list(range(size))
-
-    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] = {}
         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
+            if num_reqs > len(self.free_slots):
-            self._cached[task_id] = cached
+                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:
+    def write(self, indices, values):
-        with self._lock:
+        self.req_to_token[indices] = values
-            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)
 
 
-class Storage:
+class KVStorage:
-    """KV-cache tensor storage with paged write/gather."""
+    """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.size = size
-        self.k_cache = torch.empty(
+        self.k_buffer = torch.empty(
-            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
+            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
-            device=device,
-            dtype=dtype,
         )
-        self.v_cache = torch.empty(
+        self.v_buffer = torch.empty(
-            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
+            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
-            device=device,
-            dtype=dtype,
         )
 
-    def write(
+    def get_key_buffer(self, layer_id: int) -> Tensor:
-        self,
+        return self.k_buffer[layer_id]
-        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 gather(
+    def get_value_buffer(self, layer_id: int) -> Tensor:
-        self, layer_id: int, page_table: Tensor, total_len: int
+        return self.v_buffer[layer_id]
-    ) -> Tuple[Tensor, Tensor]:
+
-        safe = page_table.clamp(min=0)
+    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
-        k = self.k_cache[layer_id, safe]
+        self.k_buffer[layer_id][loc] = k
-        v = self.v_cache[layer_id, safe]
+        self.v_buffer[layer_id][loc] = v
-        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
 
 
-class CacheView(ABC):
+@dataclass
-    """Abstract view passed to attention layers for KV-cache I/O."""
+class KVCache:
+    """Pure data struct passed to model for KV cache I/O.
 
-    @abstractmethod
+    The attention layer does raw buffer indexing — no methods, no abstraction.
-    def write(self, layer_id: int, k: Tensor, v: Tensor): ...
 
-    @abstractmethod
+    Attributes:
-    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ...
+        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):
+class PagePool:
-    """Abstract KV-cache facade for scheduler/executor."""
+    """Top-level KV cache manager.
 
-    @abstractmethod
+    Combines KVStorage + ReqToTokenPool + Allocator + PrefixCache.
-    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ...
 
-    @abstractmethod
+    Args:
-    def task_free(self, task_id: str): ...
+        n_layers: Number of transformer layers.
-
+        n_kv_heads: Number of KV attention heads.
-    @abstractmethod
+        head_dim: Dimension per head.
-    def task_extend(self, task_id: str, pos: int) -> bool: ...
+        max_batch_size: Maximum concurrent requests.
-
+        max_seq_len: Maximum sequence length per request.
-    @abstractmethod
+        device, dtype: Tensor device and dtype.
-    def bind_tasks(
+        page_size: Page size for paged mode (1 = token-level).
-        self,
+        n_tokens: Total token slots for paged mode. None = contiguous mode
-        task_ids: List[str],
+            (pre-allocates max_batch_size * max_seq_len).
-        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."""
 
     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(
+        self.device = device
-            n_layers,
+        self.dtype = dtype
-            max_batch_size,
+        self.n_layers = n_layers
-            max_seq_len,
+        self.n_kv_heads = n_kv_heads
-            n_kv_heads,
+        self.head_dim = head_dim
-            head_dim,
+
-            device=device,
+        self.contiguous = n_tokens is None
-            dtype=dtype,
+        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(
+        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
-            n_layers,
+
-            max_batch_size,
+        if self.contiguous:
-            max_seq_len,
+            for i in range(max_batch_size):
-            n_kv_heads,
+                self._req_pool.req_to_token[i] = torch.arange(
-            head_dim,
+                    i * max_seq_len, (i + 1) * max_seq_len, device=device
-            device=device,
+                )
-            dtype=dtype,
+            self._alloc: Optional[Allocator] = None
-        )
+            self._prefix: Optional[PrefixCache] = None
-        self._slot_len: Dict[int, int] = {}
+        else:
-        self._task_slot: Dict[str, int] = {}
+            n_pages = self.n_tokens // page_size
-        self._free_slots = list(range(max_batch_size))
+            self._alloc = Allocator(n_pages)
-        self._device = device
+            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)
+        req_idx = req_slots[0]
-        self._task_slot[task_id] = slot
+        self._task_req[task_id] = req_idx
-        self._slot_len[slot] = 0
+
+        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)
+        req_idx = self._task_req.pop(task_id, None)
-        if slot is not None:
+        if req_idx is None:
-            self._slot_len.pop(slot, None)
+            return
-            self._free_slots.append(slot)
+        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)
+        return self._task_cached.get(task_id, 0)
-        if slot is None:
+
-            return 0
+    def task_record_hashes(
-        return self._slot_len.get(slot, 0)
+        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,
+        start_pos: Optional[int] = None,
-    ) -> ContiguousCacheView:
+    ) -> KVCache:
-        slots = [self._task_slot[tid] for tid in task_ids]
+        req_indices = [self._task_req[tid] for tid in task_ids]
-        batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
+        req_pool_indices = torch.tensor(req_indices, dtype=torch.long, device=device)
-        for slot in slots:
+        seq_lens_t = torch.tensor(seq_lens, dtype=torch.long, device=device)
-            if total_len > self._slot_len.get(slot, 0):
+
-                self._slot_len[slot] = total_len
+        if start_pos is not None:
-        return ContiguousCacheView(
+            seq_len = seq_lens[0]
-            self, batch_indices, total_len, write_positions=write_positions
+            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

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

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(
+            self._cache = PagePool(
-                config.num_hidden_layers,
+                n_layers=config.num_hidden_layers,
-                max_batch_size,
+                n_kv_heads=config.num_key_value_heads,
-                self.max_seq_len,
+                head_dim=head_dim,
-                config.num_key_value_heads,
+                max_batch_size=max_batch_size,
-                head_dim,
+                max_seq_len=self.max_seq_len,
-                self.device,
+                device=self.device,
-                self.dtype,
+                dtype=self.dtype,
             )
 
         self._task_mgr = TaskManager(

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

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:
+        if kv_cache is not None:
-            paged_cache.write(self.layer_id, k, v)
+            kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k
-            k, v = paged_cache.gather(self.layer_id)
+            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:
+        if kv_cache is not None:
-            paged_cache.write(self.layer_id, k, v)
+            kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k
-            k, v = paged_cache.gather(self.layer_id)
+            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)

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

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)

tests/inference/test_cache.py

@@ -4,17 +4,14 @@ import torch
 
 from astrai.inference import (
     Allocator,
-    PageCache,
+    KVStorage,
     PagePool,
     PrefixCache,
-    Storage,
+    ReqToTokenPool,
-    TaskTable,
     page_hash,
 )
 
-
+# ---- page_hash ----
-def make_pool(n_pages: int, page_size: int) -> PagePool:
-    return PagePool(Allocator(n_pages), PrefixCache(page_size))
 
 
 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():
+# ---- Allocator ----
-    pool = make_pool(4, 64)
+
-    a = pool.alloc()
+
-    b = pool.alloc()
+def test_allocator_alloc_free_cycle():
+    alloc = Allocator(4)
+    a = alloc.alloc()
+    b = alloc.alloc()
     assert a != b
-    pool.free(a)
+    alloc.free(a)
-    pool.free(b)
+    alloc.free(b)
-    c = pool.alloc()
+    c = alloc.alloc()
     assert c in (a, b)
 
 
-def test_page_pool_alloc_when_full():
+def test_allocator_alloc_when_full():
-    pool = make_pool(2, 64)
+    alloc = Allocator(2)
-    pool.alloc()
+    alloc.alloc()
-    pool.alloc()
+    alloc.alloc()
-    assert pool.alloc() == -1
+    assert alloc.alloc() == -1
 
 
-def test_page_pool_lru_eviction():
+def test_allocator_lru_eviction():
-    pool = make_pool(2, 64)
+    alloc = Allocator(2)
-    p0 = pool.alloc()
+    p0 = alloc.alloc()
-    p1 = pool.alloc()
+    p1 = alloc.alloc()
-    pool.record(p0, list(range(64)), 0)
+    alloc.free(p0, keep_cached=True)
-    pool.record(p1, list(range(64, 128)), 0)
+    alloc.free(p1, keep_cached=True)
-    pool.free(p0)
+    alloc.alloc()
-    pool.free(p1)
+    assert p0 in alloc._lru or p1 in alloc._lru
-    pool.alloc()
-    assert p0 in pool._alloc._lru or p1 in pool._alloc._lru
 
 
-def test_page_pool_inc_ref_and_free():
+def test_allocator_inc_ref_and_free():
-    pool = make_pool(2, 64)
+    alloc = Allocator(2)
-    p = pool.alloc()
+    p = alloc.alloc()
-    pool.inc_ref(p)
+    alloc.inc_ref(p)
-    assert pool._alloc._refs[p] == 2
+    assert alloc._refs[p] == 2
-    pool.free(p)
+    alloc.free(p)
-    assert pool._alloc._refs[p] == 1
+    assert alloc._refs[p] == 1
-    pool.free(p)
+    alloc.free(p)
-    assert pool._alloc._refs[p] == 0
+    assert alloc._refs[p] == 0
 
 
-def test_page_pool_keep_cached_realloc():
+# ---- PrefixCache ----
-    """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
 
 
 def test_prefix_cache_lookup_returns_hits():
     token_ids = list(range(256))
-    pool = make_pool(16, 64)
+    prefix = PrefixCache(64)
-    pages = [pool.alloc() for _ in range(4)]
+    pages = [0, 1, 2, 3]
     for i, p in enumerate(pages):
-        pool.record(p, token_ids, i)
+        prefix.record(p, token_ids, i)
-        pool.free(p)
+    hits = prefix.lookup(token_ids)
-    hits = pool.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)
+    prefix = PrefixCache(64)
-    p0 = pool.alloc()
+    prefix.record(0, token_ids, 0)
-    pool.record(p0, token_ids, 0)
+    prefix.record(1, [99] * 64, 1)
-    pool.free(p0)
+    hits = prefix.lookup(token_ids)
-    p1 = pool.alloc()
-    pool.record(p1, [99] * 64, 1)
-    pool.free(p1)
-    hits = pool.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)
+    prefix = PrefixCache(64)
-    p = pool.alloc()
+    prefix.record(0, token_ids, 0)
-    pool.record(p, token_ids, 0)
+    hits = prefix.lookup(token_ids)
-    pool.free(p)
-    hits = pool.lookup(token_ids)
     assert len(hits) == 1
 
 
 def test_prefix_cache_on_evict_clears_mappings():
-    pool = make_pool(4, 64)
+    prefix = PrefixCache(64)
-    p = pool.alloc()
+    prefix.record(0, list(range(64)), 0)
-    pool.record(p, list(range(64)), 0)
+    assert 0 in prefix._page_to_hash
-    pool.free(p)
+    prefix.evict(0)
-    assert p in pool._prefix._page_to_hash
+    assert 0 not in prefix._page_to_hash
-    pool._prefix.evict(p)
-    assert p not in pool._prefix._page_to_hash
 
 
 def test_prefix_cache_has_page():
-    pool = make_pool(4, 64)
+    prefix = PrefixCache(64)
-    p = pool.alloc()
+    assert not prefix.has_page(0)
-    assert p not in pool._prefix._page_to_hash
+    prefix.record(0, list(range(64)), 0)
-    pool.record(p, list(range(64)), 0)
+    assert prefix.has_page(0)
-    pool.free(p)
-    assert p in pool._prefix._page_to_hash
 
 
-def test_task_table_set_get():
+# ---- ReqToTokenPool ----
-    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
 
 
-def test_task_table_get_missing():
+def test_req_to_token_pool_alloc_free():
-    table = TaskTable(page_size=64)
+    pool = ReqToTokenPool(4, 128, torch.device("cpu"))
-    assert table.get("nonexistent") == []
+    slots = pool.alloc(2)
-    assert table.get_cached("nonexistent") == 0
+    assert len(slots) == 2
+    assert len(pool.free_slots) == 2
+    pool.free(slots)
+    assert len(pool.free_slots) == 4
 
 
-def test_task_table_pop():
+def test_req_to_token_pool_alloc_when_full():
-    table = TaskTable(page_size=64)
+    pool = ReqToTokenPool(2, 128, torch.device("cpu"))
-    table.set("task1", [0, 1], 64)
+    pool.alloc(2)
-    pages, cached = table.pop("task1")
+    assert pool.alloc(1) is None
-    assert pages == [0, 1]
-    assert cached == 64
-    assert table.get("task1") == []
 
 
-def test_kv_cache_task_extend_allocates():
+def test_req_to_token_pool_write():
-    cache = PageCache(
+    pool = ReqToTokenPool(4, 128, torch.device("cpu"))
-        n_layers=1,
+    slots = pool.alloc(1)
-        n_pages=8,
+    pool.write((slots[0], slice(0, 3)), torch.tensor([10, 20, 30]))
-        page_size=64,
+    assert pool.req_to_token[slots[0], 0].item() == 10
-        n_kv_heads=2,
+    assert pool.req_to_token[slots[0], 2].item() == 30
-        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_kv_cache_task_extend_fails_when_pool_full():
+# ---- KVStorage ----
-    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
 
 
-def test_task_table_table_tensor():
+def test_kv_storage_set_and_get():
-    table = TaskTable(page_size=64)
+    storage = KVStorage(
-    table.set("a", [0, 1], 0)
+        size=16,
-    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(
         n_layers=2,
-        n_pages=8,
+        n_kv_heads=4,
-        page_size=4,
-        n_kv_heads=2,
         head_dim=8,
         device=torch.device("cpu"),
         dtype=torch.float32,
     )
-    page_table = torch.tensor([[0]], dtype=torch.long)
+    loc = torch.tensor([[0, 1]], dtype=torch.long)
-    k = torch.randn(1, 2, 2, 8)
+    k = torch.randn(1, 2, 4, 8)
-    v = torch.randn(1, 2, 2, 8)
+    v = torch.randn(1, 2, 4, 8)
-
+    storage.set_kv_buffer(0, loc, k, v)
-    storage.write(0, page_table, 0, k, v)
+    assert torch.allclose(storage.get_key_buffer(0)[loc], k)
-    gk, gv = storage.gather(0, page_table, 2)
+    assert torch.allclose(storage.get_value_buffer(0)[loc], v)
-    assert torch.allclose(gk, k)
 
 
-def test_storage_write_cross_page():
+def test_kv_storage_buffer_shape():
-    storage = Storage(
+    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)
+    defaults.update(kwargs)
-    k = torch.randn(1, 8, 2, 8)
+    return PagePool(**defaults)
-    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)
 
 
-def test_storage_gather_truncates_to_total_len():
+def test_page_pool_paged_task_alloc():
-    storage = Storage(
+    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)
+    defaults.update(kwargs)
-    k = torch.randn(1, 6, 2, 8)
+    return PagePool(**defaults)
-    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)
 
 
-def test_storage_gather_clamps_negative_padding():
+def test_page_pool_paged_ps64_task_alloc():
-    storage = Storage(
+    pool = _make_paged_pool_ps64()
-        n_layers=1,
+    prompt = list(range(200))
-        n_pages=8,
+    assert pool.task_alloc("t1", prompt)
-        page_size=4,
+    assert pool.task_cached("t1") == 0
-        n_kv_heads=2,
+    n_pages = (200 + 63) // 64
-        head_dim=8,
+    assert len(pool._task_pages["t1"]) == n_pages
-        device=torch.device("cpu"),
+
-        dtype=torch.float32,
+
-    )
+def test_page_pool_paged_ps64_task_extend_crosses_page():
-    page_table = torch.tensor([[0, -1]], dtype=torch.long)
+    pool = _make_paged_pool_ps64()
-    gk, gv = storage.gather(0, page_table, 4)
+    pool.task_alloc("t1", list(range(64)))
-    assert gk.shape == (1, 4, 2, 8)
+    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)