refactor: split infer core into subpackages by concern

- Eliminate core/ directory into cache/, runtime/, network/ subpackages plus flat modules
- Split cache.py (647 lines) into cache/{buffer,strategy,pool}.py by layer
- Add explicit ContiguousStrategy, make AllocationStrategy a real ABC
- Move TaskCacheState to cache/strategy.py, drop string forward references
- Rename api/ to network/, server.py to app.py
- Move sample.py into runtime/ alongside executor and graph
- Simplify TaskCacheManager.__init__ to single pool param
- Expose pool.strategy and pool.req_pool as public properties
- Fix KVCache import in attention_backend.py (TYPE_CHECKING guard)
- Fix steady-state decode reading uninitialized position_ids on first step

astrai/extension/attention_backend.py

@@ -51,8 +51,11 @@ from astrai.extension.loader import is_available
 from astrai.factory import BaseFactory
 
 if TYPE_CHECKING:
-    from astrai.inference.core.cache import KVCache
+    from astrai.inference.cache import KVCache
 
+
+_default_backend: Optional["AttentionBackend"] = None
+_default_backend_lock = threading.Lock()
 _current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
     "attn_backend"
 )
@@ -101,10 +104,6 @@ class ATTN_BACKEND(enum.Enum):
     FLASH = "flash"
 
 
-_default_backend: Optional["AttentionBackend"] = None
-_default_backend_lock = threading.Lock()
-
-
 def _priority_backends() -> list["AttentionBackend"]:
     """Available backends in priority order: cuda -> flash -> torch."""
     backends: list[AttentionBackend] = []

astrai/inference/__init__.py

@@ -1,15 +1,29 @@
 """Inference module for continuous batching.
 
-Layers:
+Subpackages:
-  - core/:        Core inference loop (cache, executor, scheduler, task)
+  - cache/:     KV cache (buffers, strategies, pool)
-  - api/:         HTTP orchestration (ProtocolHandler, server)
+  - runtime/:   Execution + sampling (executor, CUDA graph, sampling strategies)
-  - protocols/:   Response builders (OpenAI, Anthropic)
+  - task/:      Request lifecycle + performance metrics
-  - transport/:   SSE transport utilities
+  - network/:   HTTP protocol handling (server, protocol, OpenAI/Anthropic builders)
+
+Modules:
+  - scheduler.py:  Continuous batching loop
+  - workspace.py:  Pre-allocated GPU buffers
   - engine.py:     Facade (InferenceEngine)
-  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy)
 """
 
-from astrai.inference.api import (
+from astrai.inference.cache import (
+    Allocator,
+    KVCache,
+    KVStorage,
+    PagePool,
+    RadixCache,
+    ReqToTokenPool,
+    TaskCacheManager,
+    page_hash,
+)
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network import (
     AnthropicMessage,
     BaseToolParser,
     ChatCompletionRequest,
@@ -25,26 +39,10 @@ from astrai.inference.api import (
     get_app,
     run_server,
 )
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
-from astrai.inference.core import (
+from astrai.inference.runtime.executor import Executor
-    STOP,
+from astrai.inference.runtime.sample import (
-    Allocator,
-    Executor,
-    InferenceScheduler,
-    KVCache,
-    KVStorage,
-    PagePool,
-    RadixCache,
-    ReqToTokenPool,
-    Task,
-    TaskCacheManager,
-    TaskManager,
-    TaskStatus,
-    page_hash,
-)
-from astrai.inference.engine import InferenceEngine
-from astrai.inference.sample import (
     BaseSamplingStrategy,
     FrequencyPenaltyStrategy,
     SamplingPipeline,
@@ -53,6 +51,8 @@ from astrai.inference.sample import (
     TopPStrategy,
     sample,
 )
+from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.task import STOP, Task, TaskManager, TaskStatus
 
 __all__ = [
     "InferenceEngine",

astrai/inference/cache/__init__.py

@@ -0,0 +1,27 @@
+"""KV cache subsystem: buffers, strategies, pool management."""
+
+from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.pool import PagePool, TaskCacheManager, page_hash
+from astrai.inference.cache.strategy import (
+    AllocationStrategy,
+    Allocator,
+    ContiguousStrategy,
+    PagedStrategy,
+    RadixCache,
+    TaskCacheState,
+)
+
+__all__ = [
+    "KVCache",
+    "KVStorage",
+    "ReqToTokenPool",
+    "Allocator",
+    "RadixCache",
+    "TaskCacheState",
+    "AllocationStrategy",
+    "ContiguousStrategy",
+    "PagedStrategy",
+    "PagePool",
+    "TaskCacheManager",
+    "page_hash",
+]

astrai/inference/cache/buffer.py

@@ -0,0 +1,104 @@
+"""Physical KV cache buffers.
+
+Layer 1 — ``KVStorage``:   flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
+Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
+Layer 3 — ``KVCache``:       pure dataclass passed to the model for direct buffer access
+
+These classes have no knowledge of tasks, allocation policies, or scheduling.
+They are the "dumb" physical storage layer.
+"""
+
+import threading
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+
+class ReqToTokenPool:
+    """Maps [req_idx, pos] → physical token slot in KV storage.
+
+    Each row is one request; each column is a sequence position.  The value
+    at [req_idx, pos] is the flat index into the KV storage buffers.
+    """
+
+    def __init__(self, size: int, max_context_len: int, device: torch.device):
+        self.size = size
+        self.max_context_len = max_context_len
+        self.req_to_token = torch.zeros(
+            (size, max_context_len), dtype=torch.long, device=device
+        )
+        self.free_slots = list(range(size))
+        self._lock = threading.Lock()
+
+    def alloc(self, num_reqs: int) -> Optional[List[int]]:
+        with self._lock:
+            if num_reqs > len(self.free_slots):
+                return None
+            slots = self.free_slots[:num_reqs]
+            self.free_slots = self.free_slots[num_reqs:]
+            return slots
+
+    def free(self, req_indices: List[int]):
+        with self._lock:
+            self.free_slots.extend(req_indices)
+
+    def write(self, indices, values):
+        self.req_to_token[indices] = values
+
+
+class KVStorage:
+    """Token-level KV cache storage.
+
+    Buffers: ``[n_layers, size, n_kv_heads, head_dim]``.  Each token occupies
+    one slot indexed by ``ReqToTokenPool``.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        n_layers: int,
+        n_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.size = size
+        self.k_buffer = torch.empty(
+            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
+        )
+        self.v_buffer = torch.empty(
+            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
+        )
+
+    def get_key_buffer(self, layer_id: int) -> Tensor:
+        return self.k_buffer[layer_id]
+
+    def get_value_buffer(self, layer_id: int) -> Tensor:
+        return self.v_buffer[layer_id]
+
+    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
+        self.k_buffer[layer_id, loc] = k
+        self.v_buffer[layer_id, loc] = v
+
+
+@dataclass
+class KVCache:
+    """Pure data struct passed to model for KV cache I/O.
+
+    The attention layer does raw buffer indexing — no methods, no abstraction.
+    """
+
+    k_buffer: Tensor
+    v_buffer: Tensor
+    req_to_token: Tensor
+    req_pool_indices: Tensor
+    seq_lens: Tensor
+    out_cache_loc: Tensor
+    max_len: int = 0
+    kv_indptr: Optional[Tensor] = None
+    qo_indptr: Optional[Tensor] = None
+    decode_o_part: Optional[Tensor] = None
+    decode_ml_part: Optional[Tensor] = None
+    decode_out: Optional[Tensor] = None

astrai/inference/cache/pool.py

@@ -0,0 +1,351 @@
+"""KV cache orchestration: PagePool + TaskCacheManager.
+
+PagePool owns the physical buffers (``KVStorage`` + ``ReqToTokenPool``)
+and wires them to an allocation strategy.  It assembles the ``KVCache``
+dataclass passed to the model forward.
+
+TaskCacheManager owns the ``task_id`` → ``TaskCacheState`` mapping and
+delegates physical slot allocation to the strategy, and KV bind to the pool.
+
+See ``cache_buffer.py`` for the raw buffer primitives and ``cache_strategy.py``
+for the allocation policies.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+
+from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.strategy import (
+    AllocationStrategy,
+    Allocator,
+    ContiguousStrategy,
+    PagedStrategy,
+    RadixCache,
+    TaskCacheState,
+)
+from astrai.inference.workspace import InferenceWorkspace
+
+# Re-export everything so existing ``from astrai.inference.cache import ...``
+# continues to work unchanged after the file split.
+__all__ = [
+    "KVCache",
+    "KVStorage",
+    "ReqToTokenPool",
+    "Allocator",
+    "RadixCache",
+    "AllocationStrategy",
+    "ContiguousStrategy",
+    "PagedStrategy",
+    "PagePool",
+    "TaskCacheManager",
+    "TaskCacheState",
+    "page_hash",
+]
+
+# ---- helpers ----
+
+
+def page_hash(
+    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
+) -> int:
+    start = page_idx * page_size
+    end = min(start + page_size, len(token_ids))
+    h = parent_hash
+    for i in range(start, end):
+        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
+    return h
+
+
+def _is_steady_increment(
+    prev_sig: Optional[tuple],
+    prev_vals: Optional[List[int]],
+    cur_sig: tuple,
+    cur_vals: List[int],
+) -> bool:
+    return (
+        prev_sig is not None
+        and prev_vals is not None
+        and prev_sig == cur_sig
+        and len(prev_vals) == len(cur_vals)
+        and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
+    )
+
+
+# ---- task-scoped bind state ----
+@dataclass
+class _BindState:
+    """Cached bind metadata for steady-state decode increment detection."""
+
+    sig: tuple
+    seq_lens: List[int]
+
+
+# ---- pool + manager ----
+
+
+class PagePool:
+    """Physical KV cache: buffers + req-to-token table + allocation strategy + bind.
+
+    Does not know about tasks — task lifecycle is managed by
+    :class:`TaskCacheManager`, which holds a reference to this pool.
+    """
+
+    def __init__(
+        self,
+        n_layers: int,
+        n_kv_heads: int,
+        head_dim: int,
+        max_batch_size: int,
+        max_seq_len: 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.device = device
+        self.dtype = dtype
+        self.n_layers = n_layers
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+
+        self.contiguous = n_tokens is None
+        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
+        )
+        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
+
+        if self.contiguous:
+            for i in range(max_batch_size):
+                self._req_pool.req_to_token[i] = torch.arange(
+                    i * max_seq_len, (i + 1) * max_seq_len, device=device
+                )
+            self._strategy: AllocationStrategy = ContiguousStrategy()
+        else:
+            n_pages = self.n_tokens // page_size
+            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
+            )
+
+    @property
+    def strategy(self) -> AllocationStrategy:
+        return self._strategy
+
+    @property
+    def req_pool(self) -> ReqToTokenPool:
+        return self._req_pool
+
+    def bind_tasks(
+        self,
+        req_indices: List[int],
+        seq_lens: List[int],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+        incremental: bool = False,
+    ) -> KVCache:
+        """Assemble the ``KVCache`` metadata for a batch of tasks.
+
+        Args:
+            req_indices: request slot indices (from ``ReqToTokenPool``).
+            seq_lens:    current sequence length per task.
+            workspace:   pre-allocated fixed-shape buffers (CUDA-graph safe).
+            start_pos:   if set, produce **prefill** cache (full q_len range).
+                         If ``None``, produce **decode** cache (last position).
+            incremental: if ``True``, reuse workspace state from previous step
+                         by incrementing counters in-place (decode hot path).
+
+        Returns:
+            ``KVCache`` dataclass with the correct output shapes for the
+            attention backend (prefill: ``[B, q_len]``, decode: ``[B, 1]``).
+        """
+        if device is None:
+            device = workspace.device
+        b = len(req_indices)
+
+        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
+
+        if incremental:
+            sl_buf[:b] += 1
+            kvp_buf[: b + 1] += inc_buf[: b + 1]
+        else:
+            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]
+
+        if start_pos is not None:
+            # ---- prefill: out_cache_loc covers prefix range [start_pos:seq_len] ----
+            seq_len = seq_lens[0]
+            out_cache_loc = self._req_pool.req_to_token[
+                req_pool_indices, start_pos:seq_len
+            ]
+            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 = decode_out = None
+        else:
+            # ---- decode: out_cache_loc is a single column (last position) ----
+            write_pos = seq_lens_t - 1
+            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
+            ocl_buf[:b].copy_(loc)
+            out_cache_loc = ocl_buf[:b]
+            qo_indptr = None
+            decode_o_part = getattr(workspace, "decode_o_part", None)
+            decode_ml_part = getattr(workspace, "decode_ml_part", None)
+            decode_out = getattr(workspace, "decode_out", None)
+
+        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,
+            max_len=max(seq_lens),
+            kv_indptr=kv_indptr,
+            qo_indptr=qo_indptr,
+            decode_o_part=decode_o_part,
+            decode_ml_part=decode_ml_part,
+            decode_out=decode_out,
+        )
+
+
+class TaskCacheManager:
+    """Task ↔ KV slot lifecycle manager.
+
+    Sole owner of ``task_id → TaskCacheState``.  Delegates physical slot
+    allocation to the strategy (via ``pool.strategy``) and KV bind to
+    ``pool.bind_tasks()``.
+
+    Usage::
+
+        pool = PagePool(...)
+        mgr = TaskCacheManager(pool)
+        mgr.task_alloc("req_1", [101, 202, 303])
+        ...
+        kv = mgr.bind(["req_1"], workspace)
+    """
+
+    def __init__(self, pool: PagePool):
+        self._pool = pool
+        self._strategy = pool.strategy
+        self._req_pool = pool.req_pool
+        self._max_seq_len = pool.max_seq_len
+        self._states: Dict[str, TaskCacheState] = {}
+        self._bind_state: Optional[_BindState] = None
+        self._bind_was_steady = False
+
+    # -- public task lifecycle --
+
+    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 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 list(prompt_ids) + list(output_ids[:-1])
+
+    # -- bind (assemble KVCache for the model forward) --
+
+    def bind(
+        self,
+        task_ids: List[str],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+    ) -> KVCache:
+        """Build ``KVCache`` for an ordered list of task IDs."""
+        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:
+        return self._bind_was_steady
+
+    # -- internals --
+
+    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)

astrai/inference/cache/strategy.py

@@ -0,0 +1,320 @@
+"""KV cache allocation layer.
+
+Encapsulates the physical slot allocation policy, isolated from GPU buffers
+and task lifecycle management.
+
+- ``TaskCacheState``: data contract between strategy and manager (per-task slot state)
+- ``Allocator``:       bitmask-based page allocator with LRU eviction
+- ``RadixCache``:      page-granular prefix index (exact token match)
+- ``AllocationStrategy``: ABC for physical slot allocation
+- ``ContiguousStrategy``: statically partitioned, no dynamic allocation
+- ``PagedStrategy``:    dynamic paged allocation from a shared pool
+"""
+
+import threading
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, OrderedDict
+
+import torch
+
+from astrai.inference.cache.buffer import ReqToTokenPool
+
+# ---- data contract: per-task slot state ----
+
+
+@dataclass
+class TaskCacheState:
+    """Per-task cache allocation state.
+
+    Co-locates all task-owned cache metadata so the alloc/free/extend
+    lifecycle is atomic.  Owned by ``TaskCacheManager``, consumed by
+    every ``AllocationStrategy`` method.
+    """
+
+    req_idx: int
+    length: int = 0
+    cached: int = 0
+    pages: List[int] = field(default_factory=list)
+
+
+# ---- allocation primitives ----
+
+
+class Allocator:
+    """Bitmask-based page allocator with ref-counting and LRU eviction."""
+
+    def __init__(self, n_pages: int):
+        self._free_mask = (1 << n_pages) - 1
+        self._refs: List[int] = [0] * n_pages
+        self._lru: OrderedDict[int, None] = OrderedDict()
+        self.on_evict: Optional[Callable[[int], None]] = None
+        self._lock = threading.Lock()
+
+    def alloc(self) -> int:
+        with self._lock:
+            if self._free_mask:
+                lsb = self._free_mask & -self._free_mask
+                idx = lsb.bit_length() - 1
+                self._free_mask ^= lsb
+                self._refs[idx] = 1
+                return idx
+            if self._lru:
+                idx, _ = self._lru.popitem(last=False)
+                if self.on_evict:
+                    self.on_evict(idx)
+                self._refs[idx] = 1
+                self._free_mask &= ~(1 << idx)
+                return idx
+            return -1
+
+    def free(self, idx: int, keep_cached: bool = False):
+        with self._lock:
+            self._refs[idx] -= 1
+            if self._refs[idx] == 0:
+                if keep_cached:
+                    self._lru[idx] = None
+                else:
+                    self._free_mask |= 1 << idx
+
+    def inc_ref(self, idx: int):
+        with self._lock:
+            self._refs[idx] += 1
+            self._lru.pop(idx, None)
+
+    def ref_count(self, idx: int) -> int:
+        with self._lock:
+            return self._refs[idx]
+
+    def touch(self, idx: int):
+        with self._lock:
+            if idx in self._lru:
+                self._lru.move_to_end(idx)
+
+
+class RadixNode:
+    """A page-aligned edge in the CPU-side prefix radix trie."""
+
+    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")
+
+    def __init__(self, parent=None, tokens=(), page_idx=None):
+        self.parent = parent
+        self.children: Dict[tuple, "RadixNode"] = {}
+        self.page_idx = page_idx
+        self.tokens = tuple(tokens)
+        self.lock_ref = 0
+
+
+class RadixCache:
+    """Page-granular radix prefix index with exact token matching."""
+
+    def __init__(self, page_size: int):
+        self._page_size = page_size
+        self._root = RadixNode()
+        self._page_to_node: Dict[int, RadixNode] = {}
+        self._lock = threading.Lock()
+
+    def evict(self, idx: int):
+        with self._lock:
+            node = self._page_to_node.pop(idx, None)
+            if node is None:
+                return
+            node.page_idx = None
+            parent = node.parent
+            if parent is not None:
+                parent.children.pop(node.tokens, None)
+
+    def has_page(self, idx: int) -> bool:
+        with self._lock:
+            return idx in self._page_to_node
+
+    def lookup(self, token_ids: List[int]) -> List[int]:
+        with self._lock:
+            full_pages = len(token_ids) // self._page_size
+            hits: List[int] = []
+            node = self._root
+            for i in range(full_pages):
+                start = i * self._page_size
+                page_tokens = tuple(token_ids[start : start + self._page_size])
+                child = node.children.get(page_tokens)
+                if child is None or child.page_idx is None:
+                    break
+                hits.append(child.page_idx)
+                node = child
+            return hits
+
+    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
+        with self._lock:
+            full_pages = len(token_ids) // self._page_size
+            if logical_page_idx >= full_pages:
+                return
+            old = self._page_to_node.pop(page_idx, None)
+            if old is not None and old.parent is not None:
+                old.parent.children.pop(old.tokens, None)
+
+            node = self._root
+            for i in range(logical_page_idx + 1):
+                start = i * self._page_size
+                page_tokens = tuple(token_ids[start : start + self._page_size])
+                child = node.children.get(page_tokens)
+                if child is None:
+                    child = RadixNode(node, page_tokens)
+                    node.children[page_tokens] = child
+                node = child
+            if node.page_idx is not None and node.page_idx != page_idx:
+                replaced = node.page_idx
+                self._page_to_node.pop(replaced, None)
+            node.page_idx = page_idx
+            self._page_to_node[page_idx] = node
+
+    def release(self, pages: List[int]) -> None:
+        with self._lock:
+            for page_idx in pages:
+                node = self._page_to_node.get(page_idx)
+                if node is not None and node.lock_ref:
+                    node.lock_ref -= 1
+
+
+class AllocationStrategy(ABC):
+    """Physical slot allocation policy.
+
+    Subclasses implement the actual allocation semantics.  This ABC declares
+    the contract; there are no default implementations.
+    """
+
+    @abstractmethod
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool: ...
+
+    @abstractmethod
+    def free(self, state: TaskCacheState) -> None: ...
+
+    @abstractmethod
+    def extend(self, state: TaskCacheState, pos: int) -> bool: ...
+
+    @abstractmethod
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None: ...
+
+    @abstractmethod
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None: ...
+
+
+class ContiguousStrategy(AllocationStrategy):
+    """Static contiguous allocation: slots are pre-assigned at pool init.
+
+    No dynamic allocation or prefix caching.  All operations are no-ops
+    because ``ReqToTokenPool`` is pre-filled with contiguous ranges.
+    """
+
+    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)

astrai/inference/core/__init__.py

@@ -1,35 +0,0 @@
-"""Inference core: cache, executor, scheduler, task management."""
-
-from astrai.inference.core.cache import (
-    Allocator,
-    KVCache,
-    KVStorage,
-    PagePool,
-    RadixCache,
-    ReqToTokenPool,
-    TaskCacheManager,
-    page_hash,
-)
-from astrai.inference.core.executor import Executor
-from astrai.inference.core.metrics import MetricsCollector, TaskTiming
-from astrai.inference.core.scheduler import InferenceScheduler
-from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
-
-__all__ = [
-    "Allocator",
-    "KVCache",
-    "KVStorage",
-    "PagePool",
-    "RadixCache",
-    "ReqToTokenPool",
-    "TaskCacheManager",
-    "page_hash",
-    "Executor",
-    "InferenceScheduler",
-    "MetricsCollector",
-    "TaskTiming",
-    "STOP",
-    "Task",
-    "TaskManager",
-    "TaskStatus",
-]

astrai/inference/core/cache.py

@@ -1,647 +0,0 @@
-"""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 — 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.
-
-KVCache is a pure dataclass passed to the model for direct buffer access.
-
-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 abc import ABC
-from dataclasses import dataclass, field
-from typing import Callable, Dict, List, Optional, OrderedDict
-
-import torch
-from torch import Tensor
-
-from astrai.inference.core.workspace import InferenceWorkspace
-
-
-@dataclass
-class _BindState:
-    """Cached bind metadata for steady-state decode increment detection."""
-
-    sig: tuple
-    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]],
-    cur_sig: tuple,
-    cur_vals: List[int],
-) -> bool:
-    """True when the same ordered set has every value +1 from the previous step."""
-    return (
-        prev_sig is not None
-        and prev_vals is not None
-        and prev_sig == cur_sig
-        and len(prev_vals) == len(cur_vals)
-        and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
-    )
-
-
-def page_hash(
-    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
-) -> int:
-    start = page_idx * page_size
-    end = min(start + page_size, len(token_ids))
-    h = parent_hash
-    for i in range(start, end):
-        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
-    return h
-
-
-class Allocator:
-    """Bitmask-based page allocator with ref-counting and LRU eviction."""
-
-    def __init__(self, n_pages: int):
-        self._free_mask = (1 << n_pages) - 1
-        self._refs: List[int] = [0] * n_pages
-        self._lru: OrderedDict[int, None] = OrderedDict()
-        self.on_evict: Optional[Callable[[int], None]] = None
-        self._lock = threading.Lock()
-
-    def alloc(self) -> int:
-        with self._lock:
-            if self._free_mask:
-                lsb = self._free_mask & -self._free_mask
-                idx = lsb.bit_length() - 1
-                self._free_mask ^= lsb
-                self._refs[idx] = 1
-                return idx
-            if self._lru:
-                idx, _ = self._lru.popitem(last=False)
-                if self.on_evict:
-                    self.on_evict(idx)
-                self._refs[idx] = 1
-                self._free_mask &= ~(1 << idx)
-                return idx
-            return -1
-
-    def free(self, idx: int, keep_cached: bool = False):
-        with self._lock:
-            self._refs[idx] -= 1
-            if self._refs[idx] == 0:
-                if keep_cached:
-                    self._lru[idx] = None
-                else:
-                    self._free_mask |= 1 << idx
-
-    def inc_ref(self, idx: int):
-        with self._lock:
-            self._refs[idx] += 1
-            self._lru.pop(idx, None)
-
-    def ref_count(self, idx: int) -> int:
-        with self._lock:
-            return self._refs[idx]
-
-    def touch(self, idx: int):
-        with self._lock:
-            if idx in self._lru:
-                self._lru.move_to_end(idx)
-
-
-class RadixNode:
-    """A page-aligned edge in the CPU-side prefix radix."""
-
-    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")
-
-    def __init__(self, parent=None, tokens=(), page_idx=None):
-        self.parent = parent
-        self.children: Dict[tuple, "RadixNode"] = {}
-        self.page_idx = page_idx
-        self.tokens = tuple(tokens)
-        self.lock_ref = 0
-
-
-class RadixCache:
-    """Page-granular radix prefix index with exact token matching."""
-
-    def __init__(self, page_size: int):
-        self._page_size = page_size
-        self._root = RadixNode()
-        self._page_to_node: Dict[int, RadixNode] = {}
-        self._lock = threading.Lock()
-
-    def evict(self, idx: int):
-        with self._lock:
-            node = self._page_to_node.pop(idx, None)
-            if node is None:
-                return
-            node.page_idx = None
-            parent = node.parent
-            if parent is not None:
-                parent.children.pop(node.tokens, None)
-
-    def has_page(self, idx: int) -> bool:
-        with self._lock:
-            return idx in self._page_to_node
-
-    def lookup(self, token_ids: List[int]) -> List[int]:
-        with self._lock:
-            full_pages = len(token_ids) // self._page_size
-            hits: List[int] = []
-            node = self._root
-            for i in range(full_pages):
-                start = i * self._page_size
-                page_tokens = tuple(token_ids[start : start + self._page_size])
-                child = node.children.get(page_tokens)
-                if child is None or child.page_idx is None:
-                    break
-                hits.append(child.page_idx)
-                node = child
-            return hits
-
-    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
-        with self._lock:
-            full_pages = len(token_ids) // self._page_size
-            if logical_page_idx >= full_pages:
-                return
-            old = self._page_to_node.pop(page_idx, None)
-            if old is not None and old.parent is not None:
-                old.parent.children.pop(old.tokens, None)
-
-            node = self._root
-            for i in range(logical_page_idx + 1):
-                start = i * self._page_size
-                page_tokens = tuple(token_ids[start : start + self._page_size])
-                child = node.children.get(page_tokens)
-                if child is None:
-                    child = RadixNode(node, page_tokens)
-                    node.children[page_tokens] = child
-                node = child
-            if node.page_idx is not None and node.page_idx != page_idx:
-                replaced = node.page_idx
-                self._page_to_node.pop(replaced, None)
-            node.page_idx = page_idx
-            self._page_to_node[page_idx] = node
-
-    def release(self, pages: List[int]) -> None:
-        with self._lock:
-            for page_idx in pages:
-                node = self._page_to_node.get(page_idx)
-                if node is not None and node.lock_ref:
-                    node.lock_ref -= 1
-
-
-class ReqToTokenPool:
-    """Maps [req_idx, pos] -> physical token slot in KV storage.
-
-    Each row is one request; each column is a sequence position. The value
-    at [req_idx, pos] is the flat index into the KV storage buffers.
-    """
-
-    def __init__(self, size: int, max_context_len: int, device: torch.device):
-        self.size = size
-        self.max_context_len = max_context_len
-        self.req_to_token = torch.zeros(
-            (size, max_context_len), dtype=torch.long, device=device
-        )
-        self.free_slots = list(range(size))
-        self._lock = threading.Lock()
-
-    def alloc(self, num_reqs: int) -> Optional[List[int]]:
-        with self._lock:
-            if num_reqs > len(self.free_slots):
-                return None
-            slots = self.free_slots[:num_reqs]
-            self.free_slots = self.free_slots[num_reqs:]
-            return slots
-
-    def free(self, req_indices: List[int]):
-        with self._lock:
-            self.free_slots.extend(req_indices)
-
-    def write(self, indices, values):
-        self.req_to_token[indices] = values
-
-
-class KVStorage:
-    """Token-level KV cache storage.
-
-    Buffers: [n_layers, size, n_kv_heads, head_dim]. Each token occupies
-    one slot indexed by ReqToTokenPool.
-    """
-
-    def __init__(
-        self,
-        size: int,
-        n_layers: int,
-        n_kv_heads: int,
-        head_dim: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ):
-        self.size = size
-        self.k_buffer = torch.empty(
-            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
-        )
-        self.v_buffer = torch.empty(
-            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
-        )
-
-    def get_key_buffer(self, layer_id: int) -> Tensor:
-        return self.k_buffer[layer_id]
-
-    def get_value_buffer(self, layer_id: int) -> Tensor:
-        return self.v_buffer[layer_id]
-
-    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
-        self.k_buffer[layer_id, loc] = k
-        self.v_buffer[layer_id, loc] = v
-
-
-@dataclass
-class KVCache:
-    """Pure data struct passed to model for KV cache I/O.
-
-    The attention layer does raw buffer indexing — no methods, no abstraction.
-    """
-
-    k_buffer: Tensor
-    v_buffer: Tensor
-    req_to_token: Tensor
-    req_pool_indices: Tensor
-    seq_lens: Tensor
-    out_cache_loc: Tensor
-    max_len: int = 0
-    kv_indptr: Optional[Tensor] = None
-    qo_indptr: Optional[Tensor] = None
-    decode_o_part: Optional[Tensor] = None
-    decode_ml_part: Optional[Tensor] = None
-    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:
-    """Physical KV cache: buffers + req-pool + allocation strategy + bind.
-
-    Does not know about tasks — task lifecycle is managed by
-    :class:`TaskCacheManager`, which holds a reference to this pool.
-    """
-
-    def __init__(
-        self,
-        n_layers: int,
-        n_kv_heads: int,
-        head_dim: int,
-        max_batch_size: int,
-        max_seq_len: 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.device = device
-        self.dtype = dtype
-        self.n_layers = n_layers
-        self.n_kv_heads = n_kv_heads
-        self.head_dim = head_dim
-
-        self.contiguous = n_tokens is None
-        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
-        )
-        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
-
-        if self.contiguous:
-            for i in range(max_batch_size):
-                self._req_pool.req_to_token[i] = torch.arange(
-                    i * max_seq_len, (i + 1) * max_seq_len, device=device
-                )
-            self._strategy = AllocationStrategy()
-        else:
-            n_pages = self.n_tokens // page_size
-            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,
-        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
-        b = len(req_indices)
-
-        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
-
-        if incremental:
-            sl_buf[:b] += 1
-            kvp_buf[: b + 1] += inc_buf[: b + 1]
-        else:
-            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]
-
-        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
-            ]
-            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 = decode_out = None
-        else:
-            write_pos = seq_lens_t - 1
-            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
-            ocl_buf[:b].copy_(loc)
-            out_cache_loc = ocl_buf[:b]
-            qo_indptr = None
-            decode_o_part = getattr(workspace, "decode_o_part", None)
-            decode_ml_part = getattr(workspace, "decode_ml_part", None)
-            decode_out = getattr(workspace, "decode_out", None)
-
-        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,
-            max_len=max(seq_lens),
-            kv_indptr=kv_indptr,
-            qo_indptr=qo_indptr,
-            decode_o_part=decode_o_part,
-            decode_ml_part=decode_ml_part,
-            decode_out=decode_out,
-        )
-
-
-class TaskCacheManager:
-    """Task <-> KV slot lifecycle manager.
-
-    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.
-    """
-
-    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

astrai/inference/engine.py

@@ -8,9 +8,9 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
 import torch
 import torch.nn as nn
 
-from astrai.inference.core.cache import PagePool
+from astrai.inference.cache import PagePool
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
-from astrai.inference.core.task import STOP
+from astrai.inference.task import STOP
 from astrai.tokenize import AutoTokenizer
 
 

astrai/inference/network/__init__.py

@@ -4,8 +4,7 @@
 lazy singleton FastAPI instance.
 """
 
-from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
+from astrai.inference.network.app import (
-from astrai.inference.api.server import (
     AnthropicMessage,
     ChatCompletionRequest,
     ChatMessage,
@@ -15,7 +14,8 @@ from astrai.inference.api.server import (
     get_app,
     run_server,
 )
-from astrai.inference.api.tool_parser import (
+from astrai.inference.network.protocol import GenContext, ProtocolHandler, StopChecker
+from astrai.inference.network.tool_parser import (
     BaseToolParser,
     SimpleJsonToolParser,
     ToolParserFactory,

astrai/inference/network/anthropic.py

@@ -6,13 +6,13 @@ from typing import Any, Dict, List, Tuple, Union
 
 from pydantic import BaseModel
 
-from astrai.inference.api.protocol import (
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.protocol import (
     GenContext,
     ResponseBuilder,
     StopInfo,
     sse_event,
 )
-from astrai.inference.engine import InferenceEngine
 
 
 def _extract_text(content: Union[str, List[Dict[str, Any]]]) -> str:

astrai/inference/network/app.py

@@ -18,10 +18,10 @@ import uvicorn
 from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
 
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import ProtocolHandler
 from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
+from astrai.inference.network.protocol import ProtocolHandler
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 

astrai/inference/network/openai.py

@@ -7,14 +7,14 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 from pydantic import BaseModel
 
-from astrai.inference.api.protocol import (
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.protocol import (
     GenContext,
     ResponseBuilder,
     StopInfo,
     sse_event,
 )
-from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
+from astrai.inference.network.tool_parser import BaseToolParser, ToolParserFactory
-from astrai.inference.engine import InferenceEngine
 
 logger = logging.getLogger(__name__)
 

astrai/inference/runtime/__init__.py

@@ -0,0 +1,25 @@
+"""Execution primitives: forward passes, CUDA graphs, and sampling."""
+
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.runtime.sample import (
+    BaseSamplingStrategy,
+    FrequencyPenaltyStrategy,
+    SamplingPipeline,
+    TemperatureStrategy,
+    TopKStrategy,
+    TopPStrategy,
+    sample,
+)
+
+__all__ = [
+    "Executor",
+    "CudaGraphContext",
+    "BaseSamplingStrategy",
+    "FrequencyPenaltyStrategy",
+    "SamplingPipeline",
+    "TemperatureStrategy",
+    "TopKStrategy",
+    "TopPStrategy",
+    "sample",
+]

astrai/inference/runtime/executor.py

@@ -13,11 +13,11 @@ from astrai.extension.attention_backend import (
     attn_backend,
     get_backend,
 )
-from astrai.inference.core.cache import PagePool, TaskCacheManager
+from astrai.inference.cache import PagePool, TaskCacheManager
-from astrai.inference.core.graph import CudaGraphContext
+from astrai.inference.runtime.graph import CudaGraphContext
-from astrai.inference.core.task import Task
+from astrai.inference.runtime.sample import sample
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.task import Task
-from astrai.inference.sample import sample
+from astrai.inference.workspace import InferenceWorkspace
 from astrai.model.automodel import AutoModel
 
 logger = logging.getLogger(__name__)
@@ -369,12 +369,8 @@ class Executor:
 
         kv_cache = self.task_cache.bind(task_ids, ws)
 
-        if self.task_cache.bind_was_steady:
+        if self.task_cache.bind_was_steady and self._decode_cache is not None:
-            info = (
+            info = self._decode_cache.sampling_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
         else:
             info = _build_sampling_batch_info(tasks, self.device)

astrai/inference/scheduler.py

@@ -5,10 +5,10 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 
-from astrai.inference.core.cache import PagePool, TaskCacheManager
+from astrai.inference.cache import PagePool, TaskCacheManager
-from astrai.inference.core.executor import Executor
+from astrai.inference.metrics import MetricsCollector
-from astrai.inference.core.metrics import MetricsCollector
+from astrai.inference.runtime.executor import Executor
-from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
+from astrai.inference.task import STOP, Task, TaskManager, TaskStatus
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 
@@ -59,12 +59,7 @@ class InferenceScheduler:
 
         self._metrics = MetricsCollector()
 
-        self._task_cache = TaskCacheManager(
+        self._task_cache = TaskCacheManager(self._cache)
-            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,

astrai/inference/task.py

@@ -7,7 +7,7 @@ from typing import Any, Callable, Deque, Dict, List, Optional
 
 from tokenizers.decoders import DecodeStream
 
-from astrai.inference.core.metrics import MetricsCollector
+from astrai.inference.metrics import MetricsCollector
 from astrai.tokenize.tokenizer import AutoTokenizer
 
 STOP = object()

astrai/logging.py

@@ -10,7 +10,7 @@ def setup_logging(level: str = "INFO"):
 
     Level names: ``DEBUG``, ``INFO``, ``WARNING``, ``ERROR``, ``CRITICAL``.
     ``DEBUG`` enables per-step prefill/decode timing logs
-    (:func:`astrai.inference.core.executor.timed`).
+    (:func:`astrai.inference.runtime.executor.timed`).
     """
     logger = logging.getLogger("astrai")
     if logger.handlers:

astrai/model/components/attention.py

@@ -8,7 +8,7 @@ from torch import Tensor
 from astrai.extension import attention
 from astrai.extension.rotary_backend import apply_rotary_emb
 from astrai.factory import BaseFactory
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
 

astrai/model/components/decoder_block.py

@@ -4,7 +4,7 @@ from typing import Optional, TypedDict
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.components.attention import AttnFactory
 from astrai.model.components.mlp import FFNFactory, RouterStats
 from astrai.model.components.norm import RMSNorm

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 KVCache
+from astrai.inference.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

astrai/trainer/rollout.py

@@ -6,7 +6,7 @@ Provides:
 - :class:`BaseRewardModel` — pluggable reward interface
 - :class:`RolloutGenerator` — KV-cache-backed generation of grouped
   responses + decoding (no reward); delegates the generation loop to
-  :class:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`
+  :class:`~astrai.inference.scheduler.InferenceScheduler.run_batch`
   so rollout and the production inference server share one code path
 - :class:`RolloutRunner` — orchestrates generation + scoring with a
   step-driven cache; its ``__call__`` returns ``(RolloutResult, is_fresh)``
@@ -20,7 +20,7 @@ from typing import Dict, List, Optional, Tuple
 import torch
 from torch import Tensor
 
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 
 
 @dataclass(kw_only=True)
@@ -101,7 +101,7 @@ class RolloutGenerator:
     """Pure generation + decoding for a group of responses per prompt.
 
     Delegates the prefill/decode loop to
-    :meth:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`,
+    :meth:`~astrai.inference.scheduler.InferenceScheduler.run_batch`,
     which uses a real KV cache (no O(n²) recompute).  Has no dependency
     on any reward model; can be reused in isolation for offline
     generation, qualitative sampling, or eval pipelines.

astrai/trainer/train_context.py

@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader, random_split
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import RDSampler
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size

scripts/tools/benchmark.py

@@ -7,9 +7,9 @@ import torch
 
 from astrai.config import BaseModelConfig, ConfigFactory
 from astrai.extension import ATTN_BACKEND, AttentionBackendFactory, attn_backend
-from astrai.inference.core.cache import PagePool, TaskCacheManager
+from astrai.inference.cache import PagePool, TaskCacheManager
-from astrai.inference.core.graph import CudaGraphContext
+from astrai.inference.runtime.graph import CudaGraphContext
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.workspace import InferenceWorkspace
 from astrai.model import AutoModel, AutoRegressiveLM
 
 _DTYPES = ["bfloat16", "float16", "float32"]
@@ -93,12 +93,7 @@ class GenerationBenchmark:
 
     @staticmethod
     def _make_task_cache(pool: PagePool) -> TaskCacheManager:
-        return TaskCacheManager(
+        return TaskCacheManager(pool)
-            strategy=pool._strategy,
-            req_pool=pool._req_pool,
-            max_seq_len=pool.max_seq_len,
-            pool=pool,
-        )
 
     def _run_prefill(
         self,

tests/extension/test_backend_equivalence.py

@@ -7,18 +7,13 @@ 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, TaskCacheManager
+from astrai.inference.cache import PagePool, TaskCacheManager
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.workspace import InferenceWorkspace
 from tests.extension.conftest import D, skip_no_kernel
 
 
 def _mk_task_cache(pool: PagePool) -> TaskCacheManager:
-    return TaskCacheManager(
+    return TaskCacheManager(pool)
-        strategy=pool._strategy,
-        req_pool=pool._req_pool,
-        max_seq_len=pool.max_seq_len,
-        pool=pool,
-    )
 
 
 def _ws(pool: PagePool) -> InferenceWorkspace:
@@ -184,7 +179,7 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
 @skip_no_kernel
 def test_run_batch_cuda_matches_torch_greedy(cuda_model):
     """Greedy decode (temperature=0) should produce identical tokens."""
-    from astrai.inference.core.scheduler import InferenceScheduler
+    from astrai.inference.scheduler import InferenceScheduler
     from tests.helpers import FakeTokenizer
 
     model, _ = cuda_model

tests/inference/test_cache.py

@@ -11,7 +11,7 @@ from astrai.inference import (
     TaskCacheManager,
     page_hash,
 )
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.workspace import InferenceWorkspace
 
 
 def _ws(pool: PagePool) -> InferenceWorkspace:
@@ -27,12 +27,7 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
 
 
 def _make_task_cache(pool: PagePool) -> TaskCacheManager:
-    return TaskCacheManager(
+    return TaskCacheManager(pool)
-        strategy=pool._strategy,
-        req_pool=pool._req_pool,
-        max_seq_len=pool.max_seq_len,
-        pool=pool,
-    )
 
 
 # ---- page_hash ----
@@ -345,7 +340,7 @@ def test_page_pool_paged_task_alloc():
     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]
+    assert pool.req_pool.req_to_token[state.req_idx, 0].item() == state.pages[0]
 
 
 def test_page_pool_paged_task_extend():
@@ -354,7 +349,7 @@ def test_page_pool_paged_task_extend():
     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()
+    slot = pool.req_pool.req_to_token[req_idx, 4].item()
     assert slot >= 0
 
 
@@ -364,7 +359,7 @@ def test_page_pool_paged_task_free_releases_slots():
     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
+    assert len(pool.req_pool.free_slots) == 4
 
 
 def test_page_pool_paged_bind_roundtrip():

tests/inference/test_protocol.py

@@ -5,9 +5,9 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import GenContext, StopChecker, StopInfo
+from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo
 
 
 def _make_ctx(**kwargs):

tests/inference/test_sample.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from astrai.inference.sample import (
+from astrai.inference.runtime.sample import (
     FrequencyPenaltyStrategy,
     SamplingPipeline,
     TemperatureStrategy,
@@ -268,7 +268,7 @@ def test_sample_return_logprobs_matches_manual_computation():
     logits = torch.randn(2, 30)
     tokens, logprobs = sample(logits, temperature=0.7, top_p=0.95, return_logprobs=True)
     # Recompute with the same pipeline
-    from astrai.inference.sample import (
+    from astrai.inference.runtime.sample import (
         SamplingPipeline,
         TemperatureStrategy,
         TopPStrategy,

tests/inference/test_scheduler.py

@@ -38,8 +38,8 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
     """Test concurrent add_task operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
 
-    with patch("astrai.inference.core.scheduler.AutoModel"):
+    with patch("astrai.inference.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
             scheduler = InferenceScheduler(
                 model=mock_model,
                 tokenizer=mock_tokenizer,
@@ -77,8 +77,8 @@ def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
     """Test concurrent add and remove task operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
 
-    with patch("astrai.inference.core.scheduler.AutoModel"):
+    with patch("astrai.inference.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
             scheduler = InferenceScheduler(
                 model=mock_model,
                 tokenizer=mock_tokenizer,
@@ -126,8 +126,8 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
     """Test concurrent get_stats operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
 
-    with patch("astrai.inference.core.scheduler.AutoModel"):
+    with patch("astrai.inference.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
             scheduler = InferenceScheduler(
                 model=mock_model,
                 tokenizer=mock_tokenizer,

tests/inference/test_tool_parser.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from astrai.inference.api.tool_parser import (
+from astrai.inference.network.tool_parser import (
     _TOOL_CALL_HEAD_RE,
     BaseToolParser,
     SimpleJsonToolParser,

tests/trainer/test_rollout.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.trainer.rollout import (
     BaseRewardModel,
     RawRollout,