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

perf: use int32 paged KV indices

Browse Source 
- 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:
ViperEkura
2026-08-15 13:17:06 +08:00
parent b5afe3d7a4
commit 3fb4b8ab13
11 changed files with 135 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+4 -8
View File
@@ -531,8 +531,8 @@ class CudaBackend(AttentionBackend):
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
loc = kv_cache.out_cache_loc[:, 0]
kv_cache.k_buffer[layer_id].index_copy_(0, loc, k[:, 0])
kv_cache.v_buffer[layer_id].index_copy_(0, loc, v[:, 0])
kv_cache.k_buffer[layer_id, loc] = k[:, 0]
kv_cache.v_buffer[layer_id, loc] = v[:, 0]
q_3d = q.squeeze(1)
@@ -566,12 +566,8 @@ class CudaBackend(AttentionBackend):
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
loc = kv_cache.out_cache_loc.reshape(-1)
kv_cache.k_buffer[layer_id].index_copy_(
0, loc, k.reshape(-1, k.size(2), k.size(3))
)
kv_cache.v_buffer[layer_id].index_copy_(
0, loc, v.reshape(-1, v.size(2), v.size(3))
)
kv_cache.k_buffer[layer_id, loc] = k.reshape(-1, k.size(2), k.size(3))
kv_cache.v_buffer[layer_id, loc] = v.reshape(-1, v.size(2), v.size(3))
b = q.size(0)
q_len = q.size(1)