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/config/inference_config.py
+23
View File
@@ -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
astrai/extension/backend/attention.py
+39 -81
View File
@@ -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)
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
+42 -7
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,23 +262,20 @@ 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(
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,
out_cache_loc=out_cache_loc,
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,
decode_o_part=decode_o_part,
decode_ml_part=decode_ml_part,
decode_out=decode_out,
@@ -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:
docs/developer/cuda_kernels.md
+1 -1
View File
@@ -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)
tests/extension/test_backend.py
+11 -7
View File
@@ -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
tests/inference/test_scheduler.py
+7
View File
@@ -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