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)
astrai/extension/attention_ops.py
+4 -4
View File
@@ -113,8 +113,8 @@ def attn_paged_decode(
q: [batch, n_heads, head_dim] (bf16, 3D — no seq dim)
k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
v_cache: same as k_cache
req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
req_pool_indices: [batch] (int64) — rows into req_to_token
req_to_token: [num_reqs, max_context_len] (int32) — token -> slot
req_pool_indices: [batch] (int32) — rows into req_to_token
kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
mask: 2D [batch, max_context_len] (bool, True=keep) or None
is_causal: apply causal mask
@@ -163,8 +163,8 @@ def attn_paged_prefill(
q: [total_q, n_heads, head_dim] (bf16, 3D — flattened across requests)
k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
v_cache: same as k_cache
req_to_token: [num_reqs, max_context_len] (int64)
req_pool_indices: [batch] (int64)
req_to_token: [num_reqs, max_context_len] (int32)
req_pool_indices: [batch] (int32)
kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None