From 3fa7e666767982cfda5acfa3ab83c6b68af76f30 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 8 Aug 2026 22:43:06 +0800
Subject: [PATCH] refactor: decouple task cache from PagePool and unify
 steady-state detection

- TaskCacheRegistry -> TaskCacheManager (independent, held by scheduler)
- TaskCacheState co-locates 5 parallel dicts into one dataclass
- AllocationStrategy base class + PagedStrategy subclass (page_size is a parameter)
- _rollback() helper for unified cleanup (no duplicate free paths)
- Task._kv_len + prefill_done property (explicit, no output_tokens proxy)
- Steady-state detection single-sourced in TaskCacheManager.bind()
- PagePool is now pure physical layer (no task knowledge)
- Removed dead _page_to_hash dict in RadixCache
---
 astrai/inference/__init__.py                |   2 +
 astrai/inference/core/__init__.py           |   2 +
 astrai/inference/core/cache.py              | 551 ++++++++++----------
 astrai/inference/core/executor.py           |  43 +-
 astrai/inference/core/scheduler.py          |  48 +-
 astrai/inference/core/task.py               |  19 +-
 scripts/tools/benchmark.py                  |  48 +-
 tests/extension/test_backend_equivalence.py |  43 +-
 tests/inference/test_cache.py               | 111 ++--
 tests/inference/test_task.py                |   7 +-
 10 files changed, 474 insertions(+), 400 deletions(-)

diff --git a/astrai/inference/__init__.py b/astrai/inference/__init__.py
index c85f70e..b639855 100644
--- a/astrai/inference/__init__.py
+++ b/astrai/inference/__init__.py
@@ -38,6 +38,7 @@ from astrai.inference.core import (
     RadixCache,
     ReqToTokenPool,
     Task,
+    TaskCacheManager,
     TaskManager,
     TaskStatus,
     page_hash,
@@ -67,6 +68,7 @@ __all__ = [
     "PagePool",
     "RadixCache",
     "ReqToTokenPool",
+    "TaskCacheManager",
     "page_hash",
     "sample",
     "BaseSamplingStrategy",
diff --git a/astrai/inference/core/__init__.py b/astrai/inference/core/__init__.py
index 673fefd..e5590d9 100644
--- a/astrai/inference/core/__init__.py
+++ b/astrai/inference/core/__init__.py
@@ -7,6 +7,7 @@ from astrai.inference.core.cache import (
     PagePool,
     RadixCache,
     ReqToTokenPool,
+    TaskCacheManager,
     page_hash,
 )
 from astrai.inference.core.executor import Executor
@@ -21,6 +22,7 @@ __all__ = [
     "PagePool",
     "RadixCache",
     "ReqToTokenPool",
+    "TaskCacheManager",
     "page_hash",
     "Executor",
     "InferenceScheduler",
diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py
index 258e987..6274219 100644
--- a/astrai/inference/core/cache.py
+++ b/astrai/inference/core/cache.py
@@ -1,19 +1,24 @@
 """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
+Layer 2 — ReqToTokenPool:  index table [req_idx, pos] -> physical token slot
+Layer 3 — AllocationStrategy: slot/page allocation with ref-counting and LRU
+
+PagePool owns the physical buffers and bind (KVCache assembly); it does
+not know about tasks.  TaskCacheManager owns task_id -> TaskCacheState
+mapping and delegates physical slot allocation to the strategy.
 
-PagePool orchestrates all three plus RadixCache (prefix 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
+Two strategies (selected once at construction):
+  - ContiguousStrategy: pre-allocated per-request blocks, no dynamic alloc
+  - PagedStrategy:      dynamic paged allocation; page_size is a parameter
+                        (1 = token-level, >1 = page-level with radix prefix)
 """
 
 import threading
-from dataclasses import dataclass
+from abc import ABC
+from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, OrderedDict
 
 import torch
@@ -30,6 +35,20 @@ class _BindState:
     seq_lens: List[int]
 
 
+@dataclass
+class TaskCacheState:
+    """Per-task cache allocation state.
+
+    Co-locating all task-owned cache state in one object makes the
+    alloc/free/extend lifecycle atomic.
+    """
+
+    req_idx: int
+    length: int = 0
+    cached: int = 0
+    pages: List[int] = field(default_factory=list)
+
+
 def _is_steady_increment(
     prev_sig: Optional[tuple],
     prev_vals: Optional[List[int]],
@@ -128,15 +147,11 @@ class RadixCache:
         self._page_size = page_size
         self._root = RadixNode()
         self._page_to_node: Dict[int, RadixNode] = {}
-        # Retained as an introspection-compatible map; matching never relies on
-        # this lossy value.
-        self._page_to_hash: Dict[int, int] = {}
         self._lock = threading.Lock()
 
     def evict(self, idx: int):
         with self._lock:
             node = self._page_to_node.pop(idx, None)
-            self._page_to_hash.pop(idx, None)
             if node is None:
                 return
             node.page_idx = None
@@ -169,7 +184,6 @@ class RadixCache:
             if logical_page_idx >= full_pages:
                 return
             old = self._page_to_node.pop(page_idx, None)
-            self._page_to_hash.pop(page_idx, None)
             if old is not None and old.parent is not None:
                 old.parent.children.pop(old.tokens, None)
 
@@ -185,12 +199,8 @@ class RadixCache:
             if node.page_idx is not None and node.page_idx != page_idx:
                 replaced = node.page_idx
                 self._page_to_node.pop(replaced, None)
-                self._page_to_hash.pop(replaced, None)
             node.page_idx = page_idx
             self._page_to_node[page_idx] = node
-            self._page_to_hash[page_idx] = page_hash(
-                token_ids, logical_page_idx, self._page_size
-            )
 
     def release(self, pages: List[int]) -> None:
         with self._lock:
@@ -272,21 +282,6 @@ class KVCache:
     """Pure data struct passed to model for KV cache I/O.
 
     The attention layer does raw buffer indexing — no methods, no abstraction.
-
-    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
-        max_len: max(seq_lens) as Python int — avoids GPU sync in decode
-        kv_indptr: [batch+1] int32 — prefix sum of seq_lens, precomputed once
-            per step so the attention backend avoids rebuilding it per layer.
-        qo_indptr: [batch+1] int32 — prefill qo prefix-sum (None in decode)
-        decode_o_part: split-KV o partial workspace (mirrors FlashInfer)
-        decode_ml_part: split-KV m/l partial workspace (mirrors FlashInfer)
-        decode_out: pre-allocated decode output buffer (graph-safe)
     """
 
     k_buffer: Tensor
@@ -303,21 +298,123 @@ class KVCache:
     decode_out: Optional[Tensor] = None
 
 
+class AllocationStrategy(ABC):
+    """Physical slot allocation policy.
+
+    The base class provides contiguous-mode defaults (all no-ops): req_to_token
+    is pre-filled at PagePool init, so no dynamic allocation is needed.
+    PagedStrategy overrides every method to add dynamic allocation.
+    """
+
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool:
+        return True
+
+    def free(self, state: TaskCacheState) -> None:
+        pass
+
+    def extend(self, state: TaskCacheState, pos: int) -> bool:
+        return True
+
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None:
+        pass
+
+    def record_hashes(
+        self, state: TaskCacheState, prompt_ids: List[int], start: int
+    ) -> None:
+        pass
+
+
+class PagedStrategy(AllocationStrategy):
+    """Dynamic paged allocation from a shared bitmask pool.
+
+    page_size is a parameter, not a separate strategy: at page_size=1 each
+    allocated page *is* one token slot (``page * 1 + 0``), and prefix
+    caching is simply disabled (``prefix=None``).  The unified page
+    formula ``pages[page_idx] * page_size + offset`` holds for both.
+    """
+
+    def __init__(
+        self,
+        alloc: Allocator,
+        prefix: Optional[RadixCache],
+        page_size: int,
+        req_pool: ReqToTokenPool,
+        device,
+    ):
+        self._alloc = alloc
+        self._prefix = prefix
+        self._page_size = page_size
+        self._req_pool = req_pool
+        self._device = device
+
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool:
+        if self._prefix is not None:
+            hits = self._prefix.lookup(prompt_ids)
+            state.cached = len(hits) * self._page_size
+            for p in hits:
+                self._alloc.inc_ref(p)
+            state.pages = list(hits)
+
+        remaining = len(prompt_ids) - state.cached
+        if remaining <= 0:
+            return True
+        n_new = (remaining + self._page_size - 1) // self._page_size
+        for _ in range(n_new):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            state.pages.append(p)
+        return True
+
+    def free(self, state: TaskCacheState) -> None:
+        if self._prefix is not None:
+            for p in state.pages:
+                keep = self._prefix.has_page(p)
+                self._alloc.free(p, keep_cached=keep)
+                if not keep:
+                    self._prefix.evict(p)
+        else:
+            for p in state.pages:
+                self._alloc.free(p)
+
+    def extend(self, state: TaskCacheState, pos: int) -> bool:
+        page_idx = pos // self._page_size
+        if page_idx >= len(state.pages):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            state.pages.append(p)
+        offset = pos % self._page_size
+        self._req_pool.req_to_token[state.req_idx, pos] = (
+            state.pages[page_idx] * self._page_size + offset
+        )
+        return True
+
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None:
+        total = len(prompt_ids)
+        for pos in range(state.cached, total):
+            page_idx = pos // self._page_size
+            offset = pos % self._page_size
+            if page_idx < len(state.pages):
+                self._req_pool.req_to_token[state.req_idx, pos] = (
+                    state.pages[page_idx] * self._page_size + offset
+                )
+
+    def record_hashes(
+        self, state: TaskCacheState, prompt_ids: List[int], start: int
+    ) -> None:
+        if self._prefix is None:
+            return
+        full = len(prompt_ids) // self._page_size
+        for i in range(start, min(full, len(state.pages))):
+            self._prefix.record(state.pages[i], prompt_ids, i)
+
+
 class PagePool:
-    """Top-level KV cache manager.
+    """Physical KV cache: buffers + req-pool + allocation strategy + bind.
 
-    Combines KVStorage + ReqToTokenPool + Allocator + RadixCache.
-
-    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).
+    Does not know about tasks — task lifecycle is managed by
+    :class:`TaskCacheManager`, which holds a reference to this pool.
     """
 
     def __init__(
@@ -342,10 +439,7 @@ class PagePool:
         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.n_tokens = max_batch_size * max_seq_len if self.contiguous else n_tokens
 
         self._storage = KVStorage(
             self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
@@ -357,248 +451,62 @@ class PagePool:
                 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[RadixCache] = None
+            self._strategy = AllocationStrategy()
         else:
             n_pages = self.n_tokens // page_size
-            self._alloc = Allocator(n_pages)
-            self._prefix = RadixCache(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()
-
-        # Steady-state decode validation: when the same ordered task set
-        # advances every sequence by exactly one token per step, bind_tasks
-        # updates the stable buffers in-place instead of re-cumsumming.
-        self._bind_state: Optional[_BindState] = None
-
-    # ---- task lifecycle ----
-
-    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
-        req_slots = self._req_pool.alloc(1)
-        if req_slots is None:
-            return False
-        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._task_pages.pop(task_id, None)
-                    self._task_slots.pop(task_id, None)
-                    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._task_pages.pop(task_id, None)
-                        self._task_slots.pop(task_id, None)
-                        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):
-        req_idx = self._task_req.pop(task_id, None)
-        if req_idx is None:
-            return
-        self._bind_state = None
-        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)
-            if self.page_size == 1:
-                for slot in self._task_slots.get(task_id, []):
-                    self._alloc.free(slot)
-            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:
-        req_idx = self._task_req.get(task_id)
-        if req_idx is None or pos >= self.max_seq_len:
-            return False
-
-        # Paged mode must also claim a physical slot for the new token;
-        # contiguous mode's block is pre-allocated so this is a no-op.
-        if not self.contiguous and not self._extend_slot(task_id, req_idx, pos):
-            return False
-
-        self._task_len[req_idx] = pos + 1
-        return True
-
-    def _extend_slot(self, task_id: str, req_idx: int, pos: int) -> bool:
-        """Allocate the physical slot for one extended token (paged mode)."""
-        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]
-            return True
-
-        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
-        return True
-
-    def task_cached(self, task_id: str) -> int:
-        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:
-            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)
-
-    def task_cacheable_ids(
-        self, task_id: str, prompt_ids: List[int], output_ids: List[int]
-    ):
-        """Return the sequence whose KV entries are already materialized.
-
-        The first sampled output is produced by prompt prefill, and the last
-        sampled output has not been decoded into KV yet.  Therefore the cache
-        can safely retain the prompt plus every output except the last one.
-        """
-        return list(prompt_ids) + list(output_ids[:-1])
-
-    # ---- bind for forward ----
+            alloc = Allocator(n_pages)
+            prefix = RadixCache(page_size) if page_size > 1 else None
+            if prefix is not None:
+                alloc.on_evict = prefix.evict
+            self._strategy = PagedStrategy(
+                alloc, prefix, page_size, self._req_pool, device
+            )
 
     def bind_tasks(
         self,
-        task_ids: List[str],
+        req_indices: List[int],
+        seq_lens: List[int],
         workspace: InferenceWorkspace,
         device: Optional[torch.device] = None,
         start_pos: Optional[int] = None,
+        incremental: bool = False,
     ) -> KVCache:
         if device is None:
             device = workspace.device
-        req_indices = [self._task_req[tid] for tid in task_ids]
-        # Per-request lengths come from the pool's own tracking (task_alloc
-        # sets len(prompt_ids); task_extend sets pos+1), so callers need not
-        # pass them.
-        seq_lens = [self._task_len[req_idx] for req_idx in req_indices]
-        b = len(task_ids)
-        sig = tuple(task_ids)
+        b = len(req_indices)
 
-        # Write into the caller's workspace buffers (fixed addresses, sized
-        # to max_batch/max_seq at init) — the sole owner of the per-step
-        # KV bind tensors.
         rpi_buf = workspace.req_pool_indices
         sl_buf = workspace.seq_lens
         kvp_buf = workspace.kv_indptr
         inc_buf = workspace.inc
         ocl_buf = workspace.out_cache_loc
 
-        prev = self._bind_state
-        incremental = (
-            start_pos is None
-            and prev is not None
-            and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
-        )
         if incremental:
-            # Steady-state decode: advance the stable buffers in-place.
-            # Normal-mode buffers keep ``+=`` legal regardless of whether
-            # this runs inside ``torch.inference_mode()``.
             sl_buf[:b] += 1
             kvp_buf[: b + 1] += inc_buf[: b + 1]
-            req_pool_indices = rpi_buf[:b]
-            seq_lens_t = sl_buf[:b]
-            kv_indptr = kvp_buf[: b + 1]
         else:
-            # Cold path: fill the stable buffers from fresh host tensors.
             rpi_buf[:b].copy_(
                 torch.tensor(req_indices, dtype=torch.long, device=device)
             )
             sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
             kvp_buf[: b + 1].zero_()
             kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
-            req_pool_indices = rpi_buf[:b]
-            seq_lens_t = sl_buf[:b]
-            kv_indptr = kvp_buf[: b + 1]
-        self._bind_state = _BindState(sig, list(seq_lens))
+
+        req_pool_indices = rpi_buf[:b]
+        seq_lens_t = sl_buf[:b]
+        kv_indptr = kvp_buf[: b + 1]
 
         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
             ]
-            # Ragged query segmentation for the prefill kernel, computed once
-            # (was rebuilt per layer in CudaBackend.fwd_prefill).
             q_len = seq_len - start_pos
             workspace.qo_indptr[: b + 1].copy_(
                 torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
             )
             qo_indptr = workspace.qo_indptr[: b + 1]
-            decode_o_part, decode_ml_part = None, None
-            decode_out = None
+            decode_o_part = decode_ml_part = decode_out = None
         else:
             write_pos = seq_lens_t - 1
             loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
@@ -624,37 +532,116 @@ class PagePool:
             decode_out=decode_out,
         )
 
-    # ---- 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
+class TaskCacheManager:
+    """Task <-> KV slot lifecycle manager.
 
-    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)
+    Sole owner of task state.  Owns the task_id -> TaskCacheState map and
+    delegates physical slot allocation to the strategy, and KV bind to
+    PagePool.  Held directly by the scheduler — not as a PagePool attribute.
+    """
 
-        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
+    def __init__(
+        self,
+        strategy: AllocationStrategy,
+        req_pool: ReqToTokenPool,
+        max_seq_len: int,
+        pool: PagePool,
+    ):
+        self._strategy = strategy
+        self._req_pool = req_pool
+        self._max_seq_len = max_seq_len
+        self._pool = pool
+        self._states: Dict[str, TaskCacheState] = {}
+        self._bind_state: Optional[_BindState] = None
+
+    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
+        self._bind_state = None
+        req_slots = self._req_pool.alloc(1)
+        if req_slots is None:
+            return False
+        state = TaskCacheState(req_idx=req_slots[0])
+        self._states[task_id] = state
+        if not self._strategy.alloc(state, prompt_ids):
+            self._rollback(state, task_id)
+            return False
+        self._strategy.write_indices(state, prompt_ids)
+        state.length = len(prompt_ids)
+        return True
+
+    def _rollback(self, state: TaskCacheState, task_id: str):
+        self._strategy.free(state)
+        self._req_pool.free([state.req_idx])
+        self._states.pop(task_id, None)
+
+    def task_free(self, task_id: str):
+        self._bind_state = None
+        state = self._states.pop(task_id, None)
+        if state is None:
+            return
+        self._strategy.free(state)
+        self._req_pool.free([state.req_idx])
+
+    def task_extend(self, task_id: str, pos: int) -> bool:
+        state = self._states.get(task_id)
+        if state is None or pos >= self._max_seq_len:
+            return False
+        if not self._strategy.extend(state, pos):
+            return False
+        state.length = pos + 1
+        return True
+
+    def task_cached(self, task_id: str) -> int:
+        state = self._states.get(task_id)
+        return state.cached if state is not None else 0
+
+    def task_record_hashes(
+        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
+    ):
+        state = self._states.get(task_id)
+        if state is not None:
+            self._strategy.record_hashes(state, prompt_ids, start_logical_page)
+
+    @staticmethod
+    def task_cacheable_ids(task_id: str, prompt_ids: List[int], output_ids: List[int]):
+        """Return the sequence whose KV entries are already materialized.
+
+        The first sampled output is produced by prompt prefill, and the last
+        sampled output has not been decoded into KV yet.
+        """
+        return list(prompt_ids) + list(output_ids[:-1])
+
+    def bind(
+        self,
+        task_ids: List[str],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+    ) -> KVCache:
+        states = [self._states[tid] for tid in task_ids]
+        req_indices = [s.req_idx for s in states]
+        seq_lens = [s.length for s in states]
+        sig = tuple(req_indices)
+
+        prev = self._bind_state
+        incremental = (
+            start_pos is None
+            and prev is not None
+            and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
+        )
+        self._bind_state = _BindState(sig, list(seq_lens))
+        self._bind_was_steady = incremental
+
+        return self._pool.bind_tasks(
+            req_indices,
+            seq_lens,
+            workspace,
+            device=device,
+            start_pos=start_pos,
+            incremental=incremental,
+        )
+
+    @property
+    def bind_was_steady(self) -> bool:
+        """Whether the most recent ``bind()`` was a steady-state increment."""
+        return self._bind_was_steady
diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py
index c685c68..4f45284 100644
--- a/astrai/inference/core/executor.py
+++ b/astrai/inference/core/executor.py
@@ -13,7 +13,7 @@ from astrai.extension.attention_backend import (
     attn_backend,
     get_backend,
 )
-from astrai.inference.core.cache import PagePool, _is_steady_increment
+from astrai.inference.core.cache import PagePool, TaskCacheManager
 from astrai.inference.core.graph import CudaGraphContext
 from astrai.inference.core.task import Task
 from astrai.inference.core.workspace import InferenceWorkspace
@@ -99,6 +99,7 @@ def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
 def _warmup_cuda_graphs(
     model: AutoModel,
     pool: PagePool,
+    task_cache: TaskCacheManager,
     ws: InferenceWorkspace,
     gctx: CudaGraphContext,
     max_batch_size: int,
@@ -113,12 +114,12 @@ def _warmup_cuda_graphs(
     # that follows.  Custom .so kernels do NOT need this — they are pre-built.
     warmup_len = 64
     tid = "_warmup_prefill"
-    if pool.task_alloc(tid, list(range(warmup_len))):
+    if task_cache.task_alloc(tid, list(range(warmup_len))):
         with (
             torch.inference_mode(),
             timed("warmup prefill", logger),
         ):
-            kv = pool.bind_tasks([tid], ws, start_pos=0)
+            kv = task_cache.bind([tid], ws, start_pos=0)
             ids_in = torch.arange(warmup_len, device=dev).unsqueeze(0)
             pos_in = ids_in
             model(
@@ -127,7 +128,7 @@ def _warmup_cuda_graphs(
                 kv_cache=kv,
                 position_ids=pos_in,
             )
-        pool.task_free(tid)
+        task_cache.task_free(tid)
 
     batch_sizes = [1]
     n = 2
@@ -142,12 +143,12 @@ def _warmup_cuda_graphs(
         prompt_tokens = [list(range(prompt_len)) for _ in range(b)]
         alloc_ok = True
         for tid, pt in zip(task_ids, prompt_tokens):
-            if not pool.task_alloc(tid, pt):
+            if not task_cache.task_alloc(tid, pt):
                 alloc_ok = False
                 break
         if not alloc_ok:
             for tid in task_ids:
-                pool.task_free(tid)
+                task_cache.task_free(tid)
             continue
 
         with (
@@ -159,8 +160,8 @@ def _warmup_cuda_graphs(
                 seq_pos = step
                 ws.position_ids[:b] = seq_pos
                 for tid in task_ids:
-                    pool.task_extend(tid, seq_pos)
-                kv = pool.bind_tasks(task_ids, ws)
+                    task_cache.task_extend(tid, seq_pos)
+                kv = task_cache.bind(task_ids, ws)
                 input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
                 ids_buf = ws.fill_input_ids([step] * b)
                 gctx.forward(
@@ -173,7 +174,7 @@ def _warmup_cuda_graphs(
                 )
 
         for tid in task_ids:
-            pool.task_free(tid)
+            task_cache.task_free(tid)
         torch.cuda.synchronize()
 
 
@@ -184,11 +185,13 @@ class Executor:
         self,
         model: AutoModel,
         kv_cache: PagePool,
+        task_cache: TaskCacheManager,
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
     ):
         self.model = model
         self.kv_cache = kv_cache
+        self.task_cache = task_cache
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 
@@ -228,6 +231,7 @@ class Executor:
         _warmup_cuda_graphs(
             self.model,
             self.kv_cache,
+            self.task_cache,
             self._workspace,
             self._graph_ctx,
             max_batch_size=self.kv_cache.max_batch_size,
@@ -321,7 +325,7 @@ class Executor:
                 input_ids,
                 input_mask=input_mask,
                 position_ids=position_ids,
-                kv_cache=self.kv_cache.bind_tasks(
+                kv_cache=self.task_cache.bind(
                     task_ids,
                     self._workspace,
                     start_pos=start_pos,
@@ -363,26 +367,25 @@ class Executor:
         task_ids = [t.task_id for t in tasks]
         cur_positions = [t.next_pos for t in tasks]
 
-        sig = tuple(task_ids)
-        cached = self._decode_cache
-        prev_sig = cached.task_sig if cached is not None else None
-        prev_pos = cached.positions if cached is not None else None
-        if _is_steady_increment(prev_sig, prev_pos, sig, cur_positions):
-            info = cached.sampling_info
+        kv_cache = self.task_cache.bind(task_ids, ws)
+
+        if self.task_cache.bind_was_steady:
+            info = (
+                self._decode_cache.sampling_info
+                if self._decode_cache is not None
+                else _build_sampling_batch_info(tasks, self.device)
+            )
             ws.position_ids[:b] += 1
-            self._decode_cache = DecodeSteadyState(sig, cur_positions, info)
         else:
             info = _build_sampling_batch_info(tasks, self.device)
             ws.position_ids[:b].copy_(
                 torch.tensor(cur_positions, dtype=torch.long, device=self.device)
             )
-            self._decode_cache = DecodeSteadyState(sig, cur_positions, info)
+        self._decode_cache = DecodeSteadyState(tuple(task_ids), cur_positions, info)
 
         total_len = max(cur_positions) + 1
         input_mask = ws.decode_mask(ws.position_ids[:b], total_len)
 
-        kv_cache = self.kv_cache.bind_tasks(task_ids, ws)
-
         # ---- forward (graph replay or live run + capture) ----
 
         use_graph = (
diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py
index 1bd6c7e..54c1f96 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 PagePool
+from astrai.inference.core.cache import PagePool, TaskCacheManager
 from astrai.inference.core.executor import Executor
 from astrai.inference.core.metrics import MetricsCollector
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
@@ -59,6 +59,13 @@ class InferenceScheduler:
 
         self._metrics = MetricsCollector()
 
+        self._task_cache = TaskCacheManager(
+            strategy=self._cache._strategy,
+            req_pool=self._cache._req_pool,
+            max_seq_len=self.max_seq_len,
+            pool=self._cache,
+        )
+
         self._task_mgr = TaskManager(
             tokenizer=tokenizer,
             max_batch_size=max_batch_size,
@@ -69,6 +76,7 @@ class InferenceScheduler:
         self._executor = Executor(
             model=model,
             kv_cache=self._cache,
+            task_cache=self._task_cache,
             device=self.device,
             dtype=self.dtype,
         )
@@ -81,7 +89,7 @@ class InferenceScheduler:
 
     def remove_task(self, task_id: str):
         for task in self._task_mgr.remove_task(task_id):
-            self._cache.task_free(task.task_id)
+            self._task_cache.task_free(task.task_id)
 
     def get_stats(self) -> Dict[str, Any]:
         return self._task_mgr.get_stats()
@@ -109,9 +117,7 @@ class InferenceScheduler:
             already appended to ``output_ids``) and tasks that hit the
             sequence cap and were marked ``ABORTED``.
         """
-        cache = self._cache
-
-        to_prefill = [t for t in tasks if t.output_tokens == 0 and t.prompt_ids]
+        to_prefill = [t for t in tasks if not t.prefill_done and t.prompt_ids]
         prefilled_ids = set()
         produced: List[Task] = []
         if to_prefill:
@@ -120,7 +126,9 @@ class InferenceScheduler:
 
             groups: Dict[Tuple[int, int], List[Task]] = {}
             for t in to_prefill:
-                start_pos = min(cache.task_cached(t.task_id), len(t.prompt_ids) - 1)
+                start_pos = min(
+                    self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
+                )
                 groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
 
             for (prompt_len, start_pos), group in groups.items():
@@ -132,12 +140,13 @@ class InferenceScheduler:
                 for t, out in zip(prefilled, step_out):
                     t.output_ids.append(out[0] if return_logprobs else out)
                     t.output_tokens += 1
+                    t.mark_prefill_done()
                     prefilled_ids.add(t.task_id)
                     produced.append(t)
 
-                start_logical_page = start_pos // getattr(cache, "page_size", 64)
+                start_logical_page = start_pos // self._cache.page_size
                 for t in group:
-                    cache.task_record_hashes(
+                    self._task_cache.task_record_hashes(
                         t.task_id, t.prompt_ids, start_logical_page
                     )
 
@@ -146,7 +155,7 @@ class InferenceScheduler:
         for t in tasks:
             if t.task_id in prefilled_ids:
                 continue
-            if cache.task_extend(t.task_id, t.next_pos):
+            if self._task_cache.task_extend(t.task_id, t.next_pos):
                 decoded.append(t)
             else:
                 t.status = TaskStatus.ABORTED
@@ -160,25 +169,25 @@ class InferenceScheduler:
             for t, out in zip(decoded, step_out):
                 t.output_ids.append(out[0] if return_logprobs else out)
                 t.output_tokens += 1
+                t.advance_kv()
                 produced.append(t)
 
         return produced, aborted
 
     def _run_generation_loop(self):
         stop_ids = self._task_mgr.tokenizer.stop_ids
-        cache = self._cache
         try:
             while not self._stop_event.is_set():
                 finished = self._task_mgr.remove_finished_tasks(stop_ids)
                 for task in finished:
                     if task.status == TaskStatus.FINISHED:
-                        cache.task_record_hashes(
+                        self._task_cache.task_record_hashes(
                             task.task_id,
-                            cache.task_cacheable_ids(
+                            self._task_cache.task_cacheable_ids(
                                 task.task_id, task.prompt_ids, task.output_ids
                             ),
                         )
-                    cache.task_free(task.task_id)
+                    self._task_cache.task_free(task.task_id)
 
                 active = self._task_mgr.get_active_tasks()
                 available = self._task_mgr.max_batch_size - len(active)
@@ -186,7 +195,7 @@ class InferenceScheduler:
                     candidates = self._task_mgr.pull_candidates(available)
                     failed = []
                     for task in candidates:
-                        if cache.task_alloc(task.task_id, task.prompt_ids):
+                        if self._task_cache.task_alloc(task.task_id, task.prompt_ids):
                             self._task_mgr.activate(task)
                         else:
                             failed.append(task)
@@ -216,7 +225,7 @@ class InferenceScheduler:
             logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
             for task in self._task_mgr.get_active_tasks():
                 self._task_mgr.invoke_callback(task.task_id, STOP)
-                cache.task_free(task.task_id)
+                self._task_cache.task_free(task.task_id)
             for task in self._task_mgr.get_waiting_tasks():
                 self._task_mgr.invoke_callback(task.task_id, STOP)
             self._task_mgr.clear_queues()
@@ -237,10 +246,10 @@ class InferenceScheduler:
         self._loop_thread = None
         for task in self._task_mgr.get_active_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._cache.task_free(task.task_id)
+            self._task_cache.task_free(task.task_id)
         for task in self._task_mgr.get_waiting_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._cache.task_free(task.task_id)
+            self._task_cache.task_free(task.task_id)
         self._task_mgr.clear_queues()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -279,7 +288,6 @@ class InferenceScheduler:
             ``List[Tuple[List[int], List[float]]]``.
         """
         stop_ids = self._task_mgr.tokenizer.stop_ids
-        cache = self._cache
         seq_cap = self.max_seq_len
 
         tasks: List[Task] = []
@@ -305,7 +313,7 @@ class InferenceScheduler:
                 frequency_penalty=frequency_penalty,
                 rep_window=rep_window,
             )
-            if not cache.task_alloc(task.task_id, task.prompt_ids):
+            if not self._task_cache.task_alloc(task.task_id, task.prompt_ids):
                 tasks.append(None)
                 continue
             task.input_tokens = len(task.prompt_ids)
@@ -324,7 +332,7 @@ class InferenceScheduler:
                     self._metrics.mark_finished(
                         t.task_id, t.input_tokens, t.output_tokens
                     )
-                    cache.task_free(t.task_id)
+                    self._task_cache.task_free(t.task_id)
 
         results: List[Any] = []
         for t in tasks:
diff --git a/astrai/inference/core/task.py b/astrai/inference/core/task.py
index cd04ade..db3b1c4 100644
--- a/astrai/inference/core/task.py
+++ b/astrai/inference/core/task.py
@@ -77,8 +77,18 @@ class Task:
         self.output_logprobs: List[float] = []
         self.input_tokens: int = 0
         self.output_tokens: int = 0
+        self._kv_len: int = 0
         self._decoder: Optional[StreamDecoder] = None
 
+    def mark_prefill_done(self):
+        """Prompt KV is materialized by prefill; first output sampled but
+        not yet written to KV."""
+        self._kv_len = self.input_tokens
+
+    def advance_kv(self):
+        """One more position written to KV (after a decode forward)."""
+        self._kv_len += 1
+
     def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
         """Decode the last appended output token, buffering incomplete
         multi-byte sequences across calls.
@@ -91,8 +101,13 @@ class Task:
 
     @property
     def next_pos(self) -> int:
-        # The first output is sampled from prefill and enters KV on the next step.
-        return self.input_tokens + max(0, len(self.output_ids) - 1)
+        """KV position where the next decode step will write."""
+        return self._kv_len
+
+    @property
+    def prefill_done(self) -> bool:
+        """True when all prompt KV entries are materialized."""
+        return self._kv_len >= self.input_tokens > 0
 
     def is_finished(self, stop_ids: List[int]) -> bool:
         if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py
index 9a7dbc2..eb0937f 100644
--- a/scripts/tools/benchmark.py
+++ b/scripts/tools/benchmark.py
@@ -7,7 +7,7 @@ import torch
 
 from astrai.config import BaseModelConfig, ConfigFactory
 from astrai.extension import ATTN_BACKEND, AttentionBackendFactory, attn_backend
-from astrai.inference.core.cache import PagePool
+from astrai.inference.core.cache import PagePool, TaskCacheManager
 from astrai.inference.core.graph import CudaGraphContext
 from astrai.inference.core.workspace import InferenceWorkspace
 from astrai.model import AutoModel, AutoRegressiveLM
@@ -91,9 +91,19 @@ class GenerationBenchmark:
             dtype=pool.dtype,
         )
 
+    @staticmethod
+    def _make_task_cache(pool: PagePool) -> TaskCacheManager:
+        return TaskCacheManager(
+            strategy=pool._strategy,
+            req_pool=pool._req_pool,
+            max_seq_len=pool.max_seq_len,
+            pool=pool,
+        )
+
     def _run_prefill(
         self,
         pool: PagePool,
+        task_cache: TaskCacheManager,
         batch_size: int,
         prompt_len: int,
         workspace: InferenceWorkspace,
@@ -112,9 +122,9 @@ class GenerationBenchmark:
 
         task_ids = [f"bench_{i}" for i in range(batch_size)]
         for tid in task_ids:
-            pool.task_alloc(tid, list(range(prompt_len)))
+            task_cache.task_alloc(tid, list(range(prompt_len)))
 
-        kv_cache = pool.bind_tasks(task_ids, workspace, self.device, start_pos=0)
+        kv_cache = task_cache.bind(task_ids, workspace, self.device, start_pos=0)
         with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
@@ -128,6 +138,7 @@ class GenerationBenchmark:
     def _run_decode_step(
         self,
         pool: PagePool,
+        task_cache: TaskCacheManager,
         task_ids: list,
         seq_len: int,
         workspace: InferenceWorkspace,
@@ -141,11 +152,11 @@ class GenerationBenchmark:
         )
         total_len = seq_len + 1
         for tid in task_ids:
-            pool.task_extend(tid, seq_len)
+            task_cache.task_extend(tid, seq_len)
         input_mask = position_ids[:, :, None] >= torch.arange(
             total_len, device=self.device
         )
-        kv_cache = pool.bind_tasks(task_ids, workspace, self.device)
+        kv_cache = task_cache.bind(task_ids, workspace, self.device)
         with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
@@ -164,9 +175,10 @@ class GenerationBenchmark:
 
         pool = self._make_pool(batch_size, prompt_length)
         workspace = self._make_workspace(pool, self.config)
+        task_cache = self._make_task_cache(pool)
         task_ids = [f"bench_prefill_{i}" for i in range(batch_size)]
         for tid in task_ids:
-            pool.task_alloc(tid, list(range(prompt_length)))
+            task_cache.task_alloc(tid, list(range(prompt_length)))
 
         input_ids = torch.randint(
             0, self.config.vocab_size, (batch_size, prompt_length), device=self.device
@@ -179,7 +191,7 @@ class GenerationBenchmark:
         input_mask = position_ids.unsqueeze(-1) >= torch.arange(
             prompt_length, device=self.device
         )
-        kv_cache = pool.bind_tasks(task_ids, workspace, self.device, start_pos=0)
+        kv_cache = task_cache.bind(task_ids, workspace, self.device, start_pos=0)
 
         for _ in range(3):
             with torch.inference_mode(), attn_backend(self.backend):
@@ -240,7 +252,10 @@ class GenerationBenchmark:
         max_seq_len = prompt_length + 5 + gen_length * num_trials
         pool = self._make_pool(batch_size, max_seq_len)
         workspace = self._make_workspace(pool, self.config)
-        task_ids = self._run_prefill(pool, batch_size, prompt_length, workspace)
+        task_cache = self._make_task_cache(pool)
+        task_ids = self._run_prefill(
+            pool, task_cache, batch_size, prompt_length, workspace
+        )
 
         b = batch_size
         input_ids_buf = torch.zeros(b, 1, dtype=torch.long, device=self.device)
@@ -256,8 +271,8 @@ class GenerationBenchmark:
             )
             position_ids_buf[:] = seq_len
             for tid in task_ids:
-                pool.task_extend(tid, seq_len)
-            kv_cache = pool.bind_tasks(task_ids, workspace, self.device)
+                task_cache.task_extend(tid, seq_len)
+            kv_cache = task_cache.bind(task_ids, workspace, self.device)
 
             input_mask = torch.ge(
                 position_ids_buf[:, None],
@@ -313,15 +328,22 @@ class GenerationBenchmark:
         max_seq_len = prompt_length + 5 + gen_length * num_trials
         pool = self._make_pool(batch_size, max_seq_len)
         workspace = self._make_workspace(pool, self.config)
-        task_ids = self._run_prefill(pool, batch_size, prompt_length, workspace)
+        task_cache = self._make_task_cache(pool)
+        task_ids = self._run_prefill(
+            pool, task_cache, batch_size, prompt_length, workspace
+        )
 
         for i in range(5):
-            self._run_decode_step(pool, task_ids, prompt_length + i, workspace)
+            self._run_decode_step(
+                pool, task_cache, task_ids, prompt_length + i, workspace
+            )
         torch.cuda.synchronize()
 
         t0 = time.perf_counter()
         for i in range(gen_length * num_trials):
-            self._run_decode_step(pool, task_ids, prompt_length + 5 + i, workspace)
+            self._run_decode_step(
+                pool, task_cache, task_ids, prompt_length + 5 + i, workspace
+            )
         torch.cuda.synchronize()
         elapsed = time.perf_counter() - t0
         tokens = batch_size * gen_length * num_trials
diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py
index 6cc0b2a..27f5e68 100644
--- a/tests/extension/test_backend_equivalence.py
+++ b/tests/extension/test_backend_equivalence.py
@@ -7,11 +7,20 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
 import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
-from astrai.inference.core.cache import PagePool
+from astrai.inference.core.cache import PagePool, TaskCacheManager
 from astrai.inference.core.workspace import InferenceWorkspace
 from tests.extension.conftest import D, skip_no_kernel
 
 
+def _mk_task_cache(pool: PagePool) -> TaskCacheManager:
+    return TaskCacheManager(
+        strategy=pool._strategy,
+        req_pool=pool._req_pool,
+        max_seq_len=pool.max_seq_len,
+        pool=pool,
+    )
+
+
 def _ws(pool: PagePool) -> InferenceWorkspace:
     return InferenceWorkspace(
         pool.max_batch_size,
@@ -74,20 +83,21 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
         dtype=torch.bfloat16,
     )
 
+    task_cache = _mk_task_cache(cache)
     ws = _ws(cache)
-    cache.task_alloc("t1", prompt_ids[0])
-    cache.task_alloc("t2", prompt_ids[1])
-    kv1 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
+    task_cache.task_alloc("t1", prompt_ids[0])
+    task_cache.task_alloc("t2", prompt_ids[1])
+    kv1 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     with torch.inference_mode():
         out_torch = model(
             input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
         )
 
-    cache.task_free("t1")
-    cache.task_free("t2")
-    cache.task_alloc("t1", prompt_ids[0])
-    cache.task_alloc("t2", prompt_ids[1])
-    kv2 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
+    task_cache.task_free("t1")
+    task_cache.task_free("t2")
+    task_cache.task_alloc("t1", prompt_ids[0])
+    task_cache.task_alloc("t2", prompt_ids[1])
+    kv2 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     with attn_backend(ATTN_BACKEND.CUDA):
         with torch.inference_mode():
             out_cuda = model(
@@ -138,10 +148,11 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
         input_mask[i, : len(p)] = True
         position_ids[i, : len(p)] = torch.arange(len(p), device=device)
 
+    task_cache = _mk_task_cache(cache)
     ws = _ws(cache)
-    cache.task_alloc("t1", prompt_ids[0])
-    cache.task_alloc("t2", prompt_ids[1])
-    kv = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
+    task_cache.task_alloc("t1", prompt_ids[0])
+    task_cache.task_alloc("t2", prompt_ids[1])
+    kv = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     with torch.inference_mode():
         model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
 
@@ -151,15 +162,15 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     total_len = 9
     dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
 
-    cache.task_extend("t1", 8)
-    cache.task_extend("t2", 6)
-    kv_t = cache.bind_tasks(["t1", "t2"], ws)
+    task_cache.task_extend("t1", 8)
+    task_cache.task_extend("t2", 6)
+    kv_t = task_cache.bind(["t1", "t2"], ws)
     with torch.inference_mode():
         out_torch = model(
             dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
         )
 
-    kv_c = cache.bind_tasks(["t1", "t2"], ws)
+    kv_c = task_cache.bind(["t1", "t2"], ws)
     with attn_backend(ATTN_BACKEND.CUDA):
         with torch.inference_mode():
             out_cuda = model(
diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py
index d83c91c..f383362 100644
--- a/tests/inference/test_cache.py
+++ b/tests/inference/test_cache.py
@@ -8,6 +8,7 @@ from astrai.inference import (
     PagePool,
     RadixCache,
     ReqToTokenPool,
+    TaskCacheManager,
     page_hash,
 )
 from astrai.inference.core.workspace import InferenceWorkspace
@@ -25,6 +26,15 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
     )
 
 
+def _make_task_cache(pool: PagePool) -> TaskCacheManager:
+    return TaskCacheManager(
+        strategy=pool._strategy,
+        req_pool=pool._req_pool,
+        max_seq_len=pool.max_seq_len,
+        pool=pool,
+    )
+
+
 # ---- page_hash ----
 
 
@@ -116,9 +126,9 @@ def test_prefix_cache_ignores_partial_last_page():
 def test_prefix_cache_on_evict_clears_mappings():
     prefix = RadixCache(64)
     prefix.record(0, list(range(64)), 0)
-    assert 0 in prefix._page_to_hash
+    assert prefix.has_page(0)
     prefix.evict(0)
-    assert 0 not in prefix._page_to_hash
+    assert not prefix.has_page(0)
 
 
 def test_prefix_cache_has_page():
@@ -162,7 +172,8 @@ def test_prefix_cache_does_not_record_partial_page():
 
 def test_page_pool_task_cacheable_ids_excludes_unmaterialized_tail():
     pool = _make_paged_pool_ps64()
-    assert pool.task_cacheable_ids("missing", [1, 2], [3, 4]) == [1, 2, 3]
+    task_cache = _make_task_cache(pool)
+    assert task_cache.task_cacheable_ids("missing", [1, 2], [3, 4]) == [1, 2, 3]
 
 
 # ---- ReqToTokenPool ----
@@ -243,31 +254,35 @@ def _make_contiguous_pool(**kwargs):
 
 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
+    task_cache = _make_task_cache(pool)
+    assert task_cache.task_alloc("t1", [1, 2, 3])
+    assert "t1" in task_cache._states
+    task_cache.task_free("t1")
+    assert "t1" not in task_cache._states
 
 
 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)
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", [1, 2, 3])
+    assert task_cache.task_extend("t1", 3)
+    assert task_cache.task_extend("t1", 63)
+    assert not task_cache.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
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", [1, 2, 3])
+    assert task_cache.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"], _ws(pool), start_pos=0)
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(10)))
+    task_cache.task_alloc("t2", list(range(10)))
+    kv = task_cache.bind(["t1", "t2"], _ws(pool), 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,)
@@ -275,12 +290,13 @@ def test_page_pool_contiguous_bind_tasks_prefill():
 
 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)))
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(10)))
+    task_cache.task_alloc("t2", list(range(8)))
     # Simulate one decode extension so seq_lens advance to 11 and 9.
-    assert pool.task_extend("t1", 10)
-    assert pool.task_extend("t2", 8)
-    kv = pool.bind_tasks(["t1", "t2"], _ws(pool))
+    assert task_cache.task_extend("t1", 10)
+    assert task_cache.task_extend("t2", 8)
+    kv = task_cache.bind(["t1", "t2"], _ws(pool))
     assert kv.out_cache_loc.shape == (2, 1)
     assert kv.seq_lens.tolist() == [11, 9]
 
@@ -288,9 +304,10 @@ def test_page_pool_contiguous_bind_tasks_decode():
 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)))
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(4)))
 
-    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
+    kv = task_cache.bind(["t1"], _ws(pool), 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
@@ -324,35 +341,38 @@ def _make_paged_pool(**kwargs):
 
 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]
+    task_cache = _make_task_cache(pool)
+    assert task_cache.task_alloc("t1", list(range(10)))
+    state = task_cache._states["t1"]
+    assert len(state.pages) == 10
+    assert pool._req_pool.req_to_token[state.req_idx, 0].item() == state.pages[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"]
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(4)))
+    assert task_cache.task_extend("t1", 4)
+    req_idx = task_cache._states["t1"].req_idx
     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
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(8)))
+    task_cache.task_free("t1")
+    assert "t1" not in task_cache._states
     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)))
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(4)))
 
-    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
+    kv = task_cache.bind(["t1"], _ws(pool), 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
@@ -384,26 +404,29 @@ def _make_paged_pool_ps64(**kwargs):
 
 def test_page_pool_paged_ps64_task_alloc():
     pool = _make_paged_pool_ps64()
+    task_cache = _make_task_cache(pool)
     prompt = list(range(200))
-    assert pool.task_alloc("t1", prompt)
-    assert pool.task_cached("t1") == 0
+    assert task_cache.task_alloc("t1", prompt)
+    assert task_cache.task_cached("t1") == 0
     n_pages = (200 + 63) // 64
-    assert len(pool._task_pages["t1"]) == n_pages
+    assert len(task_cache._states["t1"].pages) == 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
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(64)))
+    assert task_cache.task_extend("t1", 64)
+    assert len(task_cache._states["t1"].pages) >= 2
 
 
 def test_page_pool_paged_ps64_bind_roundtrip():
     pool = _make_paged_pool_ps64(n_layers=1, n_kv_heads=2, head_dim=4)
+    task_cache = _make_task_cache(pool)
     prompt = list(range(128))
-    pool.task_alloc("t1", prompt)
+    task_cache.task_alloc("t1", prompt)
 
-    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
+    kv = task_cache.bind(["t1"], _ws(pool), 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
diff --git a/tests/inference/test_task.py b/tests/inference/test_task.py
index 65de9d6..a93b732 100644
--- a/tests/inference/test_task.py
+++ b/tests/inference/test_task.py
@@ -20,11 +20,12 @@ def test_task_default_status_is_pending():
 def test_task_next_pos():
     task = Task("id1", [1, 2, 3])
     task.input_tokens = 5
+    task.mark_prefill_done()
     assert task.next_pos == 5
-    task.output_ids.append(4)
-    assert task.next_pos == 5
-    task.output_ids.append(5)
+    task.advance_kv()
     assert task.next_pos == 6
+    task.advance_kv()
+    assert task.next_pos == 7
 
 
 def test_task_is_finished_max_tokens():