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/
This commit is contained in:
+270
-189
@@ -4,17 +4,14 @@ import torch
|
||||
|
||||
from astrai.inference import (
|
||||
Allocator,
|
||||
PageCache,
|
||||
KVStorage,
|
||||
PagePool,
|
||||
PrefixCache,
|
||||
Storage,
|
||||
TaskTable,
|
||||
ReqToTokenPool,
|
||||
page_hash,
|
||||
)
|
||||
|
||||
|
||||
def make_pool(n_pages: int, page_size: int) -> PagePool:
|
||||
return PagePool(Allocator(n_pages), PrefixCache(page_size))
|
||||
# ---- page_hash ----
|
||||
|
||||
|
||||
def test_page_hash_full_page():
|
||||
@@ -29,251 +26,335 @@ def test_page_hash_different_page_differs():
|
||||
assert page_hash(token_ids, 0, 64) != page_hash(token_ids, 1, 64)
|
||||
|
||||
|
||||
def test_page_pool_alloc_free_cycle():
|
||||
pool = make_pool(4, 64)
|
||||
a = pool.alloc()
|
||||
b = pool.alloc()
|
||||
# ---- Allocator ----
|
||||
|
||||
|
||||
def test_allocator_alloc_free_cycle():
|
||||
alloc = Allocator(4)
|
||||
a = alloc.alloc()
|
||||
b = alloc.alloc()
|
||||
assert a != b
|
||||
pool.free(a)
|
||||
pool.free(b)
|
||||
c = pool.alloc()
|
||||
alloc.free(a)
|
||||
alloc.free(b)
|
||||
c = alloc.alloc()
|
||||
assert c in (a, b)
|
||||
|
||||
|
||||
def test_page_pool_alloc_when_full():
|
||||
pool = make_pool(2, 64)
|
||||
pool.alloc()
|
||||
pool.alloc()
|
||||
assert pool.alloc() == -1
|
||||
def test_allocator_alloc_when_full():
|
||||
alloc = Allocator(2)
|
||||
alloc.alloc()
|
||||
alloc.alloc()
|
||||
assert alloc.alloc() == -1
|
||||
|
||||
|
||||
def test_page_pool_lru_eviction():
|
||||
pool = make_pool(2, 64)
|
||||
p0 = pool.alloc()
|
||||
p1 = pool.alloc()
|
||||
pool.record(p0, list(range(64)), 0)
|
||||
pool.record(p1, list(range(64, 128)), 0)
|
||||
pool.free(p0)
|
||||
pool.free(p1)
|
||||
pool.alloc()
|
||||
assert p0 in pool._alloc._lru or p1 in pool._alloc._lru
|
||||
def test_allocator_lru_eviction():
|
||||
alloc = Allocator(2)
|
||||
p0 = alloc.alloc()
|
||||
p1 = alloc.alloc()
|
||||
alloc.free(p0, keep_cached=True)
|
||||
alloc.free(p1, keep_cached=True)
|
||||
alloc.alloc()
|
||||
assert p0 in alloc._lru or p1 in alloc._lru
|
||||
|
||||
|
||||
def test_page_pool_inc_ref_and_free():
|
||||
pool = make_pool(2, 64)
|
||||
p = pool.alloc()
|
||||
pool.inc_ref(p)
|
||||
assert pool._alloc._refs[p] == 2
|
||||
pool.free(p)
|
||||
assert pool._alloc._refs[p] == 1
|
||||
pool.free(p)
|
||||
assert pool._alloc._refs[p] == 0
|
||||
def test_allocator_inc_ref_and_free():
|
||||
alloc = Allocator(2)
|
||||
p = alloc.alloc()
|
||||
alloc.inc_ref(p)
|
||||
assert alloc._refs[p] == 2
|
||||
alloc.free(p)
|
||||
assert alloc._refs[p] == 1
|
||||
alloc.free(p)
|
||||
assert 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 = make_pool(3, 64)
|
||||
p0 = pool.alloc()
|
||||
p1 = pool.alloc()
|
||||
p2 = pool.alloc()
|
||||
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
|
||||
# ---- PrefixCache ----
|
||||
|
||||
|
||||
def test_prefix_cache_lookup_returns_hits():
|
||||
token_ids = list(range(256))
|
||||
pool = make_pool(16, 64)
|
||||
pages = [pool.alloc() for _ in range(4)]
|
||||
prefix = PrefixCache(64)
|
||||
pages = [0, 1, 2, 3]
|
||||
for i, p in enumerate(pages):
|
||||
pool.record(p, token_ids, i)
|
||||
pool.free(p)
|
||||
hits = pool.lookup(token_ids)
|
||||
prefix.record(p, token_ids, i)
|
||||
hits = prefix.lookup(token_ids)
|
||||
assert hits == pages
|
||||
|
||||
|
||||
def test_prefix_cache_lookup_stops_at_first_miss():
|
||||
token_ids = list(range(256))
|
||||
pool = make_pool(16, 64)
|
||||
p0 = pool.alloc()
|
||||
pool.record(p0, token_ids, 0)
|
||||
pool.free(p0)
|
||||
p1 = pool.alloc()
|
||||
pool.record(p1, [99] * 64, 1)
|
||||
pool.free(p1)
|
||||
hits = pool.lookup(token_ids)
|
||||
prefix = PrefixCache(64)
|
||||
prefix.record(0, token_ids, 0)
|
||||
prefix.record(1, [99] * 64, 1)
|
||||
hits = prefix.lookup(token_ids)
|
||||
assert len(hits) == 1
|
||||
assert hits[0] == p0
|
||||
assert hits[0] == 0
|
||||
|
||||
|
||||
def test_prefix_cache_ignores_partial_last_page():
|
||||
token_ids = list(range(100))
|
||||
pool = make_pool(16, 64)
|
||||
p = pool.alloc()
|
||||
pool.record(p, token_ids, 0)
|
||||
pool.free(p)
|
||||
hits = pool.lookup(token_ids)
|
||||
prefix = PrefixCache(64)
|
||||
prefix.record(0, token_ids, 0)
|
||||
hits = prefix.lookup(token_ids)
|
||||
assert len(hits) == 1
|
||||
|
||||
|
||||
def test_prefix_cache_on_evict_clears_mappings():
|
||||
pool = make_pool(4, 64)
|
||||
p = pool.alloc()
|
||||
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
|
||||
prefix = PrefixCache(64)
|
||||
prefix.record(0, list(range(64)), 0)
|
||||
assert 0 in prefix._page_to_hash
|
||||
prefix.evict(0)
|
||||
assert 0 not in prefix._page_to_hash
|
||||
|
||||
|
||||
def test_prefix_cache_has_page():
|
||||
pool = make_pool(4, 64)
|
||||
p = pool.alloc()
|
||||
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
|
||||
prefix = PrefixCache(64)
|
||||
assert not prefix.has_page(0)
|
||||
prefix.record(0, list(range(64)), 0)
|
||||
assert prefix.has_page(0)
|
||||
|
||||
|
||||
def test_task_table_set_get():
|
||||
table = TaskTable(page_size=64)
|
||||
table.set("task1", [0, 1, 2], 128)
|
||||
assert table.get("task1") == [0, 1, 2]
|
||||
assert table.get_cached("task1") == 128
|
||||
# ---- ReqToTokenPool ----
|
||||
|
||||
|
||||
def test_task_table_get_missing():
|
||||
table = TaskTable(page_size=64)
|
||||
assert table.get("nonexistent") == []
|
||||
assert table.get_cached("nonexistent") == 0
|
||||
def test_req_to_token_pool_alloc_free():
|
||||
pool = ReqToTokenPool(4, 128, torch.device("cpu"))
|
||||
slots = pool.alloc(2)
|
||||
assert len(slots) == 2
|
||||
assert len(pool.free_slots) == 2
|
||||
pool.free(slots)
|
||||
assert len(pool.free_slots) == 4
|
||||
|
||||
|
||||
def test_task_table_pop():
|
||||
table = TaskTable(page_size=64)
|
||||
table.set("task1", [0, 1], 64)
|
||||
pages, cached = table.pop("task1")
|
||||
assert pages == [0, 1]
|
||||
assert cached == 64
|
||||
assert table.get("task1") == []
|
||||
def test_req_to_token_pool_alloc_when_full():
|
||||
pool = ReqToTokenPool(2, 128, torch.device("cpu"))
|
||||
pool.alloc(2)
|
||||
assert pool.alloc(1) is None
|
||||
|
||||
|
||||
def test_kv_cache_task_extend_allocates():
|
||||
cache = PageCache(
|
||||
n_layers=1,
|
||||
n_pages=8,
|
||||
page_size=64,
|
||||
n_kv_heads=2,
|
||||
head_dim=8,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
cache._table.set("task1", [], 0)
|
||||
ok = cache.task_extend("task1", 200)
|
||||
assert ok
|
||||
assert len(cache._table.get("task1")) == 4
|
||||
def test_req_to_token_pool_write():
|
||||
pool = ReqToTokenPool(4, 128, torch.device("cpu"))
|
||||
slots = pool.alloc(1)
|
||||
pool.write((slots[0], slice(0, 3)), torch.tensor([10, 20, 30]))
|
||||
assert pool.req_to_token[slots[0], 0].item() == 10
|
||||
assert pool.req_to_token[slots[0], 2].item() == 30
|
||||
|
||||
|
||||
def test_kv_cache_task_extend_fails_when_pool_full():
|
||||
cache = PageCache(
|
||||
n_layers=1,
|
||||
n_pages=2,
|
||||
page_size=64,
|
||||
n_kv_heads=2,
|
||||
head_dim=8,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
cache._table.set("task1", [0, 1], 0)
|
||||
ok = cache.task_extend("task1", 300)
|
||||
assert not ok
|
||||
# ---- KVStorage ----
|
||||
|
||||
|
||||
def test_task_table_table_tensor():
|
||||
table = TaskTable(page_size=64)
|
||||
table.set("a", [0, 1], 0)
|
||||
table.set("b", [2, 3, 4], 0)
|
||||
t = table.table_tensor(["a", "b"], torch.device("cpu"))
|
||||
assert t.shape == (2, 3)
|
||||
assert t[0].tolist() == [0, 1, -1]
|
||||
assert t[1].tolist() == [2, 3, 4]
|
||||
|
||||
|
||||
def test_task_table_table_tensor_empty_input():
|
||||
table = TaskTable(page_size=64)
|
||||
t = table.table_tensor([], torch.device("cpu"))
|
||||
assert t.numel() == 0
|
||||
|
||||
|
||||
def test_storage_write_gather_single_page():
|
||||
storage = Storage(
|
||||
def test_kv_storage_set_and_get():
|
||||
storage = KVStorage(
|
||||
size=16,
|
||||
n_layers=2,
|
||||
n_pages=8,
|
||||
page_size=4,
|
||||
n_kv_heads=2,
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
page_table = torch.tensor([[0]], dtype=torch.long)
|
||||
k = torch.randn(1, 2, 2, 8)
|
||||
v = torch.randn(1, 2, 2, 8)
|
||||
|
||||
storage.write(0, page_table, 0, k, v)
|
||||
gk, gv = storage.gather(0, page_table, 2)
|
||||
assert torch.allclose(gk, k)
|
||||
loc = torch.tensor([[0, 1]], dtype=torch.long)
|
||||
k = torch.randn(1, 2, 4, 8)
|
||||
v = torch.randn(1, 2, 4, 8)
|
||||
storage.set_kv_buffer(0, loc, k, v)
|
||||
assert torch.allclose(storage.get_key_buffer(0)[loc], k)
|
||||
assert torch.allclose(storage.get_value_buffer(0)[loc], v)
|
||||
|
||||
|
||||
def test_storage_write_cross_page():
|
||||
storage = Storage(
|
||||
def test_kv_storage_buffer_shape():
|
||||
storage = KVStorage(
|
||||
size=32,
|
||||
n_layers=3,
|
||||
n_kv_heads=8,
|
||||
head_dim=16,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
assert storage.k_buffer.shape == (3, 32, 8, 16)
|
||||
assert storage.v_buffer.shape == (3, 32, 8, 16)
|
||||
|
||||
|
||||
# ---- PagePool (contiguous mode) ----
|
||||
|
||||
|
||||
def _make_contiguous_pool(**kwargs):
|
||||
defaults = dict(
|
||||
n_layers=2,
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
max_batch_size=4,
|
||||
max_seq_len=64,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return PagePool(**defaults)
|
||||
|
||||
|
||||
def test_page_pool_contiguous_task_alloc_free():
|
||||
pool = _make_contiguous_pool()
|
||||
assert pool.task_alloc("t1", [1, 2, 3])
|
||||
assert "t1" in pool._task_req
|
||||
pool.task_free("t1")
|
||||
assert "t1" not in pool._task_req
|
||||
|
||||
|
||||
def test_page_pool_contiguous_task_extend():
|
||||
pool = _make_contiguous_pool()
|
||||
pool.task_alloc("t1", [1, 2, 3])
|
||||
assert pool.task_extend("t1", 3)
|
||||
assert pool.task_extend("t1", 63)
|
||||
assert not pool.task_extend("t1", 64)
|
||||
|
||||
|
||||
def test_page_pool_contiguous_task_cached():
|
||||
pool = _make_contiguous_pool()
|
||||
pool.task_alloc("t1", [1, 2, 3])
|
||||
assert pool.task_cached("t1") == 0
|
||||
|
||||
|
||||
def test_page_pool_contiguous_bind_tasks_prefill():
|
||||
pool = _make_contiguous_pool()
|
||||
pool.task_alloc("t1", list(range(10)))
|
||||
pool.task_alloc("t2", list(range(10)))
|
||||
kv = pool.bind_tasks(["t1", "t2"], [10, 10], torch.device("cpu"), start_pos=0)
|
||||
assert kv.out_cache_loc.shape == (2, 10)
|
||||
assert kv.seq_lens.tolist() == [10, 10]
|
||||
assert kv.req_pool_indices.shape == (2,)
|
||||
|
||||
|
||||
def test_page_pool_contiguous_bind_tasks_decode():
|
||||
pool = _make_contiguous_pool()
|
||||
pool.task_alloc("t1", list(range(10)))
|
||||
pool.task_alloc("t2", list(range(8)))
|
||||
kv = pool.bind_tasks(["t1", "t2"], [11, 9], torch.device("cpu"))
|
||||
assert kv.out_cache_loc.shape == (2, 1)
|
||||
assert kv.seq_lens.tolist() == [11, 9]
|
||||
|
||||
|
||||
def test_page_pool_contiguous_bind_roundtrip():
|
||||
"""Write KV via bind_tasks, then gather via req_to_token indexing."""
|
||||
pool = _make_contiguous_pool(n_layers=1, n_kv_heads=2, head_dim=4)
|
||||
pool.task_alloc("t1", list(range(4)))
|
||||
|
||||
kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0)
|
||||
k = torch.randn(1, 4, 2, 4)
|
||||
v = torch.randn(1, 4, 2, 4)
|
||||
kv.k_buffer[0][kv.out_cache_loc] = k
|
||||
kv.v_buffer[0][kv.out_cache_loc] = v
|
||||
|
||||
indices = kv.req_to_token[kv.req_pool_indices, :4]
|
||||
gathered_k = kv.k_buffer[0][indices]
|
||||
gathered_v = kv.v_buffer[0][indices]
|
||||
assert torch.allclose(gathered_k, k)
|
||||
assert torch.allclose(gathered_v, v)
|
||||
|
||||
|
||||
# ---- PagePool (paged mode, page_size=1) ----
|
||||
|
||||
|
||||
def _make_paged_pool(**kwargs):
|
||||
defaults = dict(
|
||||
n_layers=1,
|
||||
n_pages=8,
|
||||
page_size=4,
|
||||
n_kv_heads=2,
|
||||
head_dim=8,
|
||||
head_dim=4,
|
||||
max_batch_size=4,
|
||||
max_seq_len=64,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
page_size=1,
|
||||
n_tokens=128,
|
||||
)
|
||||
page_table = torch.tensor([[0, 1]], dtype=torch.long)
|
||||
k = torch.randn(1, 8, 2, 8)
|
||||
v = torch.randn(1, 8, 2, 8)
|
||||
|
||||
storage.write(0, page_table, 0, k, v)
|
||||
gk, gv = storage.gather(0, page_table, 8)
|
||||
assert torch.allclose(gk, k)
|
||||
defaults.update(kwargs)
|
||||
return PagePool(**defaults)
|
||||
|
||||
|
||||
def test_storage_gather_truncates_to_total_len():
|
||||
storage = Storage(
|
||||
def test_page_pool_paged_task_alloc():
|
||||
pool = _make_paged_pool()
|
||||
assert pool.task_alloc("t1", list(range(10)))
|
||||
req_idx = pool._task_req["t1"]
|
||||
slots = pool._task_slots["t1"]
|
||||
assert len(slots) == 10
|
||||
assert pool._req_pool.req_to_token[req_idx, 0].item() == slots[0]
|
||||
|
||||
|
||||
def test_page_pool_paged_task_extend():
|
||||
pool = _make_paged_pool()
|
||||
pool.task_alloc("t1", list(range(4)))
|
||||
assert pool.task_extend("t1", 4)
|
||||
req_idx = pool._task_req["t1"]
|
||||
slot = pool._req_pool.req_to_token[req_idx, 4].item()
|
||||
assert slot >= 0
|
||||
|
||||
|
||||
def test_page_pool_paged_task_free_releases_slots():
|
||||
pool = _make_paged_pool(n_tokens=16)
|
||||
pool.task_alloc("t1", list(range(8)))
|
||||
pool.task_free("t1")
|
||||
assert "t1" not in pool._task_req
|
||||
assert len(pool._req_pool.free_slots) == 4
|
||||
|
||||
|
||||
def test_page_pool_paged_bind_roundtrip():
|
||||
pool = _make_paged_pool(n_layers=1, n_kv_heads=2, head_dim=4)
|
||||
pool.task_alloc("t1", list(range(4)))
|
||||
|
||||
kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0)
|
||||
k = torch.randn(1, 4, 2, 4)
|
||||
v = torch.randn(1, 4, 2, 4)
|
||||
kv.k_buffer[0][kv.out_cache_loc] = k
|
||||
kv.v_buffer[0][kv.out_cache_loc] = v
|
||||
|
||||
indices = kv.req_to_token[kv.req_pool_indices, :4]
|
||||
gathered_k = kv.k_buffer[0][indices]
|
||||
assert torch.allclose(gathered_k, k)
|
||||
|
||||
|
||||
# ---- PagePool (paged mode, page_size>1) ----
|
||||
|
||||
|
||||
def _make_paged_pool_ps64(**kwargs):
|
||||
defaults = dict(
|
||||
n_layers=1,
|
||||
n_pages=8,
|
||||
page_size=4,
|
||||
n_kv_heads=2,
|
||||
head_dim=8,
|
||||
head_dim=4,
|
||||
max_batch_size=4,
|
||||
max_seq_len=256,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
page_size=64,
|
||||
n_tokens=512,
|
||||
)
|
||||
page_table = torch.tensor([[0, 1]], dtype=torch.long)
|
||||
k = torch.randn(1, 6, 2, 8)
|
||||
v = torch.randn(1, 6, 2, 8)
|
||||
storage.write(0, page_table, 0, k, v)
|
||||
|
||||
gk, gv = storage.gather(0, page_table, 5)
|
||||
assert gk.shape == (1, 5, 2, 8)
|
||||
defaults.update(kwargs)
|
||||
return PagePool(**defaults)
|
||||
|
||||
|
||||
def test_storage_gather_clamps_negative_padding():
|
||||
storage = Storage(
|
||||
n_layers=1,
|
||||
n_pages=8,
|
||||
page_size=4,
|
||||
n_kv_heads=2,
|
||||
head_dim=8,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
page_table = torch.tensor([[0, -1]], dtype=torch.long)
|
||||
gk, gv = storage.gather(0, page_table, 4)
|
||||
assert gk.shape == (1, 4, 2, 8)
|
||||
def test_page_pool_paged_ps64_task_alloc():
|
||||
pool = _make_paged_pool_ps64()
|
||||
prompt = list(range(200))
|
||||
assert pool.task_alloc("t1", prompt)
|
||||
assert pool.task_cached("t1") == 0
|
||||
n_pages = (200 + 63) // 64
|
||||
assert len(pool._task_pages["t1"]) == n_pages
|
||||
|
||||
|
||||
def test_page_pool_paged_ps64_task_extend_crosses_page():
|
||||
pool = _make_paged_pool_ps64()
|
||||
pool.task_alloc("t1", list(range(64)))
|
||||
assert pool.task_extend("t1", 64)
|
||||
assert len(pool._task_pages["t1"]) >= 2
|
||||
|
||||
|
||||
def test_page_pool_paged_ps64_bind_roundtrip():
|
||||
pool = _make_paged_pool_ps64(n_layers=1, n_kv_heads=2, head_dim=4)
|
||||
prompt = list(range(128))
|
||||
pool.task_alloc("t1", prompt)
|
||||
|
||||
kv = pool.bind_tasks(["t1"], [128], torch.device("cpu"), start_pos=0)
|
||||
k = torch.randn(1, 128, 2, 4)
|
||||
v = torch.randn(1, 128, 2, 4)
|
||||
kv.k_buffer[0][kv.out_cache_loc] = k
|
||||
kv.v_buffer[0][kv.out_cache_loc] = v
|
||||
|
||||
indices = kv.req_to_token[kv.req_pool_indices, :128]
|
||||
gathered_k = kv.k_buffer[0][indices]
|
||||
assert torch.allclose(gathered_k, k)
|
||||
|
||||
Reference in New Issue
Block a user