ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: rebuild KV cache with three-layer separation architecture

Browse Source 
- 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:
ViperEkura
2026-07-30 17:19:06 +08:00
parent fc47319240
commit deb2d7e127
10 changed files with 644 additions and 607 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/model/components/attention.py
+29 -9
View File
@@ -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 CacheView
from astrai.inference.core.cache import KVCache
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[CacheView] = None,
kv_cache: Optional[KVCache] = None,
is_causal: bool = False,
) -> Tensor:
q = self._split_heads(self.q_proj(x), self.n_heads)
@@ -86,9 +86,19 @@ class GQA(nn.Module):
if self.use_qk_norm:
q, k = self.q_norm(q), self.k_norm(k)
if paged_cache is not None:
paged_cache.write(self.layer_id, k, v)
k, v = paged_cache.gather(self.layer_id)
if kv_cache is not None:
kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k
kv_cache.v_buffer[self.layer_id][kv_cache.out_cache_loc] = v
max_len = kv_cache.seq_lens.max()
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
pos_mask = (
torch.arange(max_len, device=x.device)[None, :]
< kv_cache.seq_lens[:, None]
)
indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
k = kv_cache.k_buffer[self.layer_id][indices]
v = kv_cache.v_buffer[self.layer_id][indices]
k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
@@ -161,7 +171,7 @@ class MLA(nn.Module):
x: Tensor,
rotary_emb: Tensor,
attn_mask: Tensor = None,
paged_cache: Optional[CacheView] = None,
kv_cache: Optional[KVCache] = None,
is_causal: bool = False,
) -> Tensor:
bsz, seq_len, _ = x.size()
@@ -193,9 +203,19 @@ class MLA(nn.Module):
q = self.q_norm(q)
k = self.k_norm(k)
if paged_cache is not None:
paged_cache.write(self.layer_id, k, v)
k, v = paged_cache.gather(self.layer_id)
if kv_cache is not None:
kv_cache.k_buffer[self.layer_id][kv_cache.out_cache_loc] = k
kv_cache.v_buffer[self.layer_id][kv_cache.out_cache_loc] = v
max_len = kv_cache.seq_lens.max()
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
pos_mask = (
torch.arange(max_len, device=x.device)[None, :]
< kv_cache.seq_lens[:, None]
)
indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
k = kv_cache.k_buffer[self.layer_id][indices]
v = kv_cache.v_buffer[self.layer_id][indices]
q = q.permute(0, 2, 1, 3)
k = k.permute(0, 2, 1, 3)