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

perf: precompute kv_indptr once per decode step

Browse Source 
- bind_tasks builds kv_indptr (prefix sum of seq_lens) a single time
- fwd_decode/fwd_prefill reuse it instead of rebuilding per layer
- Removes 24 cumsum launches per decode step (was ~1ms/step at B=4)
- Decode B=4: 9.60 -> 7.82 ms/step (-18.5%), +22.8% tok/s
This commit is contained in:
ViperEkura
2026-08-01 16:09:26 +08:00
parent a27c8a819d
commit 4b25664c79
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+2 -5
View File
@@ -341,8 +341,7 @@ class CudaBackend(AttentionBackend):
b = q.size(0)
q_3d = q.squeeze(1)
kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
kv_indptr = kv_cache.kv_indptr
out = attn_paged_decode(
q_3d,
@@ -376,9 +375,7 @@ class CudaBackend(AttentionBackend):
b = q.size(0)
q_len = q.size(1)
kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
kv_indptr = kv_cache.kv_indptr
qo_indptr = torch.arange(b + 1, dtype=torch.int32, device=q.device) * q_len
q_flat = q.reshape(b * q_len, q.size(2), q.size(3))