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/

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 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)

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 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

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 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)