From acbe57a0f9f4bbd15bfd89f91b640c7880d84db5 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. Migrate attention (behavior-preserving facade) and rotary onto it; new tests cover spec algebra, resolution semantics, and spec-vs-supports_call consistency. --- astrai/extension/__init__.py | 26 ++ astrai/extension/backend/attention.py | 231 ++++++++++----- astrai/extension/backend/rotary.py | 62 ++-- astrai/extension/dispatch.py | 388 ++++++++++++++++++++++++++ tests/extension/test_dispatch.py | 309 ++++++++++++++++++++ 5 files changed, 922 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..a9cdc1f 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -27,6 +27,20 @@ from astrai.extension.backend import ( attn_backend, get_backend, ) +from astrai.extension.dispatch import ( + CallContext, + ExplicitSelectionError, + ImplRecord, + Resolution, + Spec, + explain, + explain_plan, + op_backend, + register_env_alias, + register_family, + resolve, + resolve_plan, +) from astrai.extension.loader import KERNEL_NAMES, is_available from astrai.extension.ops import ( TensorLayout, @@ -52,4 +66,16 @@ __all__ = [ "is_available", "KERNEL_NAMES", "apply_rotary_emb", + "CallContext", + "ExplicitSelectionError", + "ImplRecord", + "Resolution", + "Spec", + "explain", + "explain_plan", + "op_backend", + "register_env_alias", + "register_family", + "resolve", + "resolve_plan", ] diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index bca3794..73e1435 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,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") diff --git a/astrai/extension/backend/rotary.py b/astrai/extension/backend/rotary.py index f934f0f..942d6dd 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,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) diff --git a/astrai/extension/dispatch.py b/astrai/extension/dispatch.py new file mode 100644 index 0000000..cd2a6af --- /dev/null +++ b/astrai/extension/dispatch.py @@ -0,0 +1,388 @@ +"""Operator dispatch: one selection mechanism for all op families. + +A family registers an ordered list of ``ImplRecord`` rows (name, impl +object, capability ``Spec``, machine-level ``available``). Resolution: +explicit/context selection (strict — raises when incapable) > ``ASTR_OPS`` +env entry (soft — falls through) > first capable row > family fallback. +``Spec`` is a composable predicate over ``CallContext``; 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, field, replace +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import torch + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class CallContext: + """Axes a decision may depend on, snapshotted per call. + + ``raw``/``raw_kwargs`` keep the original call so unregistered handles + can still be probed through their own capability method. + """ + + family: str + fwd: Optional[str] = None + dtype: Optional[torch.dtype] = None + device_cuda: bool = False + ndim: Optional[int] = None + head_dim: Optional[int] = None + has_cache: bool = False + has_mask: bool = False + grad_enabled: bool = True + raw: Tuple[Any, ...] = () + raw_kwargs: Mapping[str, Any] = field(default_factory=dict) + + def evolve(self, **changes) -> "CallContext": + return replace(self, **changes) + + +class Spec: + """Composable, self-describing predicate over ``CallContext``.""" + + __slots__ = ("_fn", "_desc") + + def __init__(self, fn: Callable[[CallContext], bool], desc: str): + self._fn = fn + self._desc = desc + + def matches(self, ctx: CallContext) -> bool: + return bool(self._fn(ctx)) + + @property + def description(self) -> str: + return self._desc + + def __and__(self, other: "Spec") -> "Spec": + return Spec( + lambda ctx: self._fn(ctx) and other._fn(ctx), + f"({self._desc} and {other._desc})", + ) + + def __or__(self, other: "Spec") -> "Spec": + return Spec( + lambda ctx: self._fn(ctx) or other._fn(ctx), + f"({self._desc} or {other._desc})", + ) + + def __invert__(self) -> "Spec": + return Spec(lambda ctx: not self._fn(ctx), f"not({self._desc})") + + @classmethod + def always(cls) -> "Spec": + return cls(lambda ctx: True, "always") + + @classmethod + def of(cls, fn: Callable[[CallContext], bool], desc: str) -> "Spec": + return cls(fn, desc) + + @classmethod + def dtype_in(cls, *dtypes: torch.dtype) -> "Spec": + names = [str(getattr(d, "name", d)) for d in dtypes] + return cls( + lambda ctx: ctx.dtype in dtypes, f"dtype∈{{{','.join(names)}}}" + ) + + @classmethod + def ndim(cls, n: int) -> "Spec": + return cls(lambda ctx: ctx.ndim == n, f"ndim=={n}") + + @classmethod + def head_dim_in(cls, *dims: int) -> "Spec": + return cls( + lambda ctx: ctx.head_dim in dims, + f"head_dim∈{{{','.join(map(str, dims))}}}", + ) + + @classmethod + def fwd_in(cls, *modes: Optional[str]) -> "Spec": + names = [str(m) for m in modes] + return cls(lambda ctx: ctx.fwd in modes, f"fwd∈{{{','.join(names)}}}") + + @classmethod + def inference(cls) -> "Spec": + return cls(lambda ctx: ctx.fwd is not None, "inference") + + @classmethod + def training(cls) -> "Spec": + return cls(lambda ctx: ctx.fwd is None, "training") + + @classmethod + def no_grad(cls) -> "Spec": + return cls(lambda ctx: not ctx.grad_enabled, "no_grad") + + @classmethod + def has_cache(cls) -> "Spec": + return cls(lambda ctx: ctx.has_cache, "has_cache") + + @classmethod + def mask_free(cls) -> "Spec": + return cls(lambda ctx: not ctx.has_mask, "mask_free") + + @classmethod + def cuda_device(cls) -> "Spec": + return cls(lambda ctx: ctx.device_cuda, "cuda_device") + + +@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 + 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, + 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.""" + _FAMILIES[name] = OpFamily(name, 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) -> ImplRecord: + """Wrap an unregistered object; capability probes its own method.""" + supports = getattr(handle, "supports_call", None) + if supports is not None: + spec = Spec.of( + lambda ctx: bool(supports(*ctx.raw, **ctx.raw_kwargs)), + f"{type(handle).__name__}.supports_call", + ) + else: + spec = Spec.always() + return ImplRecord(family, type(handle).__name__, handle, spec) + + +def resolve( + family: str, ctx: CallContext, explicit: Optional[Any] = None +) -> Resolution: + """Resolve one family for one call (explicit-strict / implicit-loose).""" + fam = _family(family) + + 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) + 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(ctx): + 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(ctx): + return Resolution(record, "chain") + return Resolution(fam.fallback(), "fallback") + + +def resolve_plan(ctxs: Mapping[str, CallContext]) -> Dict[str, Resolution]: + """Resolve several families at once (one decision snapshot).""" + return {family: resolve(family, ctx) for family, ctx in ctxs.items()} + + +def _describe_ctx(ctx: CallContext) -> str: + parts = [f"fwd={ctx.fwd}", f"dtype={ctx.dtype}"] + if ctx.ndim is not None: + parts.append(f"ndim={ctx.ndim}") + if ctx.head_dim is not None: + parts.append(f"head_dim={ctx.head_dim}") + parts += [f"cache={ctx.has_cache}", f"mask={ctx.has_mask}", f"grad={ctx.grad_enabled}"] + return " ".join(parts) + + +def explain(family: str, ctx: CallContext, explicit: Optional[Any] = None) -> str: + """Human-readable decision trace for one family.""" + fam = _family(family) + records = sorted(fam.provider(), key=lambda r: r.priority) + lines = [f"[{family}] {_describe_ctx(ctx)}"] + 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(ctx): + lines.append(f" {record.name}: MATCH ({record.spec.description})") + else: + lines.append(f" {record.name}: reject ({record.spec.description})") + try: + resolution = resolve(family, ctx, explicit) + 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(ctxs: Mapping[str, CallContext]) -> str: + return "\n".join(explain(family, ctx) for family, ctx in ctxs.items()) diff --git a/tests/extension/test_dispatch.py b/tests/extension/test_dispatch.py new file mode 100644 index 0000000..7901999 --- /dev/null +++ b/tests/extension/test_dispatch.py @@ -0,0 +1,309 @@ +"""Tests for the generic operator dispatcher (Spec / decision tables).""" + +import importlib + +import pytest +import torch + +import astrai.extension.dispatch as dispatch +from astrai.extension import ( + ATTN_BACKEND, + CallContext, + ExplicitSelectionError, + ImplRecord, + Spec, + 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", + Spec.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", records, lambda: ImplRecord("toy", "beta", "beta-obj", Spec.always()) + ) + yield calls + dispatch._FAMILIES.pop("toy", None) + + +def ctx(family="toy", **kw): + base = dict(family=family, grad_enabled=False) + base.update(kw) + return CallContext(**base) + + + + +def test_spec_composition_and_description(): + spec = Spec.dtype_in(torch.bfloat16) & Spec.no_grad() + assert spec.matches(ctx(dtype=torch.bfloat16, grad_enabled=False)) + assert not spec.matches(ctx(dtype=torch.bfloat16, grad_enabled=True)) + assert "dtype" in spec.description and "no_grad" in spec.description + + either = Spec.training() | Spec.has_cache() + assert either.matches(ctx(fwd=None)) + assert either.matches(ctx(fwd="decode", has_cache=True)) + assert not either.matches(ctx(fwd="decode")) + + assert (~Spec.training()).matches(ctx(fwd="decode")) + + + + +def test_chain_returns_first_capable(toy_family): + assert resolve("toy", ctx(dtype=torch.bfloat16)).record.obj == "alpha-obj" + assert resolve("toy", ctx(dtype=torch.float32)).record.obj == "beta-obj" + + +def test_unfaithful_rows_are_chain_invisible(toy_family): + resolution = resolve("toy", ctx(dtype=torch.float32)) + assert resolution.record.obj == "beta-obj" + with op_backend(toy="fp8"): + assert resolve("toy", ctx(dtype=torch.float32)).record.obj == "fp8-obj" + + +def test_explicit_selection_is_strict(toy_family): + with pytest.raises(ExplicitSelectionError): + resolve("toy", ctx(dtype=torch.float32), explicit="alpha") + assert resolve("toy", ctx(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", ctx(dtype=torch.float32)) + assert resolve("toy", ctx(dtype=torch.bfloat16)).origin == "context" + + +def test_unknown_explicit_name_raises(toy_family): + with pytest.raises(ValueError, match="Unknown toy implementation"): + resolve("toy", ctx(), 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", ctx(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", ctx(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", ctx(dtype=torch.float32)).origin == "context" + assert resolve("toy", ctx(dtype=torch.bfloat16)).origin == "context" + + + + +def test_env_entry_is_soft(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "toy=alpha") + resolution = resolve("toy", ctx(dtype=torch.float32)) + assert resolution.record.obj == "beta-obj" and resolution.origin == "chain" + assert resolve("toy", ctx(dtype=torch.bfloat16)).origin == "env" + + +def test_env_unknown_impl_ignored(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "toy=missing") + assert resolve("toy", ctx(dtype=torch.float32)).record.obj == "beta-obj" + + +def test_env_profile_reference(toy_family, monkeypatch): + monkeypatch.setenv("ASTR_OPS", "profile=reference") + resolution = resolve("toy", ctx(dtype=torch.bfloat16)) + assert resolution.origin == "profile" + monkeypatch.setenv("ASTR_OPS", "toy=alpha,profile=reference") + assert resolve("toy", ctx(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", ctx(dtype=torch.float32)).origin == "context" + + + + +def test_resolve_plan_snapshots_families(toy_family): + plan = resolve_plan( + { + "toy": ctx(dtype=torch.bfloat16), + "rotary": ctx(family="rotary", device_cuda=True, grad_enabled=False), + } + ) + 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", ctx(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", ctx(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 + call_ctx = CallContext( + family="attention", + fwd=fwd, + dtype=q.dtype, + ndim=q.dim(), + head_dim=q.size(-1), + has_cache=has_cache, + has_mask=has_mask, + grad_enabled=False, + raw=(q, cache, mask, False, fwd), + ) + + cuda = attn_mod._instance(attn_mod.CudaBackend) + assert attn_mod._SPEC_CUDA.matches(call_ctx) == cuda.supports_call( + q, cache, mask, False, fwd + ) + + flash = attn_mod._instance(attn_mod.FlashAttnBackend) + assert attn_mod._SPEC_FLASH.matches(call_ctx) == 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) + call_ctx = attn_mod._context_from_call(q, None, None, True, "prefill") + resolution = resolve("attention", call_ctx) + 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(): + call_ctx = ctx( + family="rotary", + dtype=x.dtype, + device_cuda=True, + grad_enabled=False, + ) + resolution = resolve("rotary", call_ctx) + 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() + call_ctx = ctx(family="rotary", dtype=x.dtype, device_cuda=True, grad_enabled=True) + assert resolve("rotary", call_ctx).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)