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

refactor: 重构 cache 和 inference 参数体系，分离存储与分配

Browse Source 
- 合并 GenerationRequest/GenerationParams，统一 max_tokens 参数名
- PagePool/PrefixCache 分离为 Allocator + PrefixCache + PagePool
- 拆分 KV 存储为独立 Storage 类，PagedCache → KVCache，CacheView → KvcacheView
- Allocator.inc_ref 移除 LRU 防止竞争，Storage.write 增加负页防御
- Allocator/PrefixCache/TaskTable 加 threading.Lock 保证线程安全
- server.py uvicorn.run 改为传 app 对象修复导入错误
- benchmark.py 适配 KVCache 新 API
This commit is contained in:
ViperEkura
2026-05-14 20:05:08 +08:00
parent 18fe6e9339
commit 205b40bd28
13 changed files with 394 additions and 351 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/inference/test_cache.py
+69 -92
View File
@@ -3,14 +3,20 @@
import torch
from astrai.inference import (
PagedCache,
Allocator,
KVCache,
PagePool,
PrefixCache,
Storage,
TaskTable,
page_hash,
)
def make_pool(n_pages: int, page_size: int) -> PagePool:
return PagePool(Allocator(n_pages), PrefixCache(page_size))
def test_page_hash_full_page():
token_ids = list(range(256))
h = page_hash(token_ids, 0, 64)
@@ -24,7 +30,7 @@ def test_page_hash_different_page_differs():
def test_page_pool_alloc_free_cycle():
pool = PagePool(n_pages=4)
pool = make_pool(4, 64)
a = pool.alloc()
b = pool.alloc()
assert a != b
@@ -35,130 +41,101 @@ def test_page_pool_alloc_free_cycle():
def test_page_pool_alloc_when_full():
pool = PagePool(n_pages=2)
pool = make_pool(2, 64)
pool.alloc()
pool.alloc()
assert pool.alloc() == -1
def test_page_pool_lru_eviction():
evicted = []
def on_evict(idx):
evicted.append(idx)
pool = PagePool(n_pages=2, on_evict=on_evict)
pool = make_pool(2, 64)
p0 = pool.alloc()
p1 = pool.alloc()
pool.free(p0, keep_cached=True)
pool.free(p1, keep_cached=True)
pool.record(p0, list(range(64)), 0)
pool.record(p1, list(range(64, 128)), 0)
pool.free(p0)
pool.free(p1)
pool.alloc()
assert len(evicted) == 1
assert evicted[0] == p0
assert p0 in pool._alloc._lru or p1 in pool._alloc._lru
def test_page_pool_inc_ref_and_free():
pool = PagePool(n_pages=2)
pool = make_pool(2, 64)
p = pool.alloc()
pool.inc_ref(p)
assert pool._refs[p] == 2
assert pool._alloc._refs[p] == 2
pool.free(p)
assert pool._refs[p] == 1
assert pool._alloc._refs[p] == 1
pool.free(p)
assert pool._refs[p] == 0
def test_page_pool_touch_moves_to_end():
pool = PagePool(n_pages=4)
p0 = pool.alloc()
p1 = pool.alloc()
p2 = pool.alloc()
pool.free(p0, keep_cached=True)
pool.free(p1, keep_cached=True)
pool.free(p2, keep_cached=True)
assert next(iter(pool._lru)) == p0
pool.touch(p0)
assert next(reversed(pool._lru)) == p0
def test_page_pool_remove_from_lru():
pool = PagePool(n_pages=4)
p0 = pool.alloc()
pool.free(p0, keep_cached=True)
assert p0 in pool._lru
pool.remove_from_lru(p0)
assert p0 not in pool._lru
assert pool._alloc._refs[p] == 0
def test_page_pool_keep_cached_realloc():
"""Free mask has priority over LRU; cached page returned only when no free pages."""
pool = PagePool(n_pages=3)
pool = make_pool(3, 64)
p0 = pool.alloc()
p1 = pool.alloc()
p2 = pool.alloc()
pool.free(p0, keep_cached=True)
pool.free(p1, keep_cached=True)
pool.free(p2, keep_cached=True)
for p in (p0, p1, p2):
pool.record(p, [p] * 64, 0)
pool.free(p0)
pool.free(p1)
pool.free(p2)
assert pool.alloc() == p0
def _record_then_cache(pool, prefix, page, token_ids, logical_idx):
"""Simulate the real lifecycle: record → ref stays >0, then free cached returns to LRU."""
prefix.record(page, token_ids, logical_idx, pool)
pool.free(page, keep_cached=True)
def test_prefix_cache_lookup_returns_hits():
token_ids = list(range(256))
pool = PagePool(n_pages=16)
prefix = PrefixCache(page_size=64)
pool = make_pool(16, 64)
pages = [pool.alloc() for _ in range(4)]
for i, p in enumerate(pages):
_record_then_cache(pool, prefix, p, token_ids, i)
hits = prefix.lookup(token_ids, pool)
pool.record(p, token_ids, i)
pool.free(p)
hits = pool.lookup(token_ids)
assert hits == pages
def test_prefix_cache_lookup_stops_at_first_miss():
token_ids = list(range(256))
pool = PagePool(n_pages=16)
prefix = PrefixCache(page_size=64)
pool = make_pool(16, 64)
p0 = pool.alloc()
_record_then_cache(pool, prefix, p0, token_ids, 0)
pool.record(p0, token_ids, 0)
pool.free(p0)
p1 = pool.alloc()
_record_then_cache(pool, prefix, p1, [99] * 64, 1)
hits = prefix.lookup(token_ids, pool)
pool.record(p1, [99] * 64, 1)
pool.free(p1)
hits = pool.lookup(token_ids)
assert len(hits) == 1
assert hits[0] == p0
def test_prefix_cache_ignores_partial_last_page():
token_ids = list(range(100))
pool = PagePool(n_pages=16)
prefix = PrefixCache(page_size=64)
pool = make_pool(16, 64)
p = pool.alloc()
_record_then_cache(pool, prefix, p, token_ids, 0)
hits = prefix.lookup(token_ids, pool)
pool.record(p, token_ids, 0)
pool.free(p)
hits = pool.lookup(token_ids)
assert len(hits) == 1
def test_prefix_cache_on_evict_clears_mappings():
pool = PagePool(n_pages=4)
prefix = PrefixCache(page_size=64)
pool = make_pool(4, 64)
p = pool.alloc()
_record_then_cache(pool, prefix, p, list(range(64)), 0)
assert prefix.has_page(p)
prefix.on_evict(p)
assert not prefix.has_page(p)
pool.record(p, list(range(64)), 0)
pool.free(p)
assert p in pool._prefix._page_to_hash
pool._prefix.evict(p)
assert p not in pool._prefix._page_to_hash
def test_prefix_cache_has_page():
pool = PagePool(n_pages=4)
prefix = PrefixCache(page_size=64)
pool = make_pool(4, 64)
p = pool.alloc()
assert not prefix.has_page(p)
_record_then_cache(pool, prefix, p, list(range(64)), 0)
assert prefix.has_page(p)
assert p not in pool._prefix._page_to_hash
pool.record(p, list(range(64)), 0)
pool.free(p)
assert p in pool._prefix._page_to_hash
def test_task_table_set_get():
@@ -183,8 +160,8 @@ def test_task_table_pop():
assert table.get("task1") == []
def test_paged_cache_task_extend_allocates():
cache = PagedCache(
def test_kv_cache_task_extend_allocates():
cache = KVCache(
n_layers=1,
n_pages=8,
page_size=64,
@@ -199,8 +176,8 @@ def test_paged_cache_task_extend_allocates():
assert len(cache._table.get("task1")) == 4
def test_paged_cache_task_extend_fails_when_pool_full():
cache = PagedCache(
def test_kv_cache_task_extend_fails_when_pool_full():
cache = KVCache(
n_layers=1,
n_pages=2,
page_size=64,
@@ -230,8 +207,8 @@ def test_task_table_table_tensor_empty_input():
assert t.numel() == 0
def test_paged_cache_write_gather_single_page():
cache = PagedCache(
def test_storage_write_gather_single_page():
storage = Storage(
n_layers=2,
n_pages=8,
page_size=4,
@@ -244,13 +221,13 @@ def test_paged_cache_write_gather_single_page():
k = torch.randn(1, 2, 2, 8)
v = torch.randn(1, 2, 2, 8)
cache.write(0, page_table, 0, k, v)
gk, gv = cache.gather(0, page_table, 2)
storage.write(0, page_table, 0, k, v)
gk, gv = storage.gather(0, page_table, 2)
assert torch.allclose(gk, k)
def test_paged_cache_write_cross_page():
cache = PagedCache(
def test_storage_write_cross_page():
storage = Storage(
n_layers=1,
n_pages=8,
page_size=4,
@@ -263,13 +240,13 @@ def test_paged_cache_write_cross_page():
k = torch.randn(1, 8, 2, 8)
v = torch.randn(1, 8, 2, 8)
cache.write(0, page_table, 0, k, v)
gk, gv = cache.gather(0, page_table, 8)
storage.write(0, page_table, 0, k, v)
gk, gv = storage.gather(0, page_table, 8)
assert torch.allclose(gk, k)
def test_paged_cache_gather_truncates_to_total_len():
cache = PagedCache(
def test_storage_gather_truncates_to_total_len():
storage = Storage(
n_layers=1,
n_pages=8,
page_size=4,
@@ -281,14 +258,14 @@ def test_paged_cache_gather_truncates_to_total_len():
page_table = torch.tensor([[0, 1]], dtype=torch.long)
k = torch.randn(1, 6, 2, 8)
v = torch.randn(1, 6, 2, 8)
cache.write(0, page_table, 0, k, v)
storage.write(0, page_table, 0, k, v)
gk, gv = cache.gather(0, page_table, 5)
gk, gv = storage.gather(0, page_table, 5)
assert gk.shape == (1, 5, 2, 8)
def test_paged_cache_gather_clamps_negative_padding():
cache = PagedCache(
def test_storage_gather_clamps_negative_padding():
storage = Storage(
n_layers=1,
n_pages=8,
page_size=4,
@@ -298,5 +275,5 @@ def test_paged_cache_gather_clamps_negative_padding():
dtype=torch.float32,
)
page_table = torch.tensor([[0, -1]], dtype=torch.long)
gk, gv = cache.gather(0, page_table, 4)
gk, gv = storage.gather(0, page_table, 4)
assert gk.shape == (1, 4, 2, 8)