From 4b25664c79dace9a604d40a43abf7873cf508c02 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 1 Aug 2026 16:09:26 +0800
Subject: [PATCH] perf: precompute kv_indptr once per decode step

- 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
---
 astrai/extension/attention_backend.py | 7 ++-----
 astrai/inference/core/cache.py        | 7 +++++++
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py
index 3d24d40..d9df74a 100644
--- a/astrai/extension/attention_backend.py
+++ b/astrai/extension/attention_backend.py
@@ -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))
diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py
index f5276f4..6bfef06 100644
--- a/astrai/inference/core/cache.py
+++ b/astrai/inference/core/cache.py
@@ -203,6 +203,8 @@ class KVCache:
         seq_lens: [batch_size] — per-request total sequence lengths
         out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
         max_len: max(seq_lens) as Python int — avoids GPU sync in decode
+        kv_indptr: [batch+1] int32 — prefix sum of seq_lens, precomputed once
+            per step so the attention backend avoids rebuilding it per layer.
     """
 
     k_buffer: Tensor
@@ -212,6 +214,7 @@ class KVCache:
     seq_lens: Tensor
     out_cache_loc: Tensor
     max_len: int = 0
+    kv_indptr: Optional[Tensor] = None
 
 
 class PagePool:
@@ -434,6 +437,9 @@ class PagePool:
                 req_pool_indices, write_pos
             ].unsqueeze(-1)
 
+        kv_indptr = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device=device)
+        kv_indptr[1:] = seq_lens_t.cumsum(0).to(torch.int32)
+
         return KVCache(
             k_buffer=self._storage.k_buffer,
             v_buffer=self._storage.v_buffer,
@@ -442,6 +448,7 @@ class PagePool:
             seq_lens=seq_lens_t,
             out_cache_loc=out_cache_loc,
             max_len=max(seq_lens),
+            kv_indptr=kv_indptr,
         )
 
     # ---- internals ----