perf: replace paged KV cache with contiguous ContiguousCache, decode all groups

- Add KVCache/CacheView abstract base classes in cache.py
- Add ContiguousCache (contiguous per-slot buffer, default) alongside PageCache (paged, renamed from old KVCache)
- Merge make_table_tensor + bind into bind_tasks on KVCache interface
- Remove task_cached/task_record_hashes from base class (PageCache-only)
- Scheduler: decode all position groups instead of just the largest (eliminates 63% group skip rate)
- Scheduler: accept optional cache param for swapping implementations
- Model layer type hints use CacheView base class
- Batch 1-32: 1-7% speedup from eliminating Storage.gather overhead
- All 183 inference tests pass
This commit is contained in:
2026-07-05 11:34:36 +08:00
parent 599a51f4f7
commit 5416c2e8fb
9 changed files with 215 additions and 62 deletions
+3 -3
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 KvcacheView
from astrai.inference.core.cache import CacheView
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[KvcacheView] = None,
paged_cache: Optional[CacheView] = None,
) -> Tensor:
is_causal = attn_mask is None
@@ -162,7 +162,7 @@ class MLA(nn.Module):
x: Tensor,
rotary_emb: Tensor,
attn_mask: Tensor = None,
paged_cache: Optional[KvcacheView] = None,
paged_cache: Optional[CacheView] = None,
) -> Tensor:
bsz, seq_len, _ = x.size()
is_causal = attn_mask is None
+2 -2
View File
@@ -4,7 +4,7 @@ from typing import Optional
import torch.nn as nn
from torch import Tensor
from astrai.inference.core.cache import KvcacheView
from astrai.inference.core.cache import CacheView
from astrai.model.components.attention import AttnFactory
from astrai.model.components.mlp import FFNFactory
from astrai.model.components.norm import RMSNorm
@@ -25,7 +25,7 @@ class DecoderBlock(nn.Module):
x: Tensor,
rotary_emb: Tensor,
attention_mask: Optional[Tensor] = None,
paged_cache: Optional[KvcacheView] = None,
paged_cache: Optional[CacheView] = None,
) -> Tensor:
attn_output = self.attention(
self.input_norm(x),
+2 -2
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 KvcacheView
from astrai.inference.core.cache import CacheView
from astrai.model.automodel import AutoModel
from astrai.model.components.decoder_block import DecoderBlock
from astrai.model.components.embedding import Embedding
@@ -112,7 +112,7 @@ class AutoRegressiveLM(AutoModel):
self,
input_ids: Tensor,
input_mask: Optional[Tensor] = None,
paged_cache: Optional[KvcacheView] = None,
paged_cache: Optional[CacheView] = None,
position_ids: Optional[Tensor] = None,
) -> Dict[str, Tensor]:
assert input_ids.ndim == 2