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

refactor: standardize packed 3d inference

Browse Source 
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch
This commit is contained in:
ViperEkura
2026-08-16 13:24:02 +08:00
parent 0dd9a417b7
commit 3406157431
15 changed files with 301 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/cache/pool.py
Vendored
+18 -10
View File
@@ -200,14 +200,21 @@ class PagePool:
kv_indptr = kvp_buf[: b + 1]
if start_pos is not None:
# ---- prefill: out_cache_loc covers prefix range [start_pos:seq_len] ----
seq_len = seq_lens[0]
out_cache_loc = self._req_pool.req_to_token[
req_pool_indices, start_pos:seq_len
]
q_len = seq_len - start_pos
workspace.qo_indptr[: b + 1].copy_(
torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
# Packed prefill concatenates each request's query tokens.
q_lens = [seq_len - start_pos for seq_len in seq_lens]
if any(q_len <= 0 for q_len in q_lens):
raise ValueError("prefill sequence lengths must exceed start_pos")
out_cache_loc = torch.cat(
[
self._req_pool.req_to_token[
req_pool_indices[i], start_pos : seq_lens[i]
]
for i in range(b)
]
)
workspace.qo_indptr[: b + 1].zero_()
workspace.qo_indptr[1 : b + 1].copy_(
torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(0)
)
qo_indptr = workspace.qo_indptr[: b + 1]
decode_o_part = decode_ml_part = decode_out = None
@@ -216,8 +223,9 @@ class PagePool:
write_pos = seq_lens_t - 1
loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
ocl_buf[:b].copy_(loc)
out_cache_loc = ocl_buf[:b]
qo_indptr = None
out_cache_loc = ocl_buf[:b].reshape(-1)
workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
qo_indptr = workspace.qo_indptr[: b + 1]
decode_o_part = getattr(workspace, "decode_o_part", None)
decode_ml_part = getattr(workspace, "decode_ml_part", None)
decode_out = getattr(workspace, "decode_out", None)