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 self, task_ids: List[str], total_len: int, device: torch.device
) -> CacheView: ... ) -> 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): class PageCacheView(CacheView):
"""Bundles Storage + page_table + total_len for attention layers.""" """Bundles Storage + page_table + total_len for attention layers."""

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

@ -19,13 +19,13 @@ class Executor:
self, self,
model: AutoModel, model: AutoModel,
tokenizer: AutoTokenizer, tokenizer: AutoTokenizer,
page_cache: KVCache, kv_cache: KVCache,
device: Optional[str] = None, device: Optional[str] = None,
dtype: Optional[torch.dtype] = None, dtype: Optional[torch.dtype] = None,
): ):
self.model = model self.model = model
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.page_cache = page_cache self.kv_cache = kv_cache
self.device = device or next(model.parameters()).device self.device = device or next(model.parameters()).device
self.dtype = dtype or next(model.parameters()).dtype self.dtype = dtype or next(model.parameters()).dtype
@ -52,9 +52,7 @@ class Executor:
) )
.unsqueeze(0) .unsqueeze(0)
.expand(batch_sz, -1), .expand(batch_sz, -1),
paged_cache=self.page_cache.bind_tasks( paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
task_ids, prompt_len, self.device
),
) )
def execute_decode(self, tasks: List[Task]) -> List[int]: def execute_decode(self, tasks: List[Task]) -> List[int]:
@ -81,9 +79,7 @@ class Executor:
with torch.inference_mode(): with torch.inference_mode():
outputs = self.model( outputs = self.model(
input_ids.unsqueeze(1), input_ids.unsqueeze(1),
paged_cache=self.page_cache.bind_tasks( paged_cache=self.kv_cache.bind_tasks(task_ids, total_len, self.device),
task_ids, total_len, self.device
),
position_ids=position_ids.unsqueeze(1), position_ids=position_ids.unsqueeze(1),
) )
logits = outputs["logits"][:, -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_batch_size: int = 16,
max_seq_len: Optional[int] = None, max_seq_len: Optional[int] = None,
max_prompt_len: int = 2048, max_prompt_len: int = 2048,
page_size: int = 64,
device: Optional[str] = None, device: Optional[str] = None,
dtype: Optional[torch.dtype] = None, dtype: Optional[torch.dtype] = None,
cache: Optional[KVCache] = None, cache: Optional[KVCache] = None,
@ -67,7 +66,7 @@ class InferenceScheduler:
self._executor = Executor( self._executor = Executor(
model=model, model=model,
tokenizer=tokenizer, tokenizer=tokenizer,
page_cache=self._cache, kv_cache=self._cache,
device=self.device, device=self.device,
dtype=self.dtype, dtype=self.dtype,
) )
@ -85,22 +84,6 @@ class InferenceScheduler:
def get_stats(self) -> Dict[str, Any]: def get_stats(self) -> Dict[str, Any]:
return self._task_mgr.get_stats() 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): def _run_generation_loop(self):
stop_ids = self._task_mgr.tokenizer.stop_ids stop_ids = self._task_mgr.tokenizer.stop_ids
cache = self._cache cache = self._cache
@ -131,7 +114,7 @@ class InferenceScheduler:
t t
for t in self._task_mgr.get_active_tasks() for t in self._task_mgr.get_active_tasks()
if t.output_tokens == 0 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: if to_prefill:
for t in to_prefill: for t in to_prefill:
@ -141,7 +124,7 @@ class InferenceScheduler:
for t in to_prefill: for t in to_prefill:
key = ( key = (
len(t.prompt_ids), len(t.prompt_ids),
self._cached(cache, t.task_id), cache.task_cached(t.task_id),
) )
groups.setdefault(key, []).append(t) groups.setdefault(key, []).append(t)
@ -151,8 +134,8 @@ class InferenceScheduler:
cache, "page_size", 64 cache, "page_size", 64
) )
for t in group: for t in group:
self._record_hashes( cache.task_record_hashes(
cache, t.task_id, t.prompt_ids, start_logical_page t.task_id, t.prompt_ids, start_logical_page
) )
pos_groups: Dict[int, List[Task]] = {} pos_groups: Dict[int, List[Task]] = {}
@ -168,8 +151,7 @@ class InferenceScheduler:
valid.append(t) valid.append(t)
else: else:
t.status = TaskStatus.ABORTED t.status = TaskStatus.ABORTED
if t.stream_callback: self._task_mgr.invoke_callback(t.task_id, STOP)
t.stream_callback(STOP)
if valid: if valid:
next_tokens = self._executor.execute_decode(valid) next_tokens = self._executor.execute_decode(valid)
@ -177,26 +159,23 @@ class InferenceScheduler:
for t, ntok in zip(valid, next_tokens): for t, ntok in zip(valid, next_tokens):
t.output_ids.append(ntok) t.output_ids.append(ntok)
t.output_tokens += 1 t.output_tokens += 1
if t.stream_callback: self._task_mgr.invoke_callback(
t.stream_callback( t.task_id,
self._task_mgr.tokenizer.decode([ntok]) self._task_mgr.tokenizer.decode([ntok]),
) )
for t in valid: for t in valid:
if t.is_finished(stop_ids): if t.is_finished(stop_ids):
if t.stream_callback: self._task_mgr.invoke_callback(t.task_id, STOP)
t.stream_callback(STOP)
except Exception as e: except Exception as e:
self._stop_event.set() self._stop_event.set()
logger.error(f"Scheduler loop crashed: {e}", exc_info=True) logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
for task in self._task_mgr.get_active_tasks(): for task in self._task_mgr.get_active_tasks():
if task.stream_callback: self._task_mgr.invoke_callback(task.task_id, STOP)
task.stream_callback(STOP)
cache.task_free(task.task_id) cache.task_free(task.task_id)
for task in self._task_mgr.get_waiting_tasks(): for task in self._task_mgr.get_waiting_tasks():
if task.stream_callback: self._task_mgr.invoke_callback(task.task_id, STOP)
task.stream_callback(STOP)
self._task_mgr.clear_queues() self._task_mgr.clear_queues()
def start(self): def start(self):
@ -214,12 +193,10 @@ class InferenceScheduler:
self._loop_thread.join(timeout=2.0) self._loop_thread.join(timeout=2.0)
self._loop_thread = None self._loop_thread = None
for task in self._task_mgr.get_active_tasks(): for task in self._task_mgr.get_active_tasks():
if task.stream_callback: self._task_mgr.invoke_callback(task.task_id, STOP)
task.stream_callback(STOP)
self._cache.task_free(task.task_id) self._cache.task_free(task.task_id)
for task in self._task_mgr.get_waiting_tasks(): for task in self._task_mgr.get_waiting_tasks():
if task.stream_callback: self._task_mgr.invoke_callback(task.task_id, STOP)
task.stream_callback(STOP)
self._task_mgr.clear_queues() self._task_mgr.clear_queues()
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.cuda.empty_cache() torch.cuda.empty_cache()

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

@ -33,7 +33,6 @@ class Task:
temperature: float = 1.0, temperature: float = 1.0,
top_p: float = 1.0, top_p: float = 1.0,
top_k: int = 50, top_k: int = 50,
stream_callback: Optional[Callable[[str], None]] = None,
): ):
self.task_id = task_id self.task_id = task_id
self.prompt_ids = prompt_ids self.prompt_ids = prompt_ids
@ -48,7 +47,6 @@ class Task:
self.output_tokens: int = 0 self.output_tokens: int = 0
self.arrival_time = time.time() self.arrival_time = time.time()
self.finish_time: Optional[float] = None self.finish_time: Optional[float] = None
self.stream_callback = stream_callback
@property @property
def next_pos(self) -> int: def next_pos(self) -> int:
@ -79,6 +77,7 @@ class TaskManager:
self.waiting_queue: Deque[Task] = deque() self.waiting_queue: Deque[Task] = deque()
self.active_tasks: List[Task] = [] self.active_tasks: List[Task] = []
self._callbacks: Dict[str, Callable[[str], None]] = {}
self._task_event = threading.Event() self._task_event = threading.Event()
self._lock = threading.Lock() self._lock = threading.Lock()
@ -117,12 +116,13 @@ class TaskManager:
temperature=temperature, temperature=temperature,
top_p=top_p, top_p=top_p,
top_k=top_k, top_k=top_k,
stream_callback=stream_callback,
) )
with self._lock: with self._lock:
self.waiting_queue.append(task) self.waiting_queue.append(task)
self._total_tasks += 1 self._total_tasks += 1
if stream_callback:
self._callbacks[task_id] = stream_callback
self._task_event.set() self._task_event.set()
return task_id return task_id
@ -134,8 +134,14 @@ class TaskManager:
t for t in self.waiting_queue if t.task_id != task_id 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.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
self._callbacks.pop(task_id, None)
return removed_active 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]: def get_stats(self) -> Dict[str, Any]:
return { return {
"total_tasks": self._total_tasks, "total_tasks": self._total_tasks,
@ -204,6 +210,7 @@ class TaskManager:
with self._lock: with self._lock:
self.waiting_queue.clear() self.waiting_queue.clear()
self.active_tasks.clear() self.active_tasks.clear()
self._callbacks.clear()
def wake(self): def wake(self):
self._task_event.set() 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
import torch.nn as nn import torch.nn as nn
from astrai.inference.core.cache import KVCache
from astrai.inference.core.scheduler import InferenceScheduler from astrai.inference.core.scheduler import InferenceScheduler
from astrai.inference.core.task import STOP from astrai.inference.core.task import STOP
from astrai.tokenize import AutoTokenizer from astrai.tokenize import AutoTokenizer
@ -101,6 +102,7 @@ class InferenceEngine:
max_seq_len: Optional[int] = None, max_seq_len: Optional[int] = None,
max_prompt_len: int = 2048, max_prompt_len: int = 2048,
page_size: int = 128, page_size: int = 128,
cache: Optional[KVCache] = None,
): ):
self.model = model self.model = model
self.tokenizer = tokenizer self.tokenizer = tokenizer
@ -110,7 +112,7 @@ class InferenceEngine:
max_batch_size=max_batch_size, max_batch_size=max_batch_size,
max_seq_len=max_seq_len, max_seq_len=max_seq_len,
max_prompt_len=max_prompt_len, max_prompt_len=max_prompt_len,
page_size=page_size, cache=cache,
) )
self.scheduler.start() self.scheduler.start()