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

refactor: harden inference cache state and attention dispatch

Browse Source 
- split KVCache into phase-specific PrefillKVCache/DecodeKVCache types selected by start_pos
- unify steady-state detection in TaskCacheManager
- guard decode steady-state reuse with the cached task signature so recycled req slots cannot replay a prior generation's tokens and positions
- collapse attention backend fwd_decode/fwd_prefill into a single subclass-owned forward with a shared _check_fwd guard
- fix thread-safety gap in weight update and validate prefill inputs before KV allocation
- centralize magic constants in InferenceConfig and align docs with behavior
This commit is contained in:
ViperEkura
2026-09-04 14:28:04 +08:00
parent e13fe53475
commit ae7fc3059a
12 changed files with 244 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/cache/buffer.py
Vendored
+40 -13
View File
@@ -2,7 +2,10 @@
Layer 1 — ``KVStorage``: flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
Layer 3 — ``KVCache``: pure dataclass passed to the model for direct buffer access
Layer 3 — ``BaseKVCache``: shared fields for all cache modes
``PrefillKVCache``: prefill-specific layout
``DecodeKVCache``: decode-specific layout
``KVCache``: union type for backward compatibility
These classes have no knowledge of tasks, allocation policies, or scheduling.
They are the "dumb" physical storage layer.
@@ -10,7 +13,7 @@ They are the "dumb" physical storage layer.
import threading
from dataclasses import dataclass
from typing import List, Optional
from typing import List, Optional, Union
import torch
from torch import Tensor
@@ -74,8 +77,8 @@ class KVStorage:
@dataclass
class KVCache:
"""Pure data struct passed to model for KV cache I/O.
class BaseKVCache:
"""Shared fields for all KV cache modes.
The attention layer does raw buffer indexing — no methods, no abstraction.
"""
@@ -85,12 +88,36 @@ class KVCache:
req_to_token: Tensor
req_pool_indices: Tensor
seq_lens: Tensor
out_cache_loc: Tensor
max_len: int = 0
kv_indptr: Optional[Tensor] = None
qo_indptr: Optional[Tensor] = None
q_tile_to_batch: Optional[Tensor] = None
q_tile_to_index: Optional[Tensor] = None
decode_o_part: Optional[Tensor] = None
decode_ml_part: Optional[Tensor] = None
decode_out: Optional[Tensor] = None
max_len: int
kv_indptr: Tensor # Always present in both modes
@dataclass
class PrefillKVCache(BaseKVCache):
"""Prefill-specific KV cache layout.
Handles packed ragged batching where prompts have variable lengths.
"""
out_cache_loc: Tensor # [total_q_tokens] - flattened write locations
qo_indptr: Tensor # [B+1] - prefix sum of q_lens for unpacking
q_tile_to_batch: Tensor # [num_q_tiles] - maps Q tiles to batch indices
q_tile_to_index: Tensor # [num_q_tiles] - maps Q tiles to local indices
@dataclass
class DecodeKVCache(BaseKVCache):
"""Decode-specific KV cache layout.
Single-token incremental generation with split-KV partial results.
"""
out_cache_loc: Tensor # [B] - one write position per request
qo_indptr: Tensor # [B+1] - sequential [0, 1, 2, ..., B]
decode_o_part: Tensor # [B, max_q_heads, MAX_SPLITS, head_dim]
decode_ml_part: Tensor # [B, max_q_heads, MAX_SPLITS, 2]
decode_out: Tensor # [B, max_q_heads, head_dim]
# Backward compatibility: union type for existing code
KVCache = Union[PrefillKVCache, DecodeKVCache]
astrai/inference/cache/pool.py
Vendored
+54 -19
View File
@@ -16,7 +16,13 @@ from typing import Dict, List, Optional
import torch
from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
from astrai.inference.cache.buffer import (
DecodeKVCache,
KVCache,
KVStorage,
PrefillKVCache,
ReqToTokenPool,
)
from astrai.inference.cache.strategy import (
AllocationStrategy,
Allocator,
@@ -31,6 +37,8 @@ from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace
# continues to work unchanged after the file split.
__all__ = [
"KVCache",
"PrefillKVCache",
"DecodeKVCache",
"KVStorage",
"ReqToTokenPool",
"Allocator",
@@ -232,7 +240,20 @@ class PagePool:
)
q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles]
q_tile_to_index = workspace.q_tile_to_index[:n_tiles]
decode_o_part = decode_ml_part = decode_out = None
return PrefillKVCache(
k_buffer=self._storage.k_buffer,
v_buffer=self._storage.v_buffer,
req_to_token=self._req_pool.req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens_t,
max_len=max(seq_lens),
kv_indptr=kv_indptr,
out_cache_loc=out_cache_loc,
qo_indptr=qo_indptr,
q_tile_to_batch=q_tile_to_batch,
q_tile_to_index=q_tile_to_index,
)
else:
# ---- decode: out_cache_loc is a single column (last position) ----
write_pos = seq_lens_t - 1
@@ -241,27 +262,24 @@ class PagePool:
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]
q_tile_to_batch = q_tile_to_index = None
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)
return KVCache(
k_buffer=self._storage.k_buffer,
v_buffer=self._storage.v_buffer,
req_to_token=self._req_pool.req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens_t,
out_cache_loc=out_cache_loc,
max_len=max(seq_lens),
kv_indptr=kv_indptr,
qo_indptr=qo_indptr,
q_tile_to_batch=q_tile_to_batch,
q_tile_to_index=q_tile_to_index,
decode_o_part=decode_o_part,
decode_ml_part=decode_ml_part,
decode_out=decode_out,
)
return DecodeKVCache(
k_buffer=self._storage.k_buffer,
v_buffer=self._storage.v_buffer,
req_to_token=self._req_pool.req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens_t,
max_len=max(seq_lens),
kv_indptr=kv_indptr,
out_cache_loc=out_cache_loc,
qo_indptr=qo_indptr,
decode_o_part=decode_o_part,
decode_ml_part=decode_ml_part,
decode_out=decode_out,
)
class TaskCacheManager:
@@ -385,8 +403,25 @@ class TaskCacheManager:
@property
def bind_was_steady(self) -> bool:
"""True if the last bind was a steady-state increment (same tasks, +1 seq_lens)."""
return self._bind_was_steady
def last_task_signature_matches(self, task_ids: List[str]) -> bool:
"""Check if task_ids match the previous bind's signature.
Used by Executor to detect steady-state decode for device-to-device
token copy optimization.
"""
if self._bind_state is None:
return False
prev_sig = self._bind_state.sig
# sig is tuple of req_indices, need to map task_ids to req_indices
try:
current_sig = tuple(self._states[tid].req_idx for tid in task_ids)
return prev_sig == current_sig
except KeyError:
return False
# -- internals --
def _rollback(self, state: TaskCacheState, task_id: str):
astrai/inference/metrics.py
+5 -1
View File
@@ -7,6 +7,10 @@ from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Deque, Dict, Generator, List, Literal, Optional
from astrai.config.inference_config import InferenceConfig
_config = InferenceConfig()
@dataclass
class TaskTiming:
@@ -123,7 +127,7 @@ class MetricsCollector:
stats = metrics.get_stats()
"""
def __init__(self, max_recent: int = 128):
def __init__(self, max_recent: int = _config.max_recent_tasks):
self._timings: Dict[str, TaskTiming] = {}
self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
self._lock = threading.Lock()
astrai/inference/runtime/executor.py
+40 -6
View File
@@ -7,6 +7,7 @@ from typing import List, Optional
import torch
from torch import Tensor
from astrai.config.inference_config import InferenceConfig
from astrai.extension.backend.attention import (
CudaBackend,
get_backend,
@@ -19,6 +20,7 @@ from astrai.inference.workspace import InferenceWorkspace
from astrai.model.automodel import AutoModel
logger = logging.getLogger(__name__)
_config = InferenceConfig()
@contextmanager
@@ -114,7 +116,7 @@ def _warmup_cuda_graphs(
# shapes on first call (F.linear is the dominant cost). This also warms
# up the CUDA context (driver init) and compiles the graph-capture trace
# that follows. Custom .so kernels do NOT need this — they are pre-built.
warmup_len = 64
warmup_len = _config.prefill_warmup_len
tid = "_warmup_prefill"
if task_cache.task_alloc(tid, list(range(warmup_len))):
with (
@@ -312,9 +314,20 @@ class Executor:
):
tasks = sorted(tasks, key=lambda t: t.task_id)
batch_sz = len(tasks)
# Validate batch size bounds
if batch_sz > self._workspace.max_batch_size:
raise ValueError(
f"Batch size {batch_sz} exceeds max_batch_size "
f"{self._workspace.max_batch_size}"
)
prompt_lens = [len(t.prompt_ids) for t in tasks]
# Validate inputs before any resource allocation
if any(start_pos >= prompt_len for prompt_len in prompt_lens):
raise ValueError("prefill start_pos must precede every prompt end")
q_lens = [prompt_len - start_pos for prompt_len in prompt_lens]
input_ids = torch.tensor(
@@ -380,10 +393,18 @@ class Executor:
return []
b = len(tasks)
# Validate batch size bounds
if b > self._workspace.max_batch_size:
raise ValueError(
f"Batch size {b} exceeds max_batch_size "
f"{self._workspace.max_batch_size}"
)
ws = self._workspace
task_ids = [t.task_id for t in tasks]
cur_positions = [t.next_pos for t in tasks]
task_sig = tuple(task_ids)
cur_positions = [t.next_pos for t in tasks]
# ---- pre-replay: update input buffers in-place ----
@@ -392,9 +413,16 @@ class Executor:
# slots — fill input ids device-to-device. inference_mode guards
# the read because the source was produced under sampling's
# inference-mode context.
#
# ``cache_valid`` checks the decode cache's own task signature:
# req-index signatures in the cache manager are recycled when freed
# slots are reallocated to new tasks, so a fresh batch whose prefill
# re-bind coincides with a stale signature would otherwise replay a
# previous generation's tokens into ``input_ids``.
task_sig_match = self.task_cache.last_task_signature_matches(task_ids)
cached = self._decode_cache
sig_match = cached is not None and cached.task_sig == task_sig
if sig_match and cached.last_tokens is not None:
cache_valid = cached is not None and cached.task_sig == task_sig
if task_sig_match and cache_valid and cached.last_tokens is not None:
with torch.inference_mode():
input_ids = ws.fill_input_ids_from_device(cached.last_tokens)
else:
@@ -404,9 +432,15 @@ class Executor:
kv_cache = self.task_cache.bind(task_ids, ws)
reuse_decode_state = self.task_cache.bind_was_steady and sig_match
# Reuse sampling state only if all conditions hold:
# 1. KV bind detected steady increment (same req_indices, seq_lens +1)
# 2. Task signature matches (same task_ids in same order)
# 3. We have a valid cached decode state for THIS task set
reuse_decode_state = (
cache_valid and self.task_cache.bind_was_steady and task_sig_match
)
if reuse_decode_state:
info = self._decode_cache.sampling_info
info = cached.sampling_info
ws.position_ids[:b] += 1
else:
info = _build_sampling_batch_info(tasks, self.device)
astrai/inference/scheduler.py
+1
View File
@@ -150,6 +150,7 @@ class InferenceScheduler:
)
def _ensure_weight_update_ready(self) -> None:
"""Check weight update preconditions. Must be called under _weight_lock."""
if self._loop_thread is not None and self._loop_thread.is_alive():
raise RuntimeError("Stop the scheduler before updating model weights")
if self._task_mgr.get_active_tasks() or self._task_mgr.get_waiting_tasks():
astrai/inference/task.py
+3 -1
View File
@@ -18,6 +18,7 @@ from typing import (
from tokenizers.decoders import DecodeStream
from astrai.config.inference_config import InferenceConfig
from astrai.inference.metrics import MetricsCollector
from astrai.tokenize.tokenizer import AutoTokenizer
@@ -25,6 +26,7 @@ if TYPE_CHECKING:
from astrai.extension import AttentionBackend
STOP = object()
_config = InferenceConfig()
@dataclass(frozen=True)
@@ -85,7 +87,7 @@ class Task:
top_p: float = 1.0,
top_k: int = 50,
frequency_penalty: float = 0.0,
rep_window: int = 64,
rep_window: int = _config.default_rep_window,
backend: Optional["AttentionBackend"] = None,
):
self.task_id = task_id
astrai/inference/workspace.py
+20 -23
View File
@@ -9,8 +9,11 @@ the hot loop — a prerequisite for CUDA-graph capture.
import torch
from torch import Tensor
_MAX_SPLITS = 32
Q_TILE_ROWS = 64
from astrai.config.inference_config import InferenceConfig
_CONFIG = InferenceConfig()
MAX_SPLITS = _CONFIG.max_splits
Q_TILE_ROWS = _CONFIG.q_tile_rows
class InferenceWorkspace:
@@ -22,8 +25,9 @@ class InferenceWorkspace:
- ``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).
- ``input_ids``: per-step token IDs filled from host — values are
staged through a pinned buffer and bulk-copied into the stable
device buffer (fixed address for CUDA-graph capture).
- 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.
@@ -71,17 +75,13 @@ class InferenceWorkspace:
# Per-step token IDs. Values come from host Python lists every
# step, so the device buffer is pre-allocated (stable address for
# CUDA-graph capture) and filled via a host staging buffer. A
# double buffer keeps a copy in flight from being overwritten by
# the next fill.
# CUDA-graph capture) and filled via a host staging buffer.
self.input_ids = torch.empty(
(max_batch_size,), dtype=torch.long, device=device
)
self._pin = [
torch.empty((max_batch_size,), dtype=torch.long),
torch.empty((max_batch_size,), dtype=torch.long),
]
self._pin_idx = 0
self._pin = torch.empty(
(max_batch_size,), dtype=torch.long, pin_memory=True
)
# KV-cache bind metadata (fixed shape, written by
# ``PagePool.bind_tasks`` when the Executor passes this
@@ -124,15 +124,15 @@ class InferenceWorkspace:
# Split-KV partial-result buffers for decode (persistent, one
# global alloc per process — mirrors FlashInfer's workspace
# pattern). Shape:
# [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)
# [max_batch_size, max_q_heads, _MAX_SPLITS, 2] (ml_part)
# [max_batch_size, max_q_heads, MAX_SPLITS, head_dim] (o_part)
# [max_batch_size, max_q_heads, MAX_SPLITS, 2] (ml_part)
self.decode_o_part = torch.empty(
(max_batch_size, max_q_heads, _MAX_SPLITS, head_dim),
(max_batch_size, max_q_heads, MAX_SPLITS, head_dim),
dtype=torch.float32,
device=device,
)
self.decode_ml_part = torch.empty(
(max_batch_size, max_q_heads, _MAX_SPLITS, 2),
(max_batch_size, max_q_heads, MAX_SPLITS, 2),
dtype=torch.float32,
device=device,
)
@@ -148,16 +148,13 @@ class InferenceWorkspace:
def fill_input_ids(self, ids: "list[int]") -> Tensor:
"""Write ``ids`` into the device buffer and return ``[B]``.
Host values are staged through the double buffer and copied into the
stable device buffer (``copy_`` without pinning is synchronous, so
the alternating buffers guard against an in-flight transfer).
Host values are staged through a pinned buffer and copied synchronously
into the stable device buffer.
"""
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])
self._pin[i] = v
self.input_ids[:b].copy_(self._pin[:b])
return self.input_ids[:b]
def fill_input_ids_from_device(self, tokens: Tensor) -> Tensor: