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

perf: preallocate inference decode buffers

Browse Source 
- add InferenceWorkspace with fixed-shape per-step buffers (input_ids, decode mask, KV bind metadata) for CUDA-graph capture
- bind_tasks derives seq_lens from the pool's own _task_len tracking, dropping the seq_lens parameter
- update decode metadata in-place (position_ids, seq_lens, kv_indptr) instead of re-allocating per step
- task_extend advances _task_len in contiguous mode so the pool tracks current length
- skip log_softmax when logprobs are not requested
This commit is contained in:
ViperEkura
2026-08-03 00:55:26 +08:00
parent d033b2ef0f
commit a03504a280
6 changed files with 240 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/extension/test_backend_equivalence.py
+16 -11
View File
@@ -8,9 +8,16 @@ import torch
from astrai.extension import ATTN_BACKEND, attn_backend
from astrai.inference.core.cache import PagePool
from astrai.inference.core.workspace import InferenceWorkspace
from tests.extension.conftest import D, skip_no_kernel
def _ws(pool: PagePool) -> InferenceWorkspace:
return InferenceWorkspace(
pool.max_batch_size, pool.max_seq_len, pool.device, pool.dtype
)
@skip_no_kernel
def test_training_forward_matches_torch(cuda_model):
"""Training forward (kv_cache=None) should produce identical logits.
@@ -62,11 +69,10 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
dtype=torch.bfloat16,
)
ws = _ws(cache)
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv1 = cache.bind_tasks(
["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
)
kv1 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
with torch.inference_mode():
out_torch = model(
input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
@@ -76,9 +82,7 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
cache.task_free("t2")
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv2 = cache.bind_tasks(
["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
)
kv2 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode():
out_cuda = model(
@@ -129,11 +133,10 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
input_mask[i, : len(p)] = True
position_ids[i, : len(p)] = torch.arange(len(p), device=device)
ws = _ws(cache)
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv = cache.bind_tasks(
["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
)
kv = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
with torch.inference_mode():
model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
@@ -143,13 +146,15 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
total_len = 9
dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
kv_t = cache.bind_tasks(["t1", "t2"], [9, 7], device)
cache.task_extend("t1", 8)
cache.task_extend("t2", 6)
kv_t = cache.bind_tasks(["t1", "t2"], ws)
with torch.inference_mode():
out_torch = model(
dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
)
kv_c = cache.bind_tasks(["t1", "t2"], [9, 7], device)
kv_c = cache.bind_tasks(["t1", "t2"], ws)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode():
out_cuda = model(
tests/inference/test_cache.py
+17 -5
View File
@@ -10,6 +10,15 @@ from astrai.inference import (
ReqToTokenPool,
page_hash,
)
from astrai.inference.core.workspace import InferenceWorkspace
def _ws(pool: PagePool) -> InferenceWorkspace:
"""Workspace sized to the pool (bind_tasks requires it)."""
return InferenceWorkspace(
pool.max_batch_size, pool.max_seq_len, pool.device, pool.dtype
)
# ---- page_hash ----
@@ -216,7 +225,7 @@ def test_page_pool_contiguous_bind_tasks_prefill():
pool = _make_contiguous_pool()
pool.task_alloc("t1", list(range(10)))
pool.task_alloc("t2", list(range(10)))
kv = pool.bind_tasks(["t1", "t2"], [10, 10], torch.device("cpu"), start_pos=0)
kv = pool.bind_tasks(["t1", "t2"], _ws(pool), start_pos=0)
assert kv.out_cache_loc.shape == (2, 10)
assert kv.seq_lens.tolist() == [10, 10]
assert kv.req_pool_indices.shape == (2,)
@@ -226,7 +235,10 @@ def test_page_pool_contiguous_bind_tasks_decode():
pool = _make_contiguous_pool()
pool.task_alloc("t1", list(range(10)))
pool.task_alloc("t2", list(range(8)))
kv = pool.bind_tasks(["t1", "t2"], [11, 9], torch.device("cpu"))
# Simulate one decode extension so seq_lens advance to 11 and 9.
assert pool.task_extend("t1", 10)
assert pool.task_extend("t2", 8)
kv = pool.bind_tasks(["t1", "t2"], _ws(pool))
assert kv.out_cache_loc.shape == (2, 1)
assert kv.seq_lens.tolist() == [11, 9]
@@ -236,7 +248,7 @@ def test_page_pool_contiguous_bind_roundtrip():
pool = _make_contiguous_pool(n_layers=1, n_kv_heads=2, head_dim=4)
pool.task_alloc("t1", list(range(4)))
kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0)
kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
k = torch.randn(1, 4, 2, 4)
v = torch.randn(1, 4, 2, 4)
kv.k_buffer[0, kv.out_cache_loc] = k
@@ -298,7 +310,7 @@ def test_page_pool_paged_bind_roundtrip():
pool = _make_paged_pool(n_layers=1, n_kv_heads=2, head_dim=4)
pool.task_alloc("t1", list(range(4)))
kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0)
kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
k = torch.randn(1, 4, 2, 4)
v = torch.randn(1, 4, 2, 4)
kv.k_buffer[0, kv.out_cache_loc] = k
@@ -349,7 +361,7 @@ def test_page_pool_paged_ps64_bind_roundtrip():
prompt = list(range(128))
pool.task_alloc("t1", prompt)
kv = pool.bind_tasks(["t1"], [128], torch.device("cpu"), start_pos=0)
kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
k = torch.randn(1, 128, 2, 4)
v = torch.randn(1, 128, 2, 4)
kv.k_buffer[0, kv.out_cache_loc] = k