perf: move decode split partials to InferenceWorkspace

- Replace per-.cu-file static cached tensors with workspace-managed pre-allocated buffers

- InferenceWorkspace now owns decode_o_part / decode_ml_part (mirrors FlashInfer's workspace pattern)

- KVCache carries the buffers through the backend -> C++ kernel chain

- C++ kernels accept optional pre-allocated buffers; fallback to alloc_split_partials for backward compat

- Pre-allocates once at Executor init, zero allocation in the decode hot loop

- Prerequisite for CUDA-graph capture (all kernel addresses are stable)
This commit is contained in:
2026-08-06 19:12:09 +08:00
parent d0c5debbab
commit 6f67ba8942
10 changed files with 109 additions and 36 deletions
+6 -1
View File
@@ -16,7 +16,12 @@ 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
pool.max_batch_size,
pool.max_seq_len,
max_q_heads=2,
head_dim=4,
device=pool.device,
dtype=pool.dtype,
)