diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index da6c356..c06cd09 100644 --- a/astrai/extension/attention_backend.py +++ b/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] = [] diff --git a/astrai/inference/__init__.py b/astrai/inference/__init__.py index b639855..d44a411 100644 --- a/astrai/inference/__init__.py +++ b/astrai/inference/__init__.py @@ -1,15 +1,29 @@ """Inference module for continuous batching. -Layers: - - core/: Core inference loop (cache, executor, scheduler, task) - - api/: HTTP orchestration (ProtocolHandler, server) - - protocols/: Response builders (OpenAI, Anthropic) - - transport/: SSE transport utilities - - engine.py: Facade (InferenceEngine) - - sample.py: Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy) +Subpackages: + - cache/: KV cache (buffers, strategies, pool) + - runtime/: Execution + sampling (executor, CUDA graph, sampling strategies) + - task/: Request lifecycle + performance metrics + - 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) """ -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.api.openai import OpenAIResponseBuilder -from astrai.inference.core import ( - STOP, - 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 ( +from astrai.inference.network.anthropic import AnthropicResponseBuilder +from astrai.inference.network.openai import OpenAIResponseBuilder +from astrai.inference.runtime.executor import Executor +from astrai.inference.runtime.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", diff --git a/astrai/inference/cache/__init__.py b/astrai/inference/cache/__init__.py new file mode 100644 index 0000000..45a357c --- /dev/null +++ b/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", +] diff --git a/astrai/inference/cache/buffer.py b/astrai/inference/cache/buffer.py new file mode 100644 index 0000000..5f72ace --- /dev/null +++ b/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 diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py new file mode 100644 index 0000000..99e7b6f --- /dev/null +++ b/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) diff --git a/astrai/inference/cache/strategy.py b/astrai/inference/cache/strategy.py new file mode 100644 index 0000000..6b3e711 --- /dev/null +++ b/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) diff --git a/astrai/inference/core/__init__.py b/astrai/inference/core/__init__.py deleted file mode 100644 index e5590d9..0000000 --- a/astrai/inference/core/__init__.py +++ /dev/null @@ -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", -] diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py deleted file mode 100644 index 6274219..0000000 --- a/astrai/inference/core/cache.py +++ /dev/null @@ -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 diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index 0c739df..e5d55ff 100644 --- a/astrai/inference/engine.py +++ b/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.core.scheduler import InferenceScheduler -from astrai.inference.core.task import STOP +from astrai.inference.cache import PagePool +from astrai.inference.scheduler import InferenceScheduler +from astrai.inference.task import STOP from astrai.tokenize import AutoTokenizer diff --git a/astrai/inference/core/metrics.py b/astrai/inference/metrics.py similarity index 100% rename from astrai/inference/core/metrics.py rename to astrai/inference/metrics.py diff --git a/astrai/inference/api/__init__.py b/astrai/inference/network/__init__.py similarity index 80% rename from astrai/inference/api/__init__.py rename to astrai/inference/network/__init__.py index 35d7119..36b19ad 100644 --- a/astrai/inference/api/__init__.py +++ b/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.api.server import ( +from astrai.inference.network.app 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, diff --git a/astrai/inference/api/anthropic.py b/astrai/inference/network/anthropic.py similarity index 99% rename from astrai/inference/api/anthropic.py rename to astrai/inference/network/anthropic.py index fbc6827..7cf0f1c 100644 --- a/astrai/inference/api/anthropic.py +++ b/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: diff --git a/astrai/inference/api/server.py b/astrai/inference/network/app.py similarity index 96% rename from astrai/inference/api/server.py rename to astrai/inference/network/app.py index b1e9387..4292b93 100644 --- a/astrai/inference/api/server.py +++ b/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 diff --git a/astrai/inference/api/openai.py b/astrai/inference/network/openai.py similarity index 98% rename from astrai/inference/api/openai.py rename to astrai/inference/network/openai.py index 948ffb9..0688c19 100644 --- a/astrai/inference/api/openai.py +++ b/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.engine import InferenceEngine +from astrai.inference.network.tool_parser import BaseToolParser, ToolParserFactory logger = logging.getLogger(__name__) diff --git a/astrai/inference/api/protocol.py b/astrai/inference/network/protocol.py similarity index 100% rename from astrai/inference/api/protocol.py rename to astrai/inference/network/protocol.py diff --git a/astrai/inference/api/tool_parser.py b/astrai/inference/network/tool_parser.py similarity index 100% rename from astrai/inference/api/tool_parser.py rename to astrai/inference/network/tool_parser.py diff --git a/astrai/inference/runtime/__init__.py b/astrai/inference/runtime/__init__.py new file mode 100644 index 0000000..dda914e --- /dev/null +++ b/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", +] diff --git a/astrai/inference/core/executor.py b/astrai/inference/runtime/executor.py similarity index 96% rename from astrai/inference/core/executor.py rename to astrai/inference/runtime/executor.py index 4f45284..2617179 100644 --- a/astrai/inference/core/executor.py +++ b/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.core.graph import CudaGraphContext -from astrai.inference.core.task import Task -from astrai.inference.core.workspace import InferenceWorkspace -from astrai.inference.sample import sample +from astrai.inference.cache import PagePool, TaskCacheManager +from astrai.inference.runtime.graph import CudaGraphContext +from astrai.inference.runtime.sample import sample +from astrai.inference.task import Task +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: - info = ( - self._decode_cache.sampling_info - if self._decode_cache is not None - else _build_sampling_batch_info(tasks, self.device) - ) + if self.task_cache.bind_was_steady and self._decode_cache is not None: + info = self._decode_cache.sampling_info ws.position_ids[:b] += 1 else: info = _build_sampling_batch_info(tasks, self.device) diff --git a/astrai/inference/core/graph.py b/astrai/inference/runtime/graph.py similarity index 100% rename from astrai/inference/core/graph.py rename to astrai/inference/runtime/graph.py diff --git a/astrai/inference/sample.py b/astrai/inference/runtime/sample.py similarity index 100% rename from astrai/inference/sample.py rename to astrai/inference/runtime/sample.py diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/scheduler.py similarity index 96% rename from astrai/inference/core/scheduler.py rename to astrai/inference/scheduler.py index 54c1f96..1ac0e76 100644 --- a/astrai/inference/core/scheduler.py +++ b/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.core.executor import Executor -from astrai.inference.core.metrics import MetricsCollector -from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus +from astrai.inference.cache import PagePool, TaskCacheManager +from astrai.inference.metrics import MetricsCollector +from astrai.inference.runtime.executor import Executor +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( - strategy=self._cache._strategy, - req_pool=self._cache._req_pool, - max_seq_len=self.max_seq_len, - pool=self._cache, - ) + self._task_cache = TaskCacheManager(self._cache) self._task_mgr = TaskManager( tokenizer=tokenizer, diff --git a/astrai/inference/core/task.py b/astrai/inference/task.py similarity index 99% rename from astrai/inference/core/task.py rename to astrai/inference/task.py index db3b1c4..70a3f83 100644 --- a/astrai/inference/core/task.py +++ b/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() diff --git a/astrai/inference/core/workspace.py b/astrai/inference/workspace.py similarity index 100% rename from astrai/inference/core/workspace.py rename to astrai/inference/workspace.py diff --git a/astrai/logging.py b/astrai/logging.py index d1fc8f1..5ea2c46 100644 --- a/astrai/logging.py +++ b/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: diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py index 4428ad0..8e06798 100644 --- a/astrai/model/components/attention.py +++ b/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 diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py index 7c4e804..65e1497 100644 --- a/astrai/model/components/decoder_block.py +++ b/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 diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index 7c2f38c..a2f2167 100644 --- a/astrai/model/transformer.py +++ b/astrai/model/transformer.py @@ -5,7 +5,7 @@ import torch.nn as nn from torch import Tensor from astrai.config.model_config import AutoRegressiveLMConfig -from astrai.inference.core.cache import 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 diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index e1c6621..165c24e 100644 --- a/astrai/trainer/rollout.py +++ b/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. diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index b81b0e2..c4d4f98 100644 --- a/astrai/trainer/train_context.py +++ b/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 diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index eb0937f..f80afbd 100644 --- a/scripts/tools/benchmark.py +++ b/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.core.graph import CudaGraphContext -from astrai.inference.core.workspace import InferenceWorkspace +from astrai.inference.cache import PagePool, TaskCacheManager +from astrai.inference.runtime.graph import CudaGraphContext +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( - strategy=pool._strategy, - req_pool=pool._req_pool, - max_seq_len=pool.max_seq_len, - pool=pool, - ) + return TaskCacheManager(pool) def _run_prefill( self, diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index 27f5e68..b691d7b 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/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.core.workspace import InferenceWorkspace +from astrai.inference.cache import PagePool, TaskCacheManager +from astrai.inference.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, - ) + return TaskCacheManager(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 diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index f383362..0ecc846 100644 --- a/tests/inference/test_cache.py +++ b/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( - strategy=pool._strategy, - req_pool=pool._req_pool, - max_seq_len=pool.max_seq_len, - pool=pool, - ) + return TaskCacheManager(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(): diff --git a/tests/inference/test_protocol.py b/tests/inference/test_protocol.py index dd9f9a1..99b2991 100644 --- a/tests/inference/test_protocol.py +++ b/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.api.openai import OpenAIResponseBuilder -from astrai.inference.api.protocol import GenContext, StopChecker, StopInfo +from astrai.inference.network.anthropic import AnthropicResponseBuilder +from astrai.inference.network.openai import OpenAIResponseBuilder +from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo def _make_ctx(**kwargs): diff --git a/tests/inference/test_sample.py b/tests/inference/test_sample.py index c8c54bb..63489dd 100644 --- a/tests/inference/test_sample.py +++ b/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, diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 8131b35..edeae22 100644 --- a/tests/inference/test_scheduler.py +++ b/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.core.scheduler.AutoTokenizer"): + with patch("astrai.inference.scheduler.AutoModel"): + 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.core.scheduler.AutoTokenizer"): + with patch("astrai.inference.scheduler.AutoModel"): + 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.core.scheduler.AutoTokenizer"): + with patch("astrai.inference.scheduler.AutoModel"): + with patch("astrai.inference.scheduler.AutoTokenizer"): scheduler = InferenceScheduler( model=mock_model, tokenizer=mock_tokenizer, diff --git a/tests/inference/test_tool_parser.py b/tests/inference/test_tool_parser.py index 5d27f52..ddfd19f 100644 --- a/tests/inference/test_tool_parser.py +++ b/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, diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index a71e365..4e999eb 100644 --- a/tests/trainer/test_rollout.py +++ b/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,