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
+3 -3
View File
@@ -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 CacheView
from astrai.inference.core.cache import KVCache
from astrai.model.automodel import AutoModel, ModelFactory
from astrai.model.components.decoder_block import DecoderBlock
from astrai.model.components.embedding import Embedding
@@ -103,7 +103,7 @@ class AutoRegressiveLM(AutoModel):
self,
input_ids: Tensor,
input_mask: Optional[Tensor] = None,
paged_cache: Optional[CacheView] = None,
kv_cache: Optional[KVCache] = None,
position_ids: Optional[Tensor] = None,
) -> Dict[str, Tensor]:
assert input_ids.ndim == 2
@@ -114,7 +114,7 @@ class AutoRegressiveLM(AutoModel):
use_sdpa_causal_mask = attn_mask is None
for layer in self.layers:
x = layer(x, rotary_emb, attn_mask, paged_cache, use_sdpa_causal_mask)
x = layer(x, rotary_emb, attn_mask, kv_cache, use_sdpa_causal_mask)
hidden_states = self.norm(x)
logits = self.lm_head(hidden_states)