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

refactor: use single-index access and update docs for cache architecture

Browse Source 
- Replace all buffer[layer_id][loc] double indexing with buffer[layer_id, loc] single advanced indexing in cache.py and attention.py
- Revert KVStorage buffers back to 4D [n_layers, size, n_kv_heads, head_dim], remove leftover 3D reshape/view in MLA path
- Update docs/guides/inference.md, docs/developer/internals.md, docs/developer/architecture.md to reflect new PagePool/KVStorage/ReqToTokenPool/KVCache classes
This commit is contained in:
ViperEkura
2026-07-30 17:47:04 +08:00
parent deb2d7e127
commit df979b4469
6 changed files with 88 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/inference/test_cache.py
+10 -10
View File
@@ -239,12 +239,12 @@ def test_page_pool_contiguous_bind_roundtrip():
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
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]
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)
@@ -301,11 +301,11 @@ def test_page_pool_paged_bind_roundtrip():
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
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_k = kv.k_buffer[0, indices]
assert torch.allclose(gathered_k, k)
@@ -352,9 +352,9 @@ def test_page_pool_paged_ps64_bind_roundtrip():
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
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]
gathered_k = kv.k_buffer[0, indices]
assert torch.allclose(gathered_k, k)