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

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

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
+    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 +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:
+    ctx = _context_from_call(q, kv_cache, attn_mask, is_causal, fwd)
-        selected = _resolve_backend(backend)
+    explicit = _resolve_backend(backend) if backend is not None else None
-        explicit = True
+    resolution = _dispatch_resolve("attention", ctx, explicit=explicit)
-    else:
+    return resolution.record.obj.forward(
-        context_backend = _current_backend.get()
+        q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd
-        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)
 
 
 class AttentionBackend(ABC):
@@ -362,11 +368,11 @@ class AttentionBackend(ABC):
     """
 
     def __enter__(self) -> "AttentionBackend":
-        self._token = _current_backend.set(self)
+        self._token = set_override("attention", self)
         return self
 
     def __exit__(self, *exc) -> None:
-        _current_backend.reset(self._token)
+        reset_override(self._token)
 
     @classmethod
     @abstractmethod
@@ -817,3 +823,76 @@ class FlashAttnBackend(AttentionBackend):
             causal=True,
         )
         return out
+
+
+# Family registration over the generic dispatcher: the "attention" decision
+# table.  The Specs mirror each backend's ``supports_call`` exactly (a unit
+# test asserts they never drift).  The provider is re-evaluated per
+# resolution, so monkeypatching ``flash_attn_available`` (plus clearing
+# ``_priority_backends``) is honored, as before.
+
+_CLASS_TO_NAME: Dict[type, str] = {
+    CudaBackend: ATTN_BACKEND.CUDA.value,
+    FlashAttnBackend: ATTN_BACKEND.FLASH.value,
+    TorchNativeBackend: ATTN_BACKEND.TORCH_NATIVE.value,
+}
+
+_SPEC_CUDA = (
+    Spec.fwd_in("prefill", "decode")
+    & Spec.has_cache()
+    & Spec.ndim(3)
+    & Spec.dtype_in(torch.bfloat16)
+    & Spec.head_dim_in(*CudaBackend.HEAD_DIMS)
+    & Spec.of(
+        lambda ctx: is_available(f"attn_paged_{ctx.fwd}"), "paged kernels loaded"
+    )
+)
+
+_SPEC_FLASH = (
+    Spec.dtype_in(torch.float16, torch.bfloat16)
+    & Spec.of(lambda ctx: flash_attn_available(), "flash-attn available")
+    & (
+        (
+            Spec.inference()
+            & Spec.ndim(3)
+            & Spec.of(
+                lambda ctx: _flash_attn is not None
+                and hasattr(_flash_attn, "flash_attn_varlen_func"),
+                "varlen api present",
+            )
+        )
+        | (Spec.training() & Spec.mask_free())
+    )
+)
+
+
+def _attention_records() -> list:
+    specs = {
+        CudaBackend: _SPEC_CUDA,
+        FlashAttnBackend: _SPEC_FLASH,
+        TorchNativeBackend: Spec.always(),
+    }
+    return [
+        ImplRecord(
+            family="attention",
+            name=_CLASS_TO_NAME[type(backend)],
+            obj=backend,
+            spec=specs[type(backend)],
+            priority=position,
+        )
+        for position, backend in enumerate(_priority_backends())
+    ]
+
+
+def _reference_record() -> ImplRecord:
+    return ImplRecord(
+        family="attention",
+        name=ATTN_BACKEND.TORCH_NATIVE.value,
+        obj=_instance(TorchNativeBackend),
+        spec=Spec.always(),
+        priority=999,
+    )
+
+
+register_family("attention", _attention_records, _reference_record)
+register_env_alias("attention", "ASTR_BACKEND")

astrai/extension/backend/rotary.py

@@ -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,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}
+_SPEC_CUDA = Spec.cuda_device() & Spec.dtype_in(torch.bfloat16) & Spec.no_grad()
-
-
-def _cuda_available() -> bool:
-    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 +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 (
+    ctx = CallContext(
-        _cuda_available()
+        family="rotary",
-        and not torch.is_grad_enabled()
+        dtype=x.dtype,
-        and x.is_cuda
+        device_cuda=x.is_cuda,
-        and x.dtype == torch.bfloat16
+        grad_enabled=torch.is_grad_enabled(),
-    ):
+        raw=(x, freqs_cis),
-        return _cuda_rotary(x, freqs_cis)
+    )
-    return _torch_apply(x, freqs_cis)
+    return resolve("rotary", ctx).record.obj(x, freqs_cis)

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

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)