diff --git a/astrai/inference/__init__.py b/astrai/inference/__init__.py index 15e9fb2..e3b4b9d 100644 --- a/astrai/inference/__init__.py +++ b/astrai/inference/__init__.py @@ -30,10 +30,14 @@ from astrai.inference.api.openai import OpenAIResponseBuilder from astrai.inference.core import ( STOP, Allocator, + CacheView, + ContiguousCache, + ContiguousCacheView, Executor, InferenceScheduler, KVCache, - KvcacheView, + PageCache, + PageCacheView, PagePool, PrefixCache, Storage, @@ -63,8 +67,12 @@ __all__ = [ "TaskManager", "TaskStatus", "Allocator", + "CacheView", "KVCache", - "KvcacheView", + "ContiguousCache", + "ContiguousCacheView", + "PageCache", + "PageCacheView", "PagePool", "PrefixCache", "Storage", diff --git a/astrai/inference/core/__init__.py b/astrai/inference/core/__init__.py index 183af3c..a545e9b 100644 --- a/astrai/inference/core/__init__.py +++ b/astrai/inference/core/__init__.py @@ -2,8 +2,12 @@ from astrai.inference.core.cache import ( Allocator, + CacheView, + ContiguousCache, + ContiguousCacheView, KVCache, - KvcacheView, + PageCache, + PageCacheView, PagePool, PrefixCache, Storage, @@ -16,8 +20,12 @@ from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus __all__ = [ "Allocator", + "CacheView", "KVCache", - "KvcacheView", + "ContiguousCache", + "ContiguousCacheView", + "PageCache", + "PageCacheView", "PagePool", "PrefixCache", "Storage", diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index 921fb91..663978b 100644 --- a/astrai/inference/core/cache.py +++ b/astrai/inference/core/cache.py @@ -1,4 +1,5 @@ import threading +from abc import ABC, abstractmethod from collections import OrderedDict from typing import Callable, Dict, List, Optional, Tuple @@ -275,7 +276,35 @@ class Storage: return k, v -class KvcacheView: +class CacheView(ABC): + """Abstract view passed to attention layers for KV-cache I/O.""" + + @abstractmethod + def write(self, layer_id: int, k: Tensor, v: Tensor): ... + + @abstractmethod + def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ... + + +class KVCache(ABC): + """Abstract KV-cache facade for scheduler/executor.""" + + @abstractmethod + def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ... + + @abstractmethod + def task_free(self, task_id: str): ... + + @abstractmethod + def task_extend(self, task_id: str, pos: int) -> bool: ... + + @abstractmethod + def bind_tasks( + self, task_ids: List[str], total_len: int, device: torch.device + ) -> CacheView: ... + + +class PageCacheView(CacheView): """Bundles Storage + page_table + total_len for attention layers.""" def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0): @@ -291,8 +320,8 @@ class KvcacheView: return self._storage.gather(layer_id, self._page_table, self._total_len) -class KVCache: - """Facade: page management + KV-cache I/O for continuous batching.""" +class PageCache(KVCache): + """Paged KV-cache with prefix sharing.""" def __init__( self, @@ -362,8 +391,102 @@ class KVCache: for i in range(start_logical_page, full_pages): self._pool.record(page_table[i], prompt_ids, i) - def make_table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor: - return self._table.table_tensor(task_ids, device) + def bind_tasks( + self, task_ids: List[str], total_len: int, device: torch.device + ) -> PageCacheView: + page_table = self._table.table_tensor(task_ids, device) + return PageCacheView(self._storage, page_table, total_len) - def bind(self, page_table: Tensor, total_len: int = 0) -> KvcacheView: - return KvcacheView(self._storage, page_table, total_len) + +class ContiguousCacheView(CacheView): + """Contiguous KV-cache view for attention layers.""" + + def __init__( + self, cache: "ContiguousCache", batch_indices: Tensor, total_len: int = 0 + ): + self._cache = cache + self._batch_indices = batch_indices + self._total_len = total_len + + def write(self, layer_id: int, k: Tensor, v: Tensor): + seq_len = k.size(1) + start_pos = self._total_len - seq_len + indices = self._batch_indices + self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k + self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v + new_len = start_pos + seq_len + for s in indices.tolist(): + cur = self._cache._slot_len.get(s, 0) + if new_len > cur: + self._cache._slot_len[s] = new_len + + def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: + max_len = max( + self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist() + ) + indices = self._batch_indices + k = self._cache.k[layer_id, indices, :max_len] + v = self._cache.v[layer_id, indices, :max_len] + return k, v + + +class ContiguousCache(KVCache): + """Contiguous per-slot KV cache (default implementation).""" + + def __init__( + self, + n_layers: int, + max_batch_size: int, + max_seq_len: int, + n_kv_heads: int, + head_dim: int, + device: torch.device, + dtype: torch.dtype, + ): + self.max_seq_len = max_seq_len + self.k = torch.zeros( + n_layers, + max_batch_size, + max_seq_len, + n_kv_heads, + head_dim, + device=device, + dtype=dtype, + ) + self.v = torch.zeros( + n_layers, + max_batch_size, + max_seq_len, + n_kv_heads, + head_dim, + device=device, + dtype=dtype, + ) + self._slot_len: Dict[int, int] = {} + self._task_slot: Dict[str, int] = {} + self._free_slots = list(range(max_batch_size)) + self._device = device + + def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: + if not self._free_slots: + return False + slot = self._free_slots.pop(0) + self._task_slot[task_id] = slot + self._slot_len[slot] = 0 + return True + + def task_free(self, task_id: str): + slot = self._task_slot.pop(task_id, None) + if slot is not None: + self._slot_len.pop(slot, None) + self._free_slots.append(slot) + + def task_extend(self, task_id: str, pos: int) -> bool: + return pos < self.max_seq_len + + def bind_tasks( + self, task_ids: List[str], total_len: int, device: torch.device + ) -> ContiguousCacheView: + slots = [self._task_slot[tid] for tid in task_ids] + batch_indices = torch.tensor(slots, dtype=torch.long, device=device) + return ContiguousCacheView(self, batch_indices, total_len) diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index 9157cfe..b963f18 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -43,7 +43,6 @@ class Executor: ) task_ids = [t.task_id for t in tasks] - page_tables = self.page_cache.make_table_tensor(task_ids, self.device) with torch.inference_mode(): self.model( @@ -53,7 +52,9 @@ class Executor: ) .unsqueeze(0) .expand(batch_sz, -1), - paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len), + paged_cache=self.page_cache.bind_tasks( + task_ids, prompt_len, self.device + ), ) def execute_decode(self, tasks: List[Task]) -> List[int]: @@ -72,7 +73,6 @@ class Executor: total_len = position_ids.max().item() + 1 task_ids = [t.task_id for t in tasks] - page_tables = self.page_cache.make_table_tensor(task_ids, self.device) temperatures = torch.tensor([t.temperature for t in tasks], device=self.device) top_ks = torch.tensor([t.top_k for t in tasks], device=self.device) @@ -81,7 +81,9 @@ class Executor: with torch.inference_mode(): outputs = self.model( input_ids.unsqueeze(1), - paged_cache=self.page_cache.bind(page_tables, total_len=total_len), + paged_cache=self.page_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 c06f3cb..ff4e215 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple import torch -from astrai.inference.core.cache import KVCache +from astrai.inference.core.cache import ContiguousCache, KVCache from astrai.inference.core.executor import Executor from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus from astrai.model.automodel import AutoModel @@ -14,7 +14,7 @@ logger = logging.getLogger(__name__) class InferenceScheduler: - """Four-phase continuous batching loop: cleanup -> refill -> prefill -> decode.""" + """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups).""" def __init__( self, @@ -26,6 +26,7 @@ class InferenceScheduler: page_size: int = 64, device: Optional[str] = None, dtype: Optional[torch.dtype] = None, + cache: Optional[KVCache] = None, ): config = model.config @@ -41,19 +42,20 @@ class InferenceScheduler: self.device = device or next(model.parameters()).device self.dtype = dtype or next(model.parameters()).dtype - n_pages = ( - max_batch_size * (self.max_seq_len + page_size) + page_size - 1 - ) // page_size + head_dim = config.dim // config.n_heads - self._page_cache = KVCache( - config.n_layers, - n_pages, - page_size, - config.n_kv_heads, - config.dim // config.n_heads, - self.device, - self.dtype, - ) + if cache is not None: + self._cache = cache + else: + self._cache = ContiguousCache( + config.n_layers, + max_batch_size, + self.max_seq_len, + config.n_kv_heads, + head_dim, + self.device, + self.dtype, + ) self._task_mgr = TaskManager( tokenizer=tokenizer, @@ -65,7 +67,7 @@ class InferenceScheduler: self._executor = Executor( model=model, tokenizer=tokenizer, - page_cache=self._page_cache, + page_cache=self._cache, device=self.device, dtype=self.dtype, ) @@ -78,18 +80,35 @@ class InferenceScheduler: def remove_task(self, task_id: str): for task in self._task_mgr.remove_task(task_id): - self._page_cache.task_free(task.task_id) + self._cache.task_free(task.task_id) 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 try: while not self._stop_event.is_set(): finished = self._task_mgr.remove_finished_tasks(stop_ids) for task in finished: - self._page_cache.task_free(task.task_id) + cache.task_free(task.task_id) active = self._task_mgr.get_active_tasks() available = self._task_mgr.max_batch_size - len(active) @@ -97,7 +116,7 @@ class InferenceScheduler: candidates = self._task_mgr.pull_candidates(available) failed = [] for task in candidates: - if self._page_cache.task_alloc(task.task_id, task.prompt_ids): + if cache.task_alloc(task.task_id, task.prompt_ids): self._task_mgr.activate(task) else: failed.append(task) @@ -112,7 +131,7 @@ class InferenceScheduler: t for t in self._task_mgr.get_active_tasks() if t.output_tokens == 0 - and self._page_cache.task_cached(t.task_id) < len(t.prompt_ids) + and self._cached(cache, t.task_id) < len(t.prompt_ids) ] if to_prefill: for t in to_prefill: @@ -122,31 +141,30 @@ class InferenceScheduler: for t in to_prefill: key = ( len(t.prompt_ids), - self._page_cache.task_cached(t.task_id), + self._cached(cache, t.task_id), ) groups.setdefault(key, []).append(t) for (prompt_len, start_pos), group in groups.items(): self._executor.execute_prefill(group, prompt_len, start_pos) - start_logical_page = start_pos // self._page_cache.page_size + start_logical_page = start_pos // getattr( + cache, "page_size", 64 + ) for t in group: - self._page_cache.task_record_hashes( - t.task_id, - t.prompt_ids, - start_logical_page=start_logical_page, + self._record_hashes( + cache, t.task_id, t.prompt_ids, start_logical_page ) pos_groups: Dict[int, List[Task]] = {} for t in self._task_mgr.get_active_tasks(): pos_groups.setdefault(t.next_pos, []).append(t) - if pos_groups: - best_key = max(pos_groups, key=lambda k: len(pos_groups[k])) - group = sorted(pos_groups[best_key], key=lambda t: t.task_id) + for next_pos in sorted(pos_groups.keys()): + group = sorted(pos_groups[next_pos], key=lambda t: t.task_id) valid: List[Task] = [] for t in group: - if self._page_cache.task_extend(t.task_id, t.next_pos): + if cache.task_extend(t.task_id, t.next_pos): valid.append(t) else: t.status = TaskStatus.ABORTED @@ -159,16 +177,10 @@ class InferenceScheduler: for t, ntok in zip(valid, next_tokens): t.output_ids.append(ntok) t.output_tokens += 1 - pos = t.input_tokens + t.output_tokens - extend_ok = self._page_cache.task_extend(t.task_id, pos) if t.stream_callback: t.stream_callback( self._task_mgr.tokenizer.decode([ntok]) ) - if not extend_ok: - t.status = TaskStatus.ABORTED - if t.stream_callback: - t.stream_callback(STOP) for t in valid: if t.is_finished(stop_ids): @@ -181,7 +193,7 @@ class InferenceScheduler: for task in self._task_mgr.get_active_tasks(): if task.stream_callback: task.stream_callback(STOP) - self._page_cache.task_free(task.task_id) + cache.task_free(task.task_id) for task in self._task_mgr.get_waiting_tasks(): if task.stream_callback: task.stream_callback(STOP) @@ -204,7 +216,7 @@ class InferenceScheduler: for task in self._task_mgr.get_active_tasks(): if task.stream_callback: task.stream_callback(STOP) - self._page_cache.task_free(task.task_id) + self._cache.task_free(task.task_id) for task in self._task_mgr.get_waiting_tasks(): if task.stream_callback: task.stream_callback(STOP) diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py index 3a64b4c..e94977a 100644 --- a/astrai/model/components/attention.py +++ b/astrai/model/components/attention.py @@ -6,7 +6,7 @@ import torch.nn.functional as F from torch import Tensor from astrai.factory import BaseFactory -from astrai.inference.core.cache import KvcacheView +from astrai.inference.core.cache import CacheView from astrai.model.components.linear import Linear from astrai.model.components.norm import RMSNorm from astrai.model.components.rope import apply_rotary_emb @@ -75,7 +75,7 @@ class GQA(nn.Module): x: Tensor, rotary_emb: Tensor, attn_mask: Tensor = None, - paged_cache: Optional[KvcacheView] = None, + paged_cache: Optional[CacheView] = None, ) -> Tensor: is_causal = attn_mask is None @@ -162,7 +162,7 @@ class MLA(nn.Module): x: Tensor, rotary_emb: Tensor, attn_mask: Tensor = None, - paged_cache: Optional[KvcacheView] = None, + paged_cache: Optional[CacheView] = None, ) -> Tensor: bsz, seq_len, _ = x.size() is_causal = attn_mask is None diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py index 18d71a3..1c79d2a 100644 --- a/astrai/model/components/decoder_block.py +++ b/astrai/model/components/decoder_block.py @@ -4,7 +4,7 @@ from typing import Optional import torch.nn as nn from torch import Tensor -from astrai.inference.core.cache import KvcacheView +from astrai.inference.core.cache import CacheView from astrai.model.components.attention import AttnFactory from astrai.model.components.mlp import FFNFactory from astrai.model.components.norm import RMSNorm @@ -25,7 +25,7 @@ class DecoderBlock(nn.Module): x: Tensor, rotary_emb: Tensor, attention_mask: Optional[Tensor] = None, - paged_cache: Optional[KvcacheView] = None, + paged_cache: Optional[CacheView] = None, ) -> Tensor: attn_output = self.attention( self.input_norm(x), diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index 9810545..123d1ef 100644 --- a/astrai/model/transformer.py +++ b/astrai/model/transformer.py @@ -5,7 +5,7 @@ import torch.nn as nn from torch import Tensor from astrai.config.model_config import AutoRegressiveLMConfig -from astrai.inference.core.cache import KvcacheView +from astrai.inference.core.cache import CacheView from astrai.model.automodel import AutoModel from astrai.model.components.decoder_block import DecoderBlock from astrai.model.components.embedding import Embedding @@ -112,7 +112,7 @@ class AutoRegressiveLM(AutoModel): self, input_ids: Tensor, input_mask: Optional[Tensor] = None, - paged_cache: Optional[KvcacheView] = None, + paged_cache: Optional[CacheView] = None, position_ids: Optional[Tensor] = None, ) -> Dict[str, Tensor]: assert input_ids.ndim == 2 diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index a42f09a..e6f5dbb 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -4,7 +4,7 @@ import torch from astrai.inference import ( Allocator, - KVCache, + PageCache, PagePool, PrefixCache, Storage, @@ -161,7 +161,7 @@ def test_task_table_pop(): def test_kv_cache_task_extend_allocates(): - cache = KVCache( + cache = PageCache( n_layers=1, n_pages=8, page_size=64, @@ -177,7 +177,7 @@ def test_kv_cache_task_extend_allocates(): def test_kv_cache_task_extend_fails_when_pool_full(): - cache = KVCache( + cache = PageCache( n_layers=1, n_pages=2, page_size=64,