ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: clean up inference design patterns 

1. KVCache base: add default task_cached/task_record_hashes, remove getattr from scheduler
2. Remove page_size param from scheduler constructor (ContiguousCache-only)
3. InferenceEngine expose cache param for KVCache injection
4. Rename page_cache -> kv_cache in Executor
5. Move stream_callback from Task to TaskManager._callbacks dict
6. TaskManager.clear_queues clears callbacks
This commit is contained in:
ViperEkura 2026-07-05 11:41:54 +08:00
parent 5416c2e8fb
commit 849e1e00a3
5 changed files with 39 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
astrai/inference/core/cache.py
View File

@ -303,6 +303,13 @@ class KVCache(ABC):
self, task_ids: List[str], total_len: int, device: torch.device
) -> CacheView: ...
def task_cached(self, task_id: str) -> int:
return 0
def task_record_hashes(
self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
): ...
class PageCacheView(CacheView):
"""Bundles Storage + page_table + total_len for attention layers."""

12
astrai/inference/core/executor.py
View File

@ -19,13 +19,13 @@ class Executor:
self,
model: AutoModel,
tokenizer: AutoTokenizer,
page_cache: KVCache,
kv_cache: KVCache,
device: Optional[str] = None,
dtype: Optional[torch.dtype] = None,
):
self.model = model
self.tokenizer = tokenizer
self.page_cache = page_cache
self.kv_cache = kv_cache
self.device = device or next(model.parameters()).device
self.dtype = dtype or next(model.parameters()).dtype
@ -52,9 +52,7 @@ class Executor:
)
.unsqueeze(0)
.expand(batch_sz, -1),
paged_cache=self.page_cache.bind_tasks(
task_ids, prompt_len, self.device
),
paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
)
def execute_decode(self, tasks: List[Task]) -> List[int]:
@ -81,9 +79,7 @@ class Executor:
with torch.inference_mode():
outputs = self.model(
input_ids.unsqueeze(1),
paged_cache=self.page_cache.bind_tasks(
task_ids, total_len, self.device
),
paged_cache=self.kv_cache.bind_tasks(task_ids, total_len, self.device),
position_ids=position_ids.unsqueeze(1),
)
logits = outputs["logits"][:, -1, :]

53
astrai/inference/core/scheduler.py
View File

@ -23,7 +23,6 @@ class InferenceScheduler:
max_batch_size: int = 16,
max_seq_len: Optional[int] = None,
max_prompt_len: int = 2048,
page_size: int = 64,
device: Optional[str] = None,
dtype: Optional[torch.dtype] = None,
cache: Optional[KVCache] = None,
@ -67,7 +66,7 @@ class InferenceScheduler:
self._executor = Executor(
model=model,
tokenizer=tokenizer,
page_cache=self._cache,
kv_cache=self._cache,
device=self.device,
dtype=self.dtype,
)
@ -85,22 +84,6 @@ class InferenceScheduler:
def get_stats(self) -> Dict[str, Any]:
return self._task_mgr.get_stats()
@staticmethod
def _cached(cache: KVCache, task_id: str) -> int:
fn = getattr(cache, "task_cached", None)
return fn(task_id) if fn else 0
@staticmethod
def _record_hashes(
cache: KVCache,
task_id: str,
prompt_ids: List[int],
start_logical_page: int = 0,
):
fn = getattr(cache, "task_record_hashes", None)
if fn:
fn(task_id, prompt_ids, start_logical_page=start_logical_page)
def _run_generation_loop(self):
stop_ids = self._task_mgr.tokenizer.stop_ids
cache = self._cache
@ -131,7 +114,7 @@ class InferenceScheduler:
t
for t in self._task_mgr.get_active_tasks()
if t.output_tokens == 0
and self._cached(cache, t.task_id) < len(t.prompt_ids)
and cache.task_cached(t.task_id) < len(t.prompt_ids)
]
if to_prefill:
for t in to_prefill:
@ -141,7 +124,7 @@ class InferenceScheduler:
for t in to_prefill:
key = (
len(t.prompt_ids),
self._cached(cache, t.task_id),
cache.task_cached(t.task_id),
)
groups.setdefault(key, []).append(t)
@ -151,8 +134,8 @@ class InferenceScheduler:
cache, "page_size", 64
)
for t in group:
self._record_hashes(
cache, t.task_id, t.prompt_ids, start_logical_page
cache.task_record_hashes(
t.task_id, t.prompt_ids, start_logical_page
)
pos_groups: Dict[int, List[Task]] = {}
@ -168,8 +151,7 @@ class InferenceScheduler:
valid.append(t)
else:
t.status = TaskStatus.ABORTED
if t.stream_callback:
t.stream_callback(STOP)
self._task_mgr.invoke_callback(t.task_id, STOP)
if valid:
next_tokens = self._executor.execute_decode(valid)
@ -177,26 +159,23 @@ class InferenceScheduler:
for t, ntok in zip(valid, next_tokens):
t.output_ids.append(ntok)
t.output_tokens += 1
if t.stream_callback:
t.stream_callback(
self._task_mgr.tokenizer.decode([ntok])
)
self._task_mgr.invoke_callback(
t.task_id,
self._task_mgr.tokenizer.decode([ntok]),
)
for t in valid:
if t.is_finished(stop_ids):
if t.stream_callback:
t.stream_callback(STOP)
self._task_mgr.invoke_callback(t.task_id, STOP)
except Exception as e:
self._stop_event.set()
logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
for task in self._task_mgr.get_active_tasks():
if task.stream_callback:
task.stream_callback(STOP)
self._task_mgr.invoke_callback(task.task_id, STOP)
cache.task_free(task.task_id)
for task in self._task_mgr.get_waiting_tasks():
if task.stream_callback:
task.stream_callback(STOP)
self._task_mgr.invoke_callback(task.task_id, STOP)
self._task_mgr.clear_queues()
def start(self):
@ -214,12 +193,10 @@ class InferenceScheduler:
self._loop_thread.join(timeout=2.0)
self._loop_thread = None
for task in self._task_mgr.get_active_tasks():
if task.stream_callback:
task.stream_callback(STOP)
self._task_mgr.invoke_callback(task.task_id, STOP)
self._cache.task_free(task.task_id)
for task in self._task_mgr.get_waiting_tasks():
if task.stream_callback:
task.stream_callback(STOP)
self._task_mgr.invoke_callback(task.task_id, STOP)
self._task_mgr.clear_queues()
if torch.cuda.is_available():
torch.cuda.empty_cache()

13
astrai/inference/core/task.py
View File

@ -33,7 +33,6 @@ class Task:
temperature: float = 1.0,
top_p: float = 1.0,
top_k: int = 50,
stream_callback: Optional[Callable[[str], None]] = None,
):
self.task_id = task_id
self.prompt_ids = prompt_ids
@ -48,7 +47,6 @@ class Task:
self.output_tokens: int = 0
self.arrival_time = time.time()
self.finish_time: Optional[float] = None
self.stream_callback = stream_callback
@property
def next_pos(self) -> int:
@ -79,6 +77,7 @@ class TaskManager:
self.waiting_queue: Deque[Task] = deque()
self.active_tasks: List[Task] = []
self._callbacks: Dict[str, Callable[[str], None]] = {}
self._task_event = threading.Event()
self._lock = threading.Lock()
@ -117,12 +116,13 @@ class TaskManager:
temperature=temperature,
top_p=top_p,
top_k=top_k,
stream_callback=stream_callback,
)
with self._lock:
self.waiting_queue.append(task)
self._total_tasks += 1
if stream_callback:
self._callbacks[task_id] = stream_callback
self._task_event.set()
return task_id
@ -134,8 +134,14 @@ class TaskManager:
t for t in self.waiting_queue if t.task_id != task_id
)
self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
self._callbacks.pop(task_id, None)
return removed_active
def invoke_callback(self, task_id: str, token: str):
cb = self._callbacks.get(task_id)
if cb:
cb(token)
def get_stats(self) -> Dict[str, Any]:
return {
"total_tasks": self._total_tasks,
@ -204,6 +210,7 @@ class TaskManager:
with self._lock:
self.waiting_queue.clear()
self.active_tasks.clear()
self._callbacks.clear()
def wake(self):
self._task_event.set()

4
astrai/inference/engine.py
View File

@ -8,6 +8,7 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
import torch
import torch.nn as nn
from astrai.inference.core.cache import KVCache
from astrai.inference.core.scheduler import InferenceScheduler
from astrai.inference.core.task import STOP
from astrai.tokenize import AutoTokenizer
@ -101,6 +102,7 @@ class InferenceEngine:
max_seq_len: Optional[int] = None,
max_prompt_len: int = 2048,
page_size: int = 128,
cache: Optional[KVCache] = None,
):
self.model = model
self.tokenizer = tokenizer
@ -110,7 +112,7 @@ class InferenceEngine:
max_batch_size=max_batch_size,
max_seq_len=max_seq_len,
max_prompt_len=max_prompt_len,
page_size=page_size,
cache=cache,
)
self.scheduler.start()