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