From ae7fc3059ae98787d9d7930d3981b815d55657a5 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 4 Sep 2026 14:28:04 +0800 Subject: [PATCH] refactor: harden inference cache state and attention dispatch - 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 --- astrai/config/inference_config.py | 23 +++++ astrai/extension/backend/attention.py | 120 +++++++++----------------- astrai/inference/cache/buffer.py | 53 +++++++++--- astrai/inference/cache/pool.py | 73 ++++++++++++---- astrai/inference/metrics.py | 6 +- astrai/inference/runtime/executor.py | 46 ++++++++-- astrai/inference/scheduler.py | 1 + astrai/inference/task.py | 4 +- astrai/inference/workspace.py | 43 +++++---- docs/developer/cuda_kernels.md | 2 +- tests/extension/test_backend.py | 18 ++-- tests/inference/test_scheduler.py | 7 ++ 12 files changed, 244 insertions(+), 152 deletions(-) create mode 100644 astrai/config/inference_config.py diff --git a/astrai/config/inference_config.py b/astrai/config/inference_config.py new file mode 100644 index 0000000..baa4c63 --- /dev/null +++ b/astrai/config/inference_config.py @@ -0,0 +1,23 @@ +"""Inference engine configuration.""" + +from astrai.config.base import BaseConfig + + +class InferenceConfig(BaseConfig): + """Configuration for inference workspace and execution parameters. + + Centralizes magic constants previously scattered across inference modules. + + Args: + max_splits (int): Maximum number of splits for split-KV attention (decode partial results). Defaults to 32. + q_tile_rows (int): Number of rows per Q tile in prefill ragged batching. Defaults to 64. + prefill_warmup_len (int): Prompt length for prefill warmup (cuBLAS auto-tuning). Defaults to 64. + default_rep_window (int): Default repetition penalty window size for frequency penalty. Defaults to 64. + max_recent_tasks (int): Maximum number of recent tasks tracked for aggregate statistics. Defaults to 128. + """ + + max_splits: int = 32 + q_tile_rows: int = 64 + prefill_warmup_len: int = 64 + default_rep_window: int = 64 + max_recent_tasks: int = 128 diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index 7631f8f..7f2f327 100644 --- a/astrai/extension/backend/attention.py +++ b/astrai/extension/backend/attention.py @@ -343,9 +343,11 @@ def attention( class AttentionBackend(ABC): """Abstract base for attention computation strategies. - Subclasses implement ``fwd_decode`` (q_len == 1, with cache) and - ``fwd_prefill`` (q_len > 1, with or without cache). The public - ``forward`` method dispatches based on q_len. + Subclasses implement a single ``forward`` and branch on ``fwd`` + ("decode" / "prefill", or None for training) wherever their kernels + split — the mode taxonomy is the caller's, not the base class's, so + it lives in the implementations. ``_check_fwd`` is the shared guard + against unknown mode strings. Capability contract — every backend declares: @@ -365,7 +367,6 @@ class AttentionBackend(ABC): with attn_backend(TorchNativeBackend): # class ... with TorchNativeBackend(): # instance - ... """ def __enter__(self) -> "AttentionBackend": @@ -398,7 +399,15 @@ class AttentionBackend(ABC): Called on the canonical singleton instance (or a caller-provided one); must be side-effect free. """ + return True + @staticmethod + def _check_fwd(fwd: Optional[str]) -> None: + """Reject unknown forward modes loudly.""" + if fwd not in (None, "prefill", "decode"): + raise ValueError(f"unsupported attention forward mode: {fwd}") + + @abstractmethod def forward( self, q: Tensor, @@ -410,7 +419,7 @@ class AttentionBackend(ABC): is_causal: bool = False, fwd: Optional[str] = None, ) -> Tensor: - """Dispatch to decode or extend based on q_len. + """Run one attention call; ``fwd`` selects the mode. Args: q: [batch, q_len, n_heads, head_dim] @@ -420,41 +429,11 @@ class AttentionBackend(ABC): layer_id: transformer layer index for buffer access. attn_mask: pre-built attention mask compatible with SDPA. is_causal: whether to apply causal masking. + fwd: "prefill" / "decode" for inference, None for training. Returns: [batch, q_len, n_heads * head_dim] """ - if fwd == "decode": - return self.fwd_decode(q, k, v, kv_cache, layer_id, attn_mask, is_causal) - if fwd == "prefill" or fwd is None: - return self.fwd_prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal) - raise ValueError(f"unsupported attention forward mode: {fwd}") - - @abstractmethod - def fwd_decode( - self, - q: Tensor, - k: Tensor, - v: Tensor, - kv_cache: Optional["KVCache"], - layer_id: int, - attn_mask: Optional[Tensor] = None, - is_causal: bool = False, - ) -> Tensor: - """Single-token decode with KV cache.""" - - @abstractmethod - def fwd_prefill( - self, - q: Tensor, - k: Tensor, - v: Tensor, - kv_cache: Optional["KVCache"], - layer_id: int, - attn_mask: Optional[Tensor] = None, - is_causal: bool = False, - ) -> Tensor: - """Multi-token prefill or training forward.""" @staticmethod def supports_graph() -> bool: @@ -498,31 +477,7 @@ class TorchNativeBackend(AttentionBackend): ) -> bool: return True - def fwd_decode( - self, - q: Tensor, - k: Tensor, - v: Tensor, - kv_cache: Optional["KVCache"], - layer_id: int, - attn_mask: Optional[Tensor] = None, - is_causal: bool = False, - ) -> Tensor: - return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal) - - def fwd_prefill( - self, - q: Tensor, - k: Tensor, - v: Tensor, - kv_cache: Optional["KVCache"], - layer_id: int, - attn_mask: Optional[Tensor] = None, - is_causal: bool = False, - ) -> Tensor: - return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal) - - def _forward( + def forward( self, q: Tensor, k: Tensor, @@ -531,7 +486,9 @@ class TorchNativeBackend(AttentionBackend): layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, + fwd: Optional[str] = None, ) -> Tensor: + self._check_fwd(fwd) if q.ndim == 4: n_rep = q.size(2) // k.size(2) if n_rep > 1: @@ -633,7 +590,7 @@ class CudaBackend(AttentionBackend): def supports_graph() -> bool: return True - def fwd_decode( + def forward( self, q: Tensor, k: Tensor, @@ -642,10 +599,23 @@ class CudaBackend(AttentionBackend): layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, + fwd: Optional[str] = None, ) -> Tensor: + self._check_fwd(fwd) if kv_cache is None: raise RuntimeError("CudaBackend does not support training (kv_cache=None)") + if fwd == "decode": + return self._decode(q, k, v, kv_cache, layer_id) + return self._prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal) + def _decode( + self, + q: Tensor, + k: Tensor, + v: Tensor, + kv_cache: "KVCache", + layer_id: int, + ) -> Tensor: kv_indptr = kv_cache.kv_indptr out = attn_paged_decode( @@ -664,19 +634,16 @@ class CudaBackend(AttentionBackend): ) return out - def fwd_prefill( + def _prefill( self, q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional["KVCache"], + kv_cache: "KVCache", layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, ) -> Tensor: - if kv_cache is None: - raise RuntimeError("CudaBackend does not support training (kv_cache=None)") - loc = kv_cache.out_cache_loc kv_cache.k_buffer[layer_id, loc] = k kv_cache.v_buffer[layer_id, loc] = v @@ -734,7 +701,7 @@ class FlashAttnBackend(AttentionBackend): # back to TorchNativeBackend instead of silently ignoring the mask. return attn_mask is None - def fwd_decode( + def forward( self, q: Tensor, k: Tensor, @@ -743,20 +710,11 @@ class FlashAttnBackend(AttentionBackend): layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, + fwd: Optional[str] = None, ) -> Tensor: - return self._forward_packed(q, k, v, kv_cache, layer_id) - - def fwd_prefill( - self, - q: Tensor, - k: Tensor, - v: Tensor, - kv_cache: Optional["KVCache"], - layer_id: int, - attn_mask: Optional[Tensor] = None, - is_causal: bool = False, - ) -> Tensor: - if q.ndim == 3: + self._check_fwd(fwd) + # Decode is always packed; prefill/training split by layout. + if fwd == "decode" or q.ndim == 3: return self._forward_packed(q, k, v, kv_cache, layer_id) return self._forward_dense(q, k, v, attn_mask, is_causal) diff --git a/astrai/inference/cache/buffer.py b/astrai/inference/cache/buffer.py index b43a1aa..6f37adb 100644 --- a/astrai/inference/cache/buffer.py +++ b/astrai/inference/cache/buffer.py @@ -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] diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py index 2c04bb1..c6302d5 100644 --- a/astrai/inference/cache/pool.py +++ b/astrai/inference/cache/pool.py @@ -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): diff --git a/astrai/inference/metrics.py b/astrai/inference/metrics.py index 9742a7f..77371e3 100644 --- a/astrai/inference/metrics.py +++ b/astrai/inference/metrics.py @@ -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() diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py index 8d21c95..54f3585 100644 --- a/astrai/inference/runtime/executor.py +++ b/astrai/inference/runtime/executor.py @@ -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) diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 12e2a1a..879e7ac 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -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(): diff --git a/astrai/inference/task.py b/astrai/inference/task.py index 55ce442..1156563 100644 --- a/astrai/inference/task.py +++ b/astrai/inference/task.py @@ -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 diff --git a/astrai/inference/workspace.py b/astrai/inference/workspace.py index c676a90..a70d27a 100644 --- a/astrai/inference/workspace.py +++ b/astrai/inference/workspace.py @@ -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: diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 7704938..c909842 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -417,7 +417,7 @@ cycle belong under `TYPE_CHECKING`. `astrai/extension/backend/attention.py` provides the backend abstraction: -- **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len +- **`AttentionBackend`** (ABC): single abstract `forward`; each subclass branches on `fwd` ("decode" / "prefill" / None) internally, `_check_fwd` guards unknown modes - **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU. - **`FlashAttnBackend`**: Optional flash-attn dispatch via `flash_attn_varlen_func` over gathered flat K/V. - **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback) diff --git a/tests/extension/test_backend.py b/tests/extension/test_backend.py index 79b0c91..9c8e77b 100644 --- a/tests/extension/test_backend.py +++ b/tests/extension/test_backend.py @@ -154,14 +154,18 @@ class _DummyBackend(AttentionBackend): def supports_call(self, q, kv_cache, attn_mask, is_causal, fwd) -> bool: return True - def fwd_decode( - self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False - ): - return q - - def fwd_prefill( - self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False + def forward( + self, + q, + k, + v, + kv_cache=None, + layer_id=0, + attn_mask=None, + is_causal=False, + fwd=None, ): + self._check_fwd(fwd) return q diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 7f504c0..d867381 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -185,6 +185,7 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits(): executor.task_cache = MagicMock() executor.task_cache.bind.return_value = MagicMock() executor._workspace = MagicMock() + executor._workspace.max_batch_size = 16 # Add max_batch_size for validation all_logits = torch.arange(42, dtype=torch.float32).reshape(6, 7) executor.model = MagicMock(return_value={"logits": all_logits}) executor._sample_logits = MagicMock( @@ -721,11 +722,15 @@ def test_decode_does_not_reuse_previous_batch_state(): executor.device = torch.device("cpu") executor.task_cache = MagicMock() executor.task_cache.bind_was_steady = True + executor.task_cache.last_task_signature_matches.return_value = ( + False # Different task + ) executor.task_cache.bind.return_value = MagicMock() executor._graph_supported = False executor._graph_ctx = SimpleNamespace(enabled=False) workspace = MagicMock() + workspace.max_batch_size = 16 workspace.position_ids = torch.tensor([2], dtype=torch.long) workspace.fill_input_ids.return_value = torch.tensor([7], dtype=torch.long) workspace.decode_mask.return_value = torch.ones(1, 1, 9, dtype=torch.bool) @@ -766,11 +771,13 @@ def test_decode_fills_input_ids_from_device_on_matching_signature(): executor.device = torch.device("cpu") executor.task_cache = MagicMock() executor.task_cache.bind_was_steady = True + executor.task_cache.last_task_signature_matches.return_value = True # Same task executor.task_cache.bind.return_value = MagicMock() executor._graph_supported = False executor._graph_ctx = SimpleNamespace(enabled=False) workspace = MagicMock() + workspace.max_batch_size = 16 workspace.position_ids = torch.tensor([2], dtype=torch.long) workspace.fill_input_ids_from_device.return_value = torch.tensor( [9], dtype=torch.long