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
This commit is contained in:
@@ -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] = []
|
||||
|
||||
@@ -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",
|
||||
|
||||
Vendored
+27
@@ -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",
|
||||
]
|
||||
Vendored
+104
@@ -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
|
||||
Vendored
+351
@@ -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)
|
||||
Vendored
+320
@@ -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)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
@@ -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:
|
||||
@@ -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
|
||||
|
||||
@@ -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__)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
@@ -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()
|
||||
+1
-1
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user