From b4d702cd14617de26ff3b4a82759e764ab2f2788 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Tue, 1 Sep 2026 16:35:02 +0800 Subject: [PATCH] refactor: unify operator selection behind generic dispatch - 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 - make the axis schema family-owned: register_family takes an axes extractor that snapshots whatever decision axes that family needs from the call, and the core only supplies the axis() predicate vocabulary plus a tensor_axes helper - drop the central CallContext dataclass; resolve and explain take the raw call arguments, so unregistered handles are probed through supports_call on the same args - migrate attention and rotary onto family-owned axes with behavior-preserving specs and spec-vs-supports_call mirror tests - replace the non-ASCII member-of glyph in spec descriptions with plain ASCII " in " --- astrai/extension/__init__.py | 30 ++ astrai/extension/backend/attention.py | 234 ++++++++++----- astrai/extension/backend/rotary.py | 65 +++-- astrai/extension/dispatch.py | 400 ++++++++++++++++++++++++++ tests/extension/test_dispatch.py | 283 ++++++++++++++++++ 5 files changed, 918 insertions(+), 94 deletions(-) create mode 100644 astrai/extension/dispatch.py create mode 100644 tests/extension/test_dispatch.py diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index d41f44f..32e9425 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -27,6 +27,22 @@ from astrai.extension.backend import ( attn_backend, get_backend, ) +from astrai.extension.dispatch import ( + Axes, + ExplicitSelectionError, + ImplRecord, + Resolution, + Spec, + axis, + explain, + explain_plan, + op_backend, + register_env_alias, + register_family, + resolve, + resolve_plan, + tensor_axes, +) from astrai.extension.loader import KERNEL_NAMES, is_available from astrai.extension.ops import ( TensorLayout, @@ -52,4 +68,18 @@ __all__ = [ "is_available", "KERNEL_NAMES", "apply_rotary_emb", + "Axes", + "ExplicitSelectionError", + "ImplRecord", + "Resolution", + "Spec", + "axis", + "explain", + "explain_plan", + "op_backend", + "register_env_alias", + "register_family", + "resolve", + "resolve_plan", + "tensor_axes", ] diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index bca3794..7631f8f 100644 --- a/astrai/extension/backend/attention.py +++ b/astrai/extension/backend/attention.py @@ -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,22 @@ import torch import torch.nn.functional as F from torch import Tensor +from astrai.extension.dispatch import ( + Axes, + ImplRecord, + Spec, + axis, + env_selection, + get_override, + register_env_alias, + register_family, + reset_override, + set_override, + tensor_axes, +) +from astrai.extension.dispatch import ( + resolve as _dispatch_resolve, +) from astrai.extension.loader import is_available from astrai.extension.ops.attention import ( attn_paged_decode, @@ -72,15 +90,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 +166,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 +191,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 +254,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 +273,24 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor: ) +def _axes( + q: Tensor, + kv_cache: Optional["KVCache"], + attn_mask: Optional[Tensor], + is_causal: bool, + fwd: Optional[str], +) -> Axes: + """Snapshot the axes the attention decision table depends on.""" + return tensor_axes( + q, + fwd=fwd, + 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, + ) + + def attention( q: Tensor, k: Tensor, @@ -300,37 +331,13 @@ 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) + explicit = _resolve_backend(backend) if backend is not None else None + resolution = _dispatch_resolve( + "attention", q, kv_cache, attn_mask, is_causal, fwd, explicit=explicit + ) + return resolution.record.obj.forward( + q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd + ) class AttentionBackend(ABC): @@ -362,11 +369,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 +824,78 @@ 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 = ( + axis("fwd").in_("prefill", "decode") + & axis("has_cache").truthy() + & axis("ndim").eq(3) + & axis("dtype").in_(torch.bfloat16) + & axis("head_dim").in_(*CudaBackend.HEAD_DIMS) + & Spec.of( + lambda ax: is_available(f"attn_paged_{ax.get('fwd')}"), "paged kernels loaded" + ) +) + +_SPEC_FLASH = ( + axis("dtype").in_(torch.float16, torch.bfloat16) + & Spec.of(lambda ax: flash_attn_available(), "flash-attn available") + & ( + ( + axis("fwd").not_none() + & axis("ndim").eq(3) + & Spec.of( + lambda ax: ( + _flash_attn is not None + and hasattr(_flash_attn, "flash_attn_varlen_func") + ), + "varlen api present", + ) + ) + | (axis("fwd").none() & axis("has_mask").falsy()) + ) +) + + +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", _axes, _attention_records, _reference_record) +register_env_alias("attention", "ASTR_BACKEND") diff --git a/astrai/extension/backend/rotary.py b/astrai/extension/backend/rotary.py index f934f0f..dd9fefd 100644 --- a/astrai/extension/backend/rotary.py +++ b/astrai/extension/backend/rotary.py @@ -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,22 @@ freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs. import torch from torch import Tensor +from astrai.extension.dispatch import ( + ImplRecord, + Spec, + axis, + register_family, + resolve, + tensor_axes, +) 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 = ( + axis("device_cuda").truthy() + & axis("dtype").in_(torch.bfloat16) + & axis("grad_enabled").eq(False) +) def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor: @@ -33,6 +41,34 @@ 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", + lambda x, freqs_cis: tensor_axes(x), + _rotary_records, + lambda: _rotary_records()[-1], +) + + def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor: """Apply rotary embedding to x. @@ -43,11 +79,4 @@ 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) + return resolve("rotary", x, freqs_cis).record.obj(x, freqs_cis) diff --git a/astrai/extension/dispatch.py b/astrai/extension/dispatch.py new file mode 100644 index 0000000..e0148ea --- /dev/null +++ b/astrai/extension/dispatch.py @@ -0,0 +1,400 @@ +"""Operator dispatch: one selection mechanism for all op families. + +A family registers three things with the core: an ``axes`` extractor whose +signature mirrors the op call and snapshots whatever decision axes *that +family* needs, an ordered list of ``ImplRecord`` rows (name, impl object, +capability ``Spec``, machine-level ``available``), and a fallback record. +The core defines no axes itself — each ``Spec`` predicates over the axes +dict produced by the family's own extractor. Resolution: explicit/context +selection (strict — raises when incapable) > ``ASTR_OPS`` env entry (soft — +falls through) > first capable row > family fallback. The rows are the +family's decision table, printable via ``explain``. + +Records flagged ``faithful=False`` change numerics (e.g. fp8) and are only +reachable through an explicit selection, never the implicit chain. +""" + +import contextvars +import logging +import os +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import torch + +logger = logging.getLogger(__name__) + +Axes = Mapping[str, Any] +Call = Tuple[Tuple, Dict[str, Any]] + + +def _fmt(value: Any) -> str: + return str(value) + + +class Spec: + """Composable, self-describing predicate over a family's axes dict.""" + + __slots__ = ("_fn", "_desc") + + def __init__(self, fn: Callable[[Axes], bool], desc: str): + self._fn = fn + self._desc = desc + + def matches(self, ax: Axes) -> bool: + return bool(self._fn(ax)) + + @property + def description(self) -> str: + return self._desc + + def __and__(self, other: "Spec") -> "Spec": + return Spec( + lambda ax: self._fn(ax) and other._fn(ax), + f"({self._desc} and {other._desc})", + ) + + def __or__(self, other: "Spec") -> "Spec": + return Spec( + lambda ax: self._fn(ax) or other._fn(ax), + f"({self._desc} or {other._desc})", + ) + + def __invert__(self) -> "Spec": + return Spec(lambda ax: not self._fn(ax), f"not({self._desc})") + + @classmethod + def always(cls) -> "Spec": + return cls(lambda ax: True, "always") + + @classmethod + def of(cls, fn: Callable[[Axes], bool], desc: str) -> "Spec": + return cls(fn, desc) + + +class Axis: + """Named-axis predicate builder: ``axis("dtype").in_(torch.bfloat16)``. + + Axis names belong to each family; the core never defines or inspects + them beyond the predicate the builder closes over. + """ + + __slots__ = ("_name",) + + def __init__(self, name: str): + self._name = name + + def in_(self, *values: Any) -> Spec: + rendered = ", ".join(_fmt(v) for v in values) + return Spec( + lambda ax: ax.get(self._name) in values, + f"{self._name} in {{{rendered}}}", + ) + + def eq(self, value: Any) -> Spec: + return Spec( + lambda ax: ax.get(self._name) == value, f"{self._name}=={_fmt(value)}" + ) + + def is_(self, value: Any) -> Spec: + return Spec( + lambda ax: ax.get(self._name) is value, f"{self._name} is {_fmt(value)}" + ) + + def none(self) -> Spec: + return Spec(lambda ax: ax.get(self._name) is None, f"{self._name} is None") + + def not_none(self) -> Spec: + return Spec( + lambda ax: ax.get(self._name) is not None, f"{self._name} is not None" + ) + + def truthy(self) -> Spec: + return Spec(lambda ax: bool(ax.get(self._name)), self._name) + + def falsy(self) -> Spec: + return Spec(lambda ax: not ax.get(self._name), f"!{self._name}") + + +def axis(name: str) -> Axis: + """Entry point for named-axis predicates; see ``Axis``.""" + return Axis(name) + + +def tensor_axes(x: torch.Tensor, **extra: Any) -> Dict[str, Any]: + """Tensor-derived axes shared by most families; opt-in, extendable.""" + return { + "dtype": x.dtype, + "device_cuda": x.is_cuda, + "grad_enabled": torch.is_grad_enabled(), + **extra, + } + + +@dataclass(frozen=True) +class ImplRecord: + """One decision-table row: an implementation plus its capability.""" + + family: str + name: str + obj: Any + spec: Spec + available: Callable[[], bool] = lambda: True + priority: int = 100 + faithful: bool = True + + +@dataclass +class OpFamily: + name: str + axes: Callable[..., Axes] + provider: Callable[[], List[ImplRecord]] + fallback: Callable[[], ImplRecord] + + +_FAMILIES: Dict[str, OpFamily] = {} +_ENV_ALIASES: Dict[str, str] = {} + +_current_overrides: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar( + "astrai_op_overrides", default={} +) + +_env_lock = threading.Lock() +_env_cache: Dict[tuple, Optional[Dict[str, str]]] = {} +_warned: set = set() + + +def register_family( + name: str, + axes: Callable[..., Axes], + provider: Callable[[], List[ImplRecord]], + fallback: Callable[[], ImplRecord], +) -> None: + """Register (or replace) a family; ``provider`` is re-evaluated per + resolution so availability changes (tests, late imports) are honored. + ``axes`` mirrors the op call signature and snapshots that family's + decision axes; unregistered handles are probed through the same args. + """ + _FAMILIES[name] = OpFamily(name, axes, provider, fallback) + + +def register_env_alias(family: str, varname: str) -> None: + """Legacy single-value env var for a family (e.g. attention → + ASTR_BACKEND); an ASTR_OPS entry wins when both are set.""" + _ENV_ALIASES[family] = varname + + +def _family(name: str) -> OpFamily: + fam = _FAMILIES.get(name) + if fam is None: + raise KeyError(f"no operator family registered under {name!r}") + return fam + + +def _warn_once(message: str) -> None: + if message not in _warned: + _warned.add(message) + logger.warning(message) + + +def set_override(family: str, handle: Any) -> contextvars.Token: + overrides = dict(_current_overrides.get()) + overrides[family] = handle + return _current_overrides.set(overrides) + + +def reset_override(token: contextvars.Token) -> None: + _current_overrides.reset(token) + + +def get_override(family: str) -> Optional[Any]: + return _current_overrides.get().get(family) + + +@contextmanager +def op_backend(**handles: Any): + """Select implementations per family for the enclosed scope:: + + with op_backend(attention="torch_native", rotary="torch"): + engine.generate(...) + + String handles are validated eagerly against the family's currently + available implementations; object handles pass through unchecked. + """ + for family, handle in handles.items(): + if isinstance(handle, str): + fam = _FAMILIES.get(family) + if fam is None: + raise ValueError(f"unknown operator family {family!r}") + if _record_for_handle(fam, handle) is None: + raise ValueError(f"Unknown {family} implementation: {handle!r}") + tokens = [set_override(f, h) for f, h in handles.items()] + try: + yield + finally: + for token in reversed(tokens): + reset_override(token) + + +def env_overrides() -> Dict[str, str]: + """Merged ASTR_OPS + legacy-alias selections (family or "profile"). + + Cached per distinct env content; unknown families / malformed entries + warn once and are dropped (soft override, never fatal). + """ + with _env_lock: + merged: Dict[str, str] = {} + raw = os.environ.get("ASTR_OPS", "").strip() + if raw: + key = ("ASTR_OPS", raw) + if key not in _env_cache: + parsed: Dict[str, str] = {} + for item in raw.split(","): + key_part, sep, value = item.strip().partition("=") + key_part, value = key_part.strip(), value.strip() + if not sep or not key_part or not value: + _warn_once(f"ASTR_OPS: ignoring malformed entry {item!r}") + continue + parsed[key_part] = value + _env_cache[key] = parsed or None + merged.update(_env_cache[key] or {}) + for fam, varname in _ENV_ALIASES.items(): + raw = os.environ.get(varname, "").strip() + if raw: + key = (varname, raw) + if key not in _env_cache: + _env_cache[key] = {fam: raw.lower()} + merged.setdefault(fam, _env_cache[key][fam]) + for fam in [f for f in merged if f not in _FAMILIES and f != "profile"]: + _warn_once(f"ASTR_OPS: unknown operator family {fam!r}; dropping it") + merged.pop(fam) + return merged + + +def env_selection(family: str) -> Optional[str]: + return env_overrides().get(family) + + +@dataclass(frozen=True) +class Resolution: + record: ImplRecord + origin: str + + +class ExplicitSelectionError(RuntimeError): + """An explicitly selected implementation cannot handle the call.""" + + +def _record_for_handle(fam: OpFamily, handle: Any) -> Optional[ImplRecord]: + records = sorted(fam.provider(), key=lambda r: r.priority) + if isinstance(handle, str): + return next((r for r in records if r.name == handle), None) + return next((r for r in records if r.obj is handle), None) + + +def _adhoc_record(family: str, handle: Any, args: Tuple, kwargs: Dict) -> ImplRecord: + """Wrap an unregistered object; capability probes its own method on + the original call arguments.""" + supports = getattr(handle, "supports_call", None) + if supports is not None: + spec = Spec.of( + lambda ax: bool(supports(*args, **kwargs)), + f"{type(handle).__name__}.supports_call", + ) + else: + spec = Spec.always() + return ImplRecord(family, type(handle).__name__, handle, spec) + + +def resolve( + family: str, *args: Any, explicit: Optional[Any] = None, **kwargs: Any +) -> Resolution: + """Resolve one family for one call (explicit-strict / implicit-loose). + + ``args``/``kwargs`` mirror the op call: the family's ``axes`` extractor + snapshots the decision axes from them, and unregistered handles are + probed through their own ``supports_call`` with the same arguments. + """ + fam = _family(family) + ax = fam.axes(*args, **kwargs) + + handle: Optional[Any] = None + origin = "chain" + if explicit is not None: + handle, origin = explicit, "explicit" + elif get_override(family) is not None: + handle, origin = get_override(family), "context" + else: + env_name = env_selection(family) + if env_name is not None: + handle, origin = env_name, "env" + + if handle is not None: + record = _record_for_handle(fam, handle) + if record is None and not isinstance(handle, str): + record = _adhoc_record(family, handle, args, kwargs) + if record is None: + if origin in ("explicit", "context"): + raise ValueError(f"Unknown {family} implementation: {handle!r}") + _warn_once(f"ASTR_OPS: {family}={handle!r} is not registered; ignoring") + else: + if record.available() and record.spec.matches(ax): + return Resolution(record, origin) + if origin in ("explicit", "context"): + raise ExplicitSelectionError( + f"Explicitly-set backend {type(record.obj).__name__} cannot " + f"handle this {family} call; required: {record.spec.description}" + ) + + if handle is None and env_overrides().get("profile") == "reference": + return Resolution(fam.fallback(), "profile") + + for record in sorted(fam.provider(), key=lambda r: r.priority): + if record.available() and record.faithful and record.spec.matches(ax): + return Resolution(record, "chain") + return Resolution(fam.fallback(), "fallback") + + +def resolve_plan(calls: Mapping[str, Call]) -> Dict[str, Resolution]: + """Resolve several families at once (one decision snapshot).""" + return { + family: resolve(family, *args, **kwargs) + for family, (args, kwargs) in calls.items() + } + + +def _describe_axes(ax: Axes) -> str: + return " ".join(f"{key}={ax[key]}" for key in sorted(ax)) + + +def explain( + family: str, *args: Any, explicit: Optional[Any] = None, **kwargs: Any +) -> str: + """Human-readable decision trace for one family call.""" + fam = _family(family) + ax = fam.axes(*args, **kwargs) + records = sorted(fam.provider(), key=lambda r: r.priority) + lines = [f"[{family}] {_describe_axes(ax)}"] + for record in records: + if not record.available(): + lines.append(f" {record.name}: SKIP unavailable") + elif not record.faithful: + lines.append(f" {record.name}: SKIP not faithful (explicit-only)") + elif record.spec.matches(ax): + lines.append(f" {record.name}: MATCH ({record.spec.description})") + else: + lines.append(f" {record.name}: reject ({record.spec.description})") + try: + resolution = resolve(family, *args, explicit=explicit, **kwargs) + lines.append(f" => {resolution.record.name} (origin={resolution.origin})") + except (ExplicitSelectionError, ValueError) as exc: + lines.append(f" => ERROR: {exc}") + return "\n".join(lines) + + +def explain_plan(calls: Mapping[str, Call]) -> str: + return "\n".join( + explain(family, *args, **kwargs) for family, (args, kwargs) in calls.items() + ) diff --git a/tests/extension/test_dispatch.py b/tests/extension/test_dispatch.py new file mode 100644 index 0000000..bb02318 --- /dev/null +++ b/tests/extension/test_dispatch.py @@ -0,0 +1,283 @@ +"""Tests for the generic operator dispatcher (Spec / decision tables).""" + +import importlib +from types import SimpleNamespace + +import pytest +import torch + +import astrai.extension.dispatch as dispatch +from astrai.extension import ( + ATTN_BACKEND, + ExplicitSelectionError, + ImplRecord, + Spec, + axis, + explain, + op_backend, + resolve, + resolve_plan, +) +from astrai.extension.backend import apply_rotary_emb + +attn_mod = importlib.import_module("astrai.extension.backend.attention") +rotary_mod = importlib.import_module("astrai.extension.backend.rotary") + + +@pytest.fixture +def toy_family(): + """A toy family: alpha (restricted), beta, and an unfaithful fast row.""" + calls = [] + + def records(): + return [ + ImplRecord( + "toy", + "alpha", + "alpha-obj", + axis("dtype").in_(torch.bfloat16), + priority=0, + ), + ImplRecord( + "toy", + "beta", + "beta-obj", + Spec.always(), + priority=10, + ), + ImplRecord( + "toy", + "fp8", + "fp8-obj", + Spec.always(), + priority=1, + faithful=False, + ), + ] + + dispatch.register_family( + "toy", + lambda **kw: kw, + records, + lambda: ImplRecord("toy", "beta", "beta-obj", Spec.always()), + ) + yield calls + dispatch._FAMILIES.pop("toy", None) + + +def test_spec_composition_and_description(): + spec = axis("dtype").in_(torch.bfloat16) & axis("grad_enabled").eq(False) + assert spec.matches({"dtype": torch.bfloat16, "grad_enabled": False}) + assert not spec.matches({"dtype": torch.bfloat16, "grad_enabled": True}) + assert "dtype" in spec.description and "grad_enabled" in spec.description + + either = axis("fwd").none() | axis("has_cache").truthy() + assert either.matches({"fwd": None}) + assert either.matches({"fwd": "decode", "has_cache": True}) + assert not either.matches({"fwd": "decode"}) + + assert (~axis("fwd").none()).matches({"fwd": "decode"}) + + +def test_chain_returns_first_capable(toy_family): + assert resolve("toy", dtype=torch.bfloat16).record.obj == "alpha-obj" + assert resolve("toy", dtype=torch.float32).record.obj == "beta-obj" + + +def test_unfaithful_rows_are_chain_invisible(toy_family): + resolution = resolve("toy", dtype=torch.float32) + assert resolution.record.obj == "beta-obj" + with op_backend(toy="fp8"): + assert resolve("toy", dtype=torch.float32).record.obj == "fp8-obj" + + +def test_explicit_selection_is_strict(toy_family): + with pytest.raises(ExplicitSelectionError): + resolve("toy", dtype=torch.float32, explicit="alpha") + assert resolve("toy", dtype=torch.bfloat16, explicit="alpha").origin == "explicit" + + +def test_context_selection_is_strict(toy_family): + with op_backend(toy="alpha"): + with pytest.raises(ExplicitSelectionError): + resolve("toy", dtype=torch.float32) + assert resolve("toy", dtype=torch.bfloat16).origin == "context" + + +def test_unknown_explicit_name_raises(toy_family): + with pytest.raises(ValueError, match="Unknown toy implementation"): + resolve("toy", explicit="nope") + with pytest.raises(ValueError, match="Unknown toy implementation"): + with op_backend(toy="nope"): + pass + + +class _Probe: + def __init__(self, capable): + self.capable = capable + self.probed = 0 + + def supports_call(self, *args, **kwargs): + self.probed += 1 + return self.capable + + +def test_adhoc_instance_probed_via_supports_call(toy_family): + probe = _Probe(capable=True) + resolution = resolve("toy", dtype=torch.float32, explicit=probe) + assert resolution.record.obj is probe and probe.probed == 1 + + incapable = _Probe(capable=False) + with pytest.raises(ExplicitSelectionError): + resolve("toy", dtype=torch.float32, explicit=incapable) + + +def test_nested_op_backend_scopes(toy_family): + with op_backend(toy="alpha"): + with op_backend(toy="beta"): + assert resolve("toy", dtype=torch.float32).origin == "context" + assert resolve("toy", dtype=torch.bfloat16).origin == "context" + + +def test_env_entry_is_soft(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "toy=alpha") + resolution = resolve("toy", dtype=torch.float32) + assert resolution.record.obj == "beta-obj" and resolution.origin == "chain" + assert resolve("toy", dtype=torch.bfloat16).origin == "env" + + +def test_env_unknown_impl_ignored(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "toy=missing") + assert resolve("toy", dtype=torch.float32).record.obj == "beta-obj" + + +def test_env_profile_reference(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "profile=reference") + resolution = resolve("toy", dtype=torch.bfloat16) + assert resolution.origin == "profile" + monkeypatch.setenv("ASTR_OPS", "toy=alpha,profile=reference") + assert resolve("toy", dtype=torch.bfloat16).origin == "env" + + +def test_legacy_env_alias(monkeypatch): + monkeypatch.setenv("ASTR_BACKEND", "torch_native") + assert dispatch.env_selection("attention") == "torch_native" + monkeypatch.setenv("ASTR_OPS", "attention=cuda") + assert dispatch.env_selection("attention") == "cuda" + + +def test_context_beats_env(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "toy=alpha") + with op_backend(toy="beta"): + assert resolve("toy", dtype=torch.float32).origin == "context" + + +def test_resolve_plan_snapshots_families(toy_family): + plan = resolve_plan( + { + "toy": ((), {"dtype": torch.bfloat16}), + "rotary": ( + ( + SimpleNamespace(dtype=torch.bfloat16, is_cuda=True), + SimpleNamespace(), + ), + {}, + ), + } + ) + assert plan["toy"].record.obj == "alpha-obj" + assert plan["rotary"].record.name in ("cuda", "torch") + + +def test_explain_shows_rejection_reasons(toy_family): + text = explain("toy", dtype=torch.float32) + assert "alpha: reject" in text and "beta: MATCH" in text + assert "=> beta" in text + + +def test_explain_reports_strict_error(toy_family): + text = explain("toy", dtype=torch.float32, explicit="alpha") + assert "ERROR" in text + + +_DUMMY_CACHE = object() + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("head_dim", [64, 96]) +@pytest.mark.parametrize( + "fwd,has_cache,ndim,has_mask", + [ + ("decode", True, 3, False), + ("prefill", True, 3, False), + (None, False, 4, False), + (None, False, 4, True), + ], +) +def test_attention_specs_mirror_supports_call( + dtype, head_dim, fwd, has_cache, ndim, has_mask +): + shape = {3: (1, 2, head_dim), 4: (1, 2, 4, head_dim)}[ndim] + q = torch.zeros(shape, dtype=dtype) + mask = torch.zeros(1, 1, 2, 2, dtype=torch.bool) if has_mask else None + cache = _DUMMY_CACHE if has_cache else None + ax = attn_mod._axes(q, cache, mask, False, fwd) + + cuda = attn_mod._instance(attn_mod.CudaBackend) + assert attn_mod._SPEC_CUDA.matches(ax) == cuda.supports_call( + q, cache, mask, False, fwd + ) + + flash = attn_mod._instance(attn_mod.FlashAttnBackend) + assert attn_mod._SPEC_FLASH.matches(ax) == flash.supports_call( + q, cache, mask, False, fwd + ) + + +def test_attention_resolution_matches_legacy_semantics(): + q = torch.zeros(1, 2, 4, 8, dtype=torch.float32) + resolution = resolve("attention", q, None, None, True, "prefill") + assert resolution.record.name == ATTN_BACKEND.TORCH_NATIVE.value + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="rotary CUDA path needs a GPU" +) +class TestRotaryDispatch: + def _input(self): + torch.manual_seed(0) + x = torch.randn(1, 5, 3, 16, device="cuda", dtype=torch.bfloat16) + freqs = torch.randn(1, 5, 8, 2, device="cuda", dtype=torch.float32) + return x, freqs + + def test_cuda_row_selected_under_inference_mode(self): + from astrai.extension.loader import is_available + + x, freqs = self._input() + with torch.inference_mode(): + resolution = resolve("rotary", x, freqs) + expected = "cuda" if is_available("rotary_emb") else "torch" + assert resolution.record.name == expected + + def test_grad_falls_back_to_torch(self): + x, freqs = self._input() + assert resolve("rotary", x, freqs).record.name == "torch" + + def test_context_switch_to_torch(self): + x, freqs = self._input() + with torch.inference_mode(): + with op_backend(rotary="torch"): + out = apply_rotary_emb(x, freqs) + assert out.shape == x.shape and out.dtype == torch.bfloat16 + + def test_cuda_matches_torch_numerics(self): + from astrai.extension.loader import is_available + + if not is_available("rotary_emb"): + pytest.skip("rotary kernel not built") + x, freqs = self._input() + with torch.inference_mode(): + fast = apply_rotary_emb(x, freqs) + slow = rotary_mod._torch_apply + ref = slow(x, freqs) + assert torch.allclose(fast.float(), ref.float(), atol=2e-2, rtol=1e-2)