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_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