perf: use int32 paged KV indices

- store page-table, request-row, and cache-location indices as int32
- preserve CUDA graph replay with bit-exact logits and KV cache coverage
- improve B=1 decode latency by 1-6% across 1K-32K contexts on L20
This commit is contained in:
2026-08-15 13:17:06 +08:00
parent b5afe3d7a4
commit 3fb4b8ab13
11 changed files with 135 additions and 69 deletions
+3
View File
@@ -176,6 +176,7 @@ def test_page_pool_task_cacheable_ids_excludes_unmaterialized_tail():
def test_req_to_token_pool_alloc_free():
pool = ReqToTokenPool(4, 128, torch.device("cpu"))
assert pool.req_to_token.dtype == torch.int32
slots = pool.alloc(2)
assert len(slots) == 2
assert len(pool.free_slots) == 2
@@ -279,8 +280,10 @@ def test_page_pool_contiguous_bind_tasks_prefill():
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.dtype == torch.int32
assert kv.seq_lens.tolist() == [10, 10]
assert kv.req_pool_indices.shape == (2,)
assert kv.req_pool_indices.dtype == torch.int32
def test_page_pool_contiguous_bind_tasks_decode():