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
astrai/inference/core/cache.py
+77 -69
View File
@@ -20,6 +20,8 @@ from typing import Callable, Dict, List, Optional
import torch
from torch import Tensor
from astrai.inference.core.workspace import InferenceWorkspace
def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
start = page_idx * page_size
@@ -217,24 +219,6 @@ class KVCache:
kv_indptr: Optional[Tensor] = None
@dataclass
class DecodeBindCache:
"""Cached KV-addressing state for steady-state decode.
Valid for one ordered task set advancing every sequence by exactly one
token per step. ``seq_lens`` is the Python mirror used to validate the
+1 progression without a GPU round-trip; on any task-set change or
non-monotonic seq_lens the whole entry is rebuilt.
"""
sig: tuple
seq_lens: List[int]
req_pool_indices: Tensor
seq_lens_t: Tensor
kv_indptr: Tensor
inc: Tensor
class PagePool:
"""Top-level KV cache manager.
@@ -305,11 +289,13 @@ class PagePool:
self._task_pages: Dict[str, List[int]] = {}
self._lock = threading.Lock()
# Single-slot incremental cache for steady-state decode: the same
# ordered task set advances every sequence by exactly one token per
# step, so seq_lens_t and kv_indptr can be updated in-place instead
# of re-allocating + re-cumsumming. Any task-set change is a miss.
self._bind_cache: Optional[DecodeBindCache] = None
# Steady-state decode validation state: the ordered task set and its
# Python seq_lens mirror. When the same set advances every sequence
# by exactly one token per step, bind_tasks updates the stable
# buffers in-place (+=1 / +=inc) instead of re-cumsumming. Any
# task-set change is a miss and rebuilds.
self._bind_sig: Optional[tuple] = None
self._bind_seq_lens: Optional[List[int]] = None
# ---- task lifecycle ----
@@ -395,33 +381,39 @@ class PagePool:
def task_extend(self, task_id: str, pos: int) -> bool:
req_idx = self._task_req.get(task_id)
if req_idx is None:
if req_idx is None or pos >= self.max_seq_len:
return False
if self.contiguous:
return pos < self.max_seq_len
# Paged mode must also claim a physical slot for the new token;
# contiguous mode's block is pre-allocated so this is a no-op.
if not self.contiguous and not self._extend_slot(task_id, req_idx, pos):
return False
self._task_len[req_idx] = pos + 1
return True
def _extend_slot(self, task_id: str, req_idx: int, pos: int) -> bool:
"""Allocate the physical slot for one extended token (paged mode)."""
if self.page_size == 1:
slots = self._alloc_tokens(1)
if slots is None:
return False
self._task_slots.setdefault(task_id, []).extend(slots)
self._req_pool.req_to_token[req_idx, pos] = slots[0]
else:
page_idx = pos // self.page_size
existing = self._task_pages.get(task_id, [])
if page_idx >= len(existing):
p = self._alloc.alloc()
if p < 0:
return False
existing.append(p)
self._task_pages[task_id] = existing
page_offset = pos % self.page_size
page = existing[page_idx]
token_slot = page * self.page_size + page_offset
self._req_pool.req_to_token[req_idx, pos] = token_slot
return True
self._task_len[req_idx] = pos + 1
page_idx = pos // self.page_size
existing = self._task_pages.get(task_id, [])
if page_idx >= len(existing):
p = self._alloc.alloc()
if p < 0:
return False
existing.append(p)
self._task_pages[task_id] = existing
page_offset = pos % self.page_size
page = existing[page_idx]
token_slot = page * self.page_size + page_offset
self._req_pool.req_to_token[req_idx, pos] = token_slot
return True
def task_cached(self, task_id: str) -> int:
@@ -442,43 +434,59 @@ class PagePool:
def bind_tasks(
self,
task_ids: List[str],
seq_lens: List[int],
device: torch.device,
workspace: InferenceWorkspace,
device: Optional[torch.device] = None,
start_pos: Optional[int] = None,
) -> KVCache:
if device is None:
device = workspace.device
req_indices = [self._task_req[tid] for tid in task_ids]
# Per-request lengths come from the pool's own tracking (task_alloc
# sets len(prompt_ids); task_extend sets pos+1), so callers need not
# pass them.
seq_lens = [self._task_len[req_idx] for req_idx in req_indices]
b = len(task_ids)
sig = tuple(task_ids)
cache = self._bind_cache
# Write into the caller's workspace buffers (fixed addresses, sized
# to max_batch/max_seq at init) — the sole owner of the per-step
# KV bind tensors.
rpi_buf = workspace.req_pool_indices
sl_buf = workspace.seq_lens
kvp_buf = workspace.kv_indptr
inc_buf = workspace.inc
ocl_buf = workspace.out_cache_loc
incremental = (
start_pos is None
and cache is not None
and cache.sig == sig
and len(cache.seq_lens) == len(seq_lens)
and all(s == p + 1 for s, p in zip(seq_lens, cache.seq_lens))
and self._bind_sig is not None
and self._bind_sig == sig
and self._bind_seq_lens is not None
and len(self._bind_seq_lens) == b
and all(s == p + 1 for s, p in zip(seq_lens, self._bind_seq_lens))
)
if incremental:
req_pool_indices = cache.req_pool_indices
seq_lens_t = cache.seq_lens_t + 1
kv_indptr = cache.kv_indptr + cache.inc
inc = cache.inc
# Steady-state decode: advance the stable buffers in-place.
# Normal-mode buffers keep ``+=`` legal regardless of whether
# this runs inside ``torch.inference_mode()``.
sl_buf[:b] += 1
kvp_buf[: b + 1] += inc_buf[: b + 1]
req_pool_indices = rpi_buf[:b]
seq_lens_t = sl_buf[:b]
kv_indptr = kvp_buf[: b + 1]
else:
req_pool_indices = torch.tensor(
req_indices, dtype=torch.long, device=device
# Cold path: fill the stable buffers from fresh host tensors.
rpi_buf[:b].copy_(
torch.tensor(req_indices, dtype=torch.long, device=device)
)
seq_lens_t = torch.tensor(seq_lens, dtype=torch.long, device=device)
kv_indptr = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device=device)
kv_indptr[1:] = seq_lens_t.cumsum(0).to(torch.int32)
inc = torch.arange(len(seq_lens) + 1, dtype=torch.int32, device=device)
self._bind_cache = DecodeBindCache(
sig=sig,
seq_lens=list(seq_lens),
req_pool_indices=req_pool_indices,
seq_lens_t=seq_lens_t,
kv_indptr=kv_indptr,
inc=inc,
)
sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
kvp_buf[: b + 1].zero_()
kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
req_pool_indices = rpi_buf[:b]
seq_lens_t = sl_buf[:b]
kv_indptr = kvp_buf[: b + 1]
self._bind_sig = sig
self._bind_seq_lens = list(seq_lens)
if start_pos is not None:
seq_len = seq_lens[0]
@@ -487,9 +495,9 @@ class PagePool:
]
else:
write_pos = seq_lens_t - 1
out_cache_loc = self._req_pool.req_to_token[
req_pool_indices, write_pos
].unsqueeze(-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]
return KVCache(
k_buffer=self._storage.k_buffer,
astrai/inference/core/executor.py
+22 -13
View File
@@ -7,6 +7,7 @@ from torch import Tensor
from astrai.inference.core.cache import PagePool
from astrai.inference.core.task import Task
from astrai.inference.core.workspace import InferenceWorkspace
from astrai.inference.sample import sample
from astrai.model.automodel import AutoModel
from astrai.tokenize.tokenizer import AutoTokenizer
@@ -70,6 +71,17 @@ class Executor:
# any task-set change is a cache miss.
self._decode_cache: Optional[tuple] = None
# Pre-allocated fixed-shape buffers for the decode hot path
# (input_ids, decode mask, KV bind metadata). Eagerly sized at init
# so the workspace is CUDA-graph-capture friendly — no allocation
# during capture.
self._workspace = InferenceWorkspace(
max_batch_size=kv_cache.max_batch_size,
max_seq_len=kv_cache.max_seq_len,
device=self.device,
dtype=self.dtype,
)
def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
if start_pos >= prompt_len:
return
@@ -99,7 +111,9 @@ class Executor:
input_mask=input_mask,
position_ids=position_ids,
kv_cache=self.kv_cache.bind_tasks(
task_ids, [prompt_len] * batch_sz, self.device, start_pos=start_pos
task_ids,
self._workspace,
start_pos=start_pos,
),
)
@@ -123,11 +137,9 @@ class Executor:
if not tasks:
return []
input_ids = torch.tensor(
[t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
dtype=torch.long,
device=self.device,
)
input_ids = self._workspace.fill_input_ids(
[t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
).unsqueeze(1)
task_ids = [t.task_id for t in tasks]
@@ -140,7 +152,7 @@ class Executor:
and cur_positions == [p + 1 for p in cached[1]]
):
_, _, info, position_ids = cached
position_ids = position_ids + 1
position_ids += 1
self._decode_cache = (sig, cur_positions, info, position_ids)
else:
info = _build_sampling_batch_info(tasks, self.device)
@@ -150,9 +162,7 @@ class Executor:
self._decode_cache = (sig, cur_positions, info, position_ids)
total_len = max(t.next_pos for t in tasks) + 1
input_mask = position_ids[:, None, None] >= torch.arange(
total_len, device=self.device
)
input_mask = self._workspace.decode_mask(position_ids, total_len)
has_freq = bool((info.freq_penalties != 0).any())
if has_freq:
@@ -184,12 +194,11 @@ class Executor:
with torch.inference_mode():
outputs = self.model(
input_ids.unsqueeze(1),
input_ids,
input_mask=input_mask,
kv_cache=self.kv_cache.bind_tasks(
task_ids,
[t.next_pos + 1 for t in tasks],
self.device,
self._workspace,
),
position_ids=position_ids.unsqueeze(1),
)
astrai/inference/core/workspace.py
+107
View File
@@ -0,0 +1,107 @@
"""Pre-allocated buffers for the inference decode hot path.
Mirrors SGLang's pre-allocated input buffers (``input_buffers.py``): tensors
are sized once to the server's maximum dimensions and sliced to the live
batch each step, so the per-token decode loop never calls
``torch.empty``/``torch.zeros``/``torch.arange`` for the hot shapes. Fills
go through ``out=`` variants (``torch.ge``) which write into the stable
buffers instead of allocating fresh results.
All buffers are allocated eagerly at init (nothing is lazy), so the
workspace is CUDA-graph-capture friendly: the decode step reads/writes
fixed-address tensors with no allocation during capture.
"""
import torch
from torch import Tensor
class InferenceWorkspace:
"""Reusable fixed-shape per-step buffers for decode.
Families of buffers, all sized to ``max_batch_size`` / ``max_seq_len``
and sliced via views each step:
- ``decode_mask``: a ``[B, 1, total_len]`` validity mask, the RHS
``arange`` pre-computed so only a single ``torch.ge(out=)`` runs per
step.
- ``input_ids``: per-step token IDs filled from host (pinned, double-
buffered so an in-flight async H2D copy never races the next fill).
- KV-cache bind metadata (``req_pool_indices``, ``seq_lens``,
``kv_indptr``, ``inc``, ``out_cache_loc``), written by
``PagePool.bind_tasks`` when the Executor passes this workspace.
No re-allocation while the server's bounds are respected.
"""
def __init__(
self,
max_batch_size: int,
max_seq_len: int,
device: torch.device,
dtype: torch.dtype,
):
self.max_batch_size = max_batch_size
self.max_seq_len = max_seq_len
self.device = device
self.dtype = dtype
# ``position_ids[:, None, None] >= arange`` RHS, reused every step.
self.arange = torch.arange(max_seq_len, device=device)
# Decode validity mask: [max_batch, 1, max_seq_len] bool.
self.input_mask = torch.empty(
(max_batch_size, 1, max_seq_len), dtype=torch.bool, device=device
)
# Per-step token IDs. Values come from host Python lists every
# step, so the device buffer is pre-allocated and filled via an
# async copy from a double-buffered pinned host buffer (stable
# address for CUDA-graph capture; alternating buffers keep an
# in-flight copy from being overwritten by the next fill).
self.input_ids = torch.empty((max_batch_size,), dtype=torch.long, device=device)
self._pin = [
torch.empty((max_batch_size,), dtype=torch.long, pin_memory=True),
torch.empty((max_batch_size,), dtype=torch.long, pin_memory=True),
]
self._pin_idx = 0
# KV-cache bind metadata (fixed shape, written by ``PagePool.bind_tasks``
# when the Executor passes this workspace). Stable addresses make the
# decode forward CUDA-graph capturable.
self.req_pool_indices = torch.empty(
(max_batch_size,), dtype=torch.long, device=device
)
self.seq_lens = torch.empty((max_batch_size,), dtype=torch.long, device=device)
self.kv_indptr = torch.empty(
(max_batch_size + 1,), dtype=torch.int32, device=device
)
self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
self.out_cache_loc = torch.empty(
(max_batch_size, 1), dtype=torch.long, device=device
)
def fill_input_ids(self, ids: "list[int]") -> Tensor:
"""Write ``ids`` into the device buffer and return ``[B]``.
Pinned host values are copied asynchronously; the double buffer
guarantees the copy never races the next call's host writes.
"""
b = len(ids)
pin = self._pin[self._pin_idx]
self._pin_idx ^= 1
for i, v in enumerate(ids):
pin[i] = v
self.input_ids[:b].copy_(pin[:b], non_blocking=True)
return self.input_ids[:b]
def decode_mask(self, position_ids: Tensor, total_len: int) -> Tensor:
"""Return the ``[B, 1, total_len]`` validity mask for this step.
Written into the pre-allocated buffer via ``torch.ge(out=)`` — no
new tensor is allocated. ``position_ids`` is the current step's
``[B]`` positions; ``total_len`` must not exceed ``max_seq_len``.
"""
b = position_ids.size(0)
out = self.input_mask[:b, :, :total_len]
torch.ge(position_ids[:, None, None], self.arange[:total_len], out=out)
return out
astrai/inference/sample.py
+1 -1
View File
@@ -305,12 +305,12 @@ class SamplingPipeline(BaseSamplingStrategy):
return tokens, chosen
transformed = self.apply(logits, filter_value, input_ids, input_mask)
log_probs = torch.log_softmax(transformed.float(), dim=-1)
tokens = torch.multinomial(
torch.softmax(transformed, dim=-1), num_samples=1
).squeeze(-1)
if not return_logprobs:
return tokens
log_probs = torch.log_softmax(transformed.float(), dim=-1)
chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
return tokens, chosen
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