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:
@@ -4,7 +4,7 @@ from typing import Optional
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.inference.core.cache import CacheView
|
||||
from astrai.inference.core.cache import KVCache
|
||||
from astrai.model.components.attention import AttnFactory
|
||||
from astrai.model.components.mlp import FFNFactory
|
||||
from astrai.model.components.norm import RMSNorm
|
||||
@@ -33,14 +33,14 @@ class DecoderBlock(nn.Module):
|
||||
x: Tensor,
|
||||
rotary_emb: Tensor,
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
paged_cache: Optional[CacheView] = None,
|
||||
kv_cache: Optional[KVCache] = None,
|
||||
is_causal: bool = False,
|
||||
) -> Tensor:
|
||||
attn_output = self.attention(
|
||||
self.input_norm(x),
|
||||
rotary_emb,
|
||||
attention_mask,
|
||||
paged_cache,
|
||||
kv_cache,
|
||||
is_causal,
|
||||
)
|
||||
x = attn_output + x
|
||||
|
||||
Reference in New Issue
Block a user