refactor: rebuild KV cache with three-layer separation architecture

- Replace CacheView/ContiguousCache/PageCache with SGLang-inspired design: KVStorage (flat token-level NHD buffers [n_layers, size, H, D]), ReqToTokenPool (index table [req_idx, pos] -> token_slot), Allocator + PrefixCache (slot allocation with LRU and prefix sharing)
- Add KVCache as pure dataclass passed to model: k_buffer, v_buffer, req_to_token, req_pool_indices, seq_lens, out_cache_loc
- PagePool orchestrates all three layers, supports contiguous mode (pre-allocated per-request blocks, default) and paged mode (page_size=1 or >1 with dynamic allocation and prefix caching)
- Attention layers now do raw buffer indexing instead of opaque write/gather method calls on CacheView objects
- Update executor.bind_tasks signature: seq_lens list + start_pos
- Rename paged_cache -> kv_cache throughout model/ and inference/
This commit is contained in:
2026-07-30 17:19:06 +08:00
parent fc47319240
commit deb2d7e127
10 changed files with 644 additions and 607 deletions
+10 -10
View File
@@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional, Tuple
import torch
from astrai.inference.core.cache import ContiguousCache, KVCache
from astrai.inference.core.cache import PagePool
from astrai.inference.core.executor import Executor
from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
from astrai.model.automodel import AutoModel
@@ -25,7 +25,7 @@ class InferenceScheduler:
max_seq_len: Optional[int] = None,
device: Optional[str] = None,
dtype: Optional[torch.dtype] = None,
cache: Optional[KVCache] = None,
cache: Optional[PagePool] = None,
):
config = model.config
@@ -46,14 +46,14 @@ class InferenceScheduler:
if cache is not None:
self._cache = cache
else:
self._cache = ContiguousCache(
config.num_hidden_layers,
max_batch_size,
self.max_seq_len,
config.num_key_value_heads,
head_dim,
self.device,
self.dtype,
self._cache = PagePool(
n_layers=config.num_hidden_layers,
n_kv_heads=config.num_key_value_heads,
head_dim=head_dim,
max_batch_size=max_batch_size,
max_seq_len=self.max_seq_len,
device=self.device,
dtype=self.dtype,
)
self._task_mgr = TaskManager(