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

refactor: standardize packed 3d inference

Browse Source 
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch
This commit is contained in:
ViperEkura
2026-08-16 13:24:02 +08:00
parent 0dd9a417b7
commit 3406157431
15 changed files with 301 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/inference/test_cache.py
+2 -2
View File
@@ -279,7 +279,7 @@ def test_page_pool_contiguous_bind_tasks_prefill():
task_cache.task_alloc("t1", list(range(10)))
task_cache.task_alloc("t2", list(range(10)))
kv = task_cache.bind(["t1", "t2"], _ws(pool), start_pos=0)
assert kv.out_cache_loc.shape == (2, 10)
assert kv.out_cache_loc.shape == (20,)
assert kv.out_cache_loc.dtype == torch.int32
assert kv.seq_lens.tolist() == [10, 10]
assert kv.req_pool_indices.shape == (2,)
@@ -295,7 +295,7 @@ def test_page_pool_contiguous_bind_tasks_decode():
assert task_cache.task_extend("t1", 10)
assert task_cache.task_extend("t2", 8)
kv = task_cache.bind(["t1", "t2"], _ws(pool))
assert kv.out_cache_loc.shape == (2, 1)
assert kv.out_cache_loc.shape == (2,)
assert kv.seq_lens.tolist() == [11, 9]