diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index 663978b..2777231 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -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.""" diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index b963f18..6d0bd55 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -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, :] diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index ff4e215..66e14fb 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -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() diff --git a/astrai/inference/core/task.py b/astrai/inference/core/task.py index 5fcf0a4..b58af57 100644 --- a/astrai/inference/core/task.py +++ b/astrai/inference/core/task.py @@ -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() diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index c3a0c00..bbcf71e 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -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()