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

refactor: unify operator selection behind generic dispatch

Browse Source 
Add astrai/extension/dispatch.py: per-family decision tables over
composable Specs, with explicit-strict / implicit-loose resolution,
ASTR_OPS env overrides, profile presets, and explain traces. Migrate
attention (behavior-preserving facade) and rotary onto it; new tests
cover spec algebra, resolution semantics, and spec-vs-supports_call
consistency.
This commit is contained in:
ViperEkura
2026-09-01 16:35:02 +08:00
parent aabf366633
commit acbe57a0f9
5 changed files with 922 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/backend/attention.py
+155 -76
View File
@@ -21,10 +21,15 @@ Usage — mirroring ``torch.nn.attention.sdpa_kernel``:
...
Thread-safe via ``contextvars`` — each scheduler thread gets its own
active backend. Backend resolution follows a strict precedence:
active backend. Backend resolution is a thin facade over the generic
operator dispatcher (``astrai.extension.dispatch``): the three backends
are registered as the "attention" family and the decision table lives in
``_attention_records``. Resolution follows a strict precedence:
1. explicit ``attn_backend(...)`` context (wins over everything),
2. the process-wide ``ASTR_BACKEND`` environment override,
2. the process-wide ``ASTR_BACKEND`` environment override
(or an ``ASTR_OPS`` ``attention=`` entry, which wins over the legacy
variable),
3. an implicit default picked from the available backends
(cuda > flash > torch).
@@ -40,12 +45,9 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
(blhd). The backend returns ``[batch, seq_len, n_heads * head_dim]``.
"""
import contextvars
import enum
import functools
import logging
import os
import threading
from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
@@ -54,6 +56,18 @@ import torch
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.dispatch import (
CallContext,
ImplRecord,
Spec,
env_selection,
get_override,
register_env_alias,
register_family,
resolve as _dispatch_resolve,
reset_override,
set_override,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.attention import (
attn_paged_decode,
@@ -72,15 +86,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_default_backend_lock = threading.Lock()
_env_backend_name: Optional[str] = None
_env_backend: Optional["AttentionBackend"] = None
_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
contextvars.ContextVar("attn_backend", default=None)
)
# Backends are stateless — one canonical instance per class, created lazily
# and reused everywhere (resolution, fallback, context managers).
_singletons: Dict[type, "AttentionBackend"] = {}
@@ -157,28 +162,6 @@ def _resolve_default_backend() -> "AttentionBackend":
return _priority_backends()[0]
def _environment_backend() -> Optional["AttentionBackend"]:
"""Resolve the process-wide ``ASTR_BACKEND`` override, if configured."""
global _env_backend, _env_backend_name
name = os.environ.get("ASTR_BACKEND", "").strip().lower()
if not name:
return None
if name != _env_backend_name:
with _default_backend_lock:
if name != _env_backend_name:
try:
_env_backend = _resolve_backend(name)
except (ValueError, RuntimeError):
_env_backend = None
logger.warning(
"ASTR_BACKEND=%r is not a registered attention backend; "
"falling back to default resolution",
name,
)
_env_backend_name = name
return _env_backend
def _resolve_backend(
backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
) -> "AttentionBackend":
@@ -204,18 +187,44 @@ def _resolve_backend(
return _resolve_default_backend()
_ENV_WARNED: set = set()
def _environment_backend() -> Optional["AttentionBackend"]:
"""Resolve the process-wide env override (``ASTR_OPS`` or the legacy
``ASTR_BACKEND``) to a backend instance, if it names a registered one.
Invalid names warn once and are ignored, falling back to default
resolution — the override is soft, never fatal.
"""
name = env_selection("attention")
if name is None:
return None
try:
return _resolve_backend(name)
except (ValueError, RuntimeError):
message = (
f"ASTR_BACKEND/ASTR_OPS value {name!r} is not a registered "
f"attention backend; falling back to default resolution"
)
if message not in _ENV_WARNED:
_ENV_WARNED.add(message)
logger.warning(message)
return None
def get_backend(
use_default: bool = True,
) -> Optional["AttentionBackend"]:
"""Resolve the active backend: explicit context > env > default.
An ``attn_backend(...)`` context is the caller's explicit choice and
always wins. ``ASTR_BACKEND`` is a process-wide override consulted
only when no context is set. Pass ``use_default=False`` at request
submission to retain only an environment override or the caller's
:func:`attn_backend` value.
always wins. ``ASTR_BACKEND`` (or ``ASTR_OPS``) is a process-wide
override consulted only when no context is set. Pass
``use_default=False`` at request submission to retain only an
environment override or the caller's :func:`attn_backend` value.
"""
context_backend = _current_backend.get()
context_backend = get_override("attention")
if context_backend is not None:
return context_backend
env_backend = _environment_backend()
@@ -241,11 +250,11 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
...
"""
instance = _resolve_backend(backend)
token = _current_backend.set(instance)
token = set_override("attention", instance)
try:
yield instance
finally:
_current_backend.reset(token)
reset_override(token)
def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
@@ -260,6 +269,28 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
)
def _context_from_call(
q: Tensor,
kv_cache: Optional["KVCache"],
attn_mask: Optional[Tensor],
is_causal: bool,
fwd: Optional[str],
) -> CallContext:
"""Snapshot the axes the attention decision table depends on."""
return CallContext(
family="attention",
fwd=fwd,
dtype=q.dtype,
device_cuda=q.is_cuda,
ndim=q.dim(),
head_dim=q.size(-1) if q.dim() >= 1 else None,
has_cache=kv_cache is not None,
has_mask=attn_mask is not None,
grad_enabled=torch.is_grad_enabled(),
raw=(q, kv_cache, attn_mask, is_causal, fwd),
)
def attention(
q: Tensor,
k: Tensor,
@@ -300,37 +331,12 @@ def attention(
Returns:
[batch, q_len, n_heads * head_dim]
"""
if backend is not None:
selected = _resolve_backend(backend)
explicit = True
else:
context_backend = _current_backend.get()
explicit = context_backend is not None
# Resolve through the same chain as inference: explicit context >
# ASTR_BACKEND env > default. Training calls (fwd=None, no cache)
# land on the CUDA backend and fall back by capability below —
# flash when it can handle the call, else torch SDPA.
selected = get_backend()
assert selected is not None
if not selected.supports_call(q, kv_cache, attn_mask, is_causal, fwd):
if explicit:
raise RuntimeError(
f"Explicitly-set backend {type(selected).__name__} cannot "
f"handle this attention call (shape={q.shape}, "
f"dtype={q.dtype}, kv_cache={'none' if kv_cache is None else 'present'}, "
f"attn_mask={'none' if attn_mask is None else 'present'}). "
f"Remove the attn_backend() context or switch to a compatible backend."
)
selected = next(
(
candidate
for candidate in _priority_backends()
if candidate.supports_call(q, kv_cache, attn_mask, is_causal, fwd)
),
_instance(TorchNativeBackend),
)
return selected.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
ctx = _context_from_call(q, kv_cache, attn_mask, is_causal, fwd)
explicit = _resolve_backend(backend) if backend is not None else None
resolution = _dispatch_resolve("attention", ctx, explicit=explicit)
return resolution.record.obj.forward(
q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd
)
class AttentionBackend(ABC):
@@ -362,11 +368,11 @@ class AttentionBackend(ABC):
"""
def __enter__(self) -> "AttentionBackend":
self._token = _current_backend.set(self)
self._token = set_override("attention", self)
return self
def __exit__(self, *exc) -> None:
_current_backend.reset(self._token)
reset_override(self._token)
@classmethod
@abstractmethod
@@ -817,3 +823,76 @@ class FlashAttnBackend(AttentionBackend):
causal=True,
)
return out
# Family registration over the generic dispatcher: the "attention" decision
# table. The Specs mirror each backend's ``supports_call`` exactly (a unit
# test asserts they never drift). The provider is re-evaluated per
# resolution, so monkeypatching ``flash_attn_available`` (plus clearing
# ``_priority_backends``) is honored, as before.
_CLASS_TO_NAME: Dict[type, str] = {
CudaBackend: ATTN_BACKEND.CUDA.value,
FlashAttnBackend: ATTN_BACKEND.FLASH.value,
TorchNativeBackend: ATTN_BACKEND.TORCH_NATIVE.value,
}
_SPEC_CUDA = (
Spec.fwd_in("prefill", "decode")
& Spec.has_cache()
& Spec.ndim(3)
& Spec.dtype_in(torch.bfloat16)
& Spec.head_dim_in(*CudaBackend.HEAD_DIMS)
& Spec.of(
lambda ctx: is_available(f"attn_paged_{ctx.fwd}"), "paged kernels loaded"
)
)
_SPEC_FLASH = (
Spec.dtype_in(torch.float16, torch.bfloat16)
& Spec.of(lambda ctx: flash_attn_available(), "flash-attn available")
& (
(
Spec.inference()
& Spec.ndim(3)
& Spec.of(
lambda ctx: _flash_attn is not None
and hasattr(_flash_attn, "flash_attn_varlen_func"),
"varlen api present",
)
)
| (Spec.training() & Spec.mask_free())
)
)
def _attention_records() -> list:
specs = {
CudaBackend: _SPEC_CUDA,
FlashAttnBackend: _SPEC_FLASH,
TorchNativeBackend: Spec.always(),
}
return [
ImplRecord(
family="attention",
name=_CLASS_TO_NAME[type(backend)],
obj=backend,
spec=specs[type(backend)],
priority=position,
)
for position, backend in enumerate(_priority_backends())
]
def _reference_record() -> ImplRecord:
return ImplRecord(
family="attention",
name=ATTN_BACKEND.TORCH_NATIVE.value,
obj=_instance(TorchNativeBackend),
spec=Spec.always(),
priority=999,
)
register_family("attention", _attention_records, _reference_record)
register_env_alias("attention", "ASTR_BACKEND")
astrai/extension/backend/rotary.py
+44 -18
View File
@@ -1,7 +1,9 @@
"""Rotary embedding with auto-dispatch to CUDA kernel.
"""Rotary embedding dispatch (family "rotary").
Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
CUDA kernel when available, falls back to torch complex multiply otherwise.
Registered rows: the fused CUDA kernel (bf16 CUDA, inference-only) and the
torch complex-multiply fallback (autograd-safe). Selection runs through
the generic dispatcher, so ``op_backend(rotary=...)`` and
``ASTR_OPS=rotary=torch`` work exactly like for attention.
Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
@@ -10,16 +12,17 @@ freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
import torch
from torch import Tensor
from astrai.extension.dispatch import (
CallContext,
ImplRecord,
Spec,
register_family,
resolve,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.rotary import rotary_emb as _cuda_rotary
_cache = {"available": None}
def _cuda_available() -> bool:
if _cache["available"] is None:
_cache["available"] = is_available("rotary_emb")
return _cache["available"]
_SPEC_CUDA = Spec.cuda_device() & Spec.dtype_in(torch.bfloat16) & Spec.no_grad()
def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
@@ -33,6 +36,29 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
return x_out.to(dtype)
def _rotary_records() -> list:
return [
ImplRecord(
family="rotary",
name="cuda",
obj=_cuda_rotary,
spec=_SPEC_CUDA,
available=lambda: is_available("rotary_emb"),
priority=0,
),
ImplRecord(
family="rotary",
name="torch",
obj=_torch_apply,
spec=Spec.always(),
priority=99,
),
]
register_family("rotary", _rotary_records, lambda: _rotary_records()[-1])
def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
"""Apply rotary embedding to x.
@@ -43,11 +69,11 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
Returns:
[batch, seq_len, n_heads, head_dim] (bf16)
"""
if (
_cuda_available()
and not torch.is_grad_enabled()
and x.is_cuda
and x.dtype == torch.bfloat16
):
return _cuda_rotary(x, freqs_cis)
return _torch_apply(x, freqs_cis)
ctx = CallContext(
family="rotary",
dtype=x.dtype,
device_cuda=x.is_cuda,
grad_enabled=torch.is_grad_enabled(),
raw=(x, freqs_cis),
)
return resolve("rotary", ctx).record.obj(x, freqs_cis)