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

@@ -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",
 ]

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
+    always wins.  ``ASTR_BACKEND`` (or ``ASTR_OPS``) is a process-wide
-    only when no context is set.  Pass ``use_default=False`` at request
+    override consulted only when no context is set.  Pass
-    submission to retain only an environment override or the caller's
+    ``use_default=False`` at request submission to retain only an
-    :func:`attn_backend` value.
+    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:
+    explicit = _resolve_backend(backend) if backend is not None else None
-        selected = _resolve_backend(backend)
+    resolution = _dispatch_resolve(
-        explicit = True
+        "attention", q, kv_cache, attn_mask, is_causal, fwd, explicit=explicit
-    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(
+    return resolution.record.obj.forward(
-            (
+        q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd
-                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)
 
 
 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")

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
+Registered rows: the fused CUDA kernel (bf16 CUDA, inference-only) and the
-CUDA kernel when available, falls back to torch complex multiply otherwise.
+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}
+_SPEC_CUDA = (
-
+    axis("device_cuda").truthy()
-
+    & axis("dtype").in_(torch.bfloat16)
-def _cuda_available() -> bool:
+    & axis("grad_enabled").eq(False)
-    if _cache["available"] is None:
+)
-        _cache["available"] = is_available("rotary_emb")
-    return _cache["available"]
 
 
 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 (
+    return resolve("rotary", x, freqs_cis).record.obj(x, freqs_cis)
-        _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)

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()
+    )

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)