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

astrai/trainer/train_callback.py

@@ -54,8 +54,11 @@ class TrainCallback(Protocol):
     def on_batch_end(self, context: TrainContext):
         """Called at the end of each batch."""
 
-    def on_optimizer_step(self, context: TrainContext):
+    def before_optimizer_step(self, context: TrainContext):
-        """Called on every optimizer step (sync step only)."""
+        """Called immediately before every optimizer step (sync step only)."""
+
+    def after_optimizer_step(self, context: TrainContext):
+        """Called after the optimizer and scheduler step (sync step only)."""
 
     def on_error(self, context: TrainContext):
         """Called when an error occurs during training."""
@@ -82,7 +85,7 @@ class GradientClippingCallback(TrainCallback):
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
 
-    def on_optimizer_step(self, context: TrainContext):
+    def before_optimizer_step(self, context: TrainContext):
         context.grad_norm = context.executor.clip_grad_norm(
             context.model, self.max_grad_norm
         )
@@ -170,7 +173,7 @@ class CheckpointCallback(TrainCallback):
                 )
                 context.checkpoint.save(save_path)
 
-    def on_batch_end(self, context: TrainContext):
+    def after_optimizer_step(self, context: TrainContext):
         if context.optimizer_step - self.last_ckpt_step >= self.interval:
             self._save_checkpoint(context)
 
@@ -216,7 +219,7 @@ class ProgressBarCallback(TrainCallback):
         )
 
     @only_on_rank(0)
-    def on_optimizer_step(self, context: TrainContext):
+    def before_optimizer_step(self, context: TrainContext):
         postfix = {
             "step": f"{context.optimizer_step:d}",
             "loss": f"{context.loss:.4f}",
@@ -343,7 +346,7 @@ class MetricCallback(TrainCallback):
             for log in self.log_cache:
                 f.write(json.dumps(log) + "\n")
 
-    def on_optimizer_step(self, context):
+    def before_optimizer_step(self, context):
         context.grad_snr_tracker.update(context.model)
 
         if (

astrai/trainer/trainer.py

@@ -93,7 +93,7 @@ class Trainer:
                         self._call_callbacks("on_batch_end", context)
 
                         if executor.sync_gradients:
-                            self._call_callbacks("on_optimizer_step", context)
+                            self._call_callbacks("before_optimizer_step", context)
                             context.optimizer.step()
                             context.strategy.on_optimizer_step()
                             context.optimizer.zero_grad()
@@ -101,6 +101,8 @@ class Trainer:
                             if context.scheduler:
                                 context.scheduler.step()
 
+                            self._call_callbacks("after_optimizer_step", context)
+
                 self._call_callbacks("on_epoch_end", context)
 
             if context.stop_requested:

csrc/CMakeLists.txt

@@ -68,7 +68,7 @@ set(KERNEL_SRCS
     attention/prefill.cu
     attention/paged_decode.cu
     attention/paged_prefill.cu
-    rotary/rotary_emb.cu
+    rotary_emb.cu
 )
 
 if(ASTRAI_CUDA_ARCH GREATER_EQUAL 89)
@@ -90,6 +90,7 @@ foreach(i RANGE ${_kernel_last})
     target_compile_definitions(${name} PRIVATE TORCH_EXTENSION_NAME=${name})
 
     target_include_directories(${name} PRIVATE
+        "${CMAKE_CURRENT_SOURCE_DIR}/kernels"
         "${TORCH_HOME}/include"
         "${TORCH_HOME}/include/torch/csrc/api/include"
         "${PYTHON_INCLUDE_DIR}")

csrc/kernels/attention/decode_split_kv.cuh

@@ -2,8 +2,8 @@
 #include <cuda_bf16.h>
 #include <float.h>
 #include "common.h"
+#include "common/reduce.cuh"
 #include "layout_policies.cuh"
-#include "../common/reduce.cuh"
 
 namespace astrai {
 namespace attention {

csrc/kernels/attention/mma_utils.cuh

@@ -3,8 +3,8 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
-#include "../common/cp_async.cuh"
+#include "common/cp_async.cuh"
-#include "../common/mma.cuh"
+#include "common/mma.cuh"
 
 // Predicated cp.async (4-operand form) requires CUDA 11.2+.
 // bf16 mma.sync requires sm_80+ (guarded at build time by ASTRAI_NO_MMA).

csrc/kernels/attention/prefill_split_q.cuh

@@ -2,8 +2,8 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "common.h"
+#include "common/reduce.cuh"
 #include "layout_policies.cuh"
-#include "../common/reduce.cuh"
 
 namespace astrai {
 namespace attention {

csrc/kernels/fp8/gemm.cuh

@@ -11,8 +11,8 @@
 #include <cuda_runtime.h>
 #include <type_traits>
 
-#include "../common/cp_async.cuh"
 #include "common.h"
+#include "common/cp_async.cuh"
 #include "gemm/epilogue.cuh"
 #include "gemm/load.cuh"
 #include "gemm/mainloop.cuh"

csrc/kernels/fp8/gemm/epilogue.cuh

@@ -2,7 +2,7 @@
 // Collective epilogue: fused bias, the bf16 scatter of the fp32 accumulators
 // through the reclaimed operand shared memory, and the coalesced copy-out.
 
-#include "../common.h"
+#include "fp8/common.h"
 #include "policy.cuh"
 
 namespace astrai {

csrc/kernels/fp8/gemm/load.cuh

@@ -5,8 +5,8 @@
 // The staging invariants and the swizzle derivation live in
 // docs/developer/cuda_kernels.md.
 
-#include "../../common/cp_async.cuh"
+#include "common/cp_async.cuh"
-#include "../common.h"
+#include "fp8/common.h"
 #include "policy.cuh"
 
 namespace astrai {

csrc/kernels/fp8/gemm/mainloop.cuh

@@ -7,8 +7,8 @@
 
 #include <type_traits>
 
-#include "../../common/mma.cuh"
+#include "common/mma.cuh"
-#include "../common.h"
+#include "fp8/common.h"
 #include "load.cuh"
 #include "policy.cuh"
 

csrc/kernels/fp8/gemm/policy.cuh

@@ -5,7 +5,7 @@
 
 #include <type_traits>
 
-#include "../common.h"
+#include "fp8/common.h"
 
 namespace astrai {
 namespace fp8 {

csrc/kernels/fp8/ops.cu

@@ -8,7 +8,7 @@
 #include <mutex>
 #include <unordered_map>
 
-#include "../common/device.cuh"
+#include "common/device.cuh"
 #include "gemm.cuh"
 #include "quantize.cuh"
 

csrc/kernels/fp8/quantize.cuh

@@ -10,7 +10,7 @@
 #include <cstdint>
 
 #include "common.h"
-#include "../common/reduce.cuh"
+#include "common/reduce.cuh"
 
 namespace astrai {
 namespace fp8 {

csrc/kernels/rotary_emb.cu

@@ -11,34 +11,41 @@ __global__ void rotary_emb_kernel(
     int n_heads,
     int head_dim
 ) {
-    const int half_dim = head_dim >> 1;
+    // Each head tiles into exact 2-pair chunks: one 8B x access and one 16B
-    const int total = n_tokens * n_heads * half_dim;
+    // cos/sin access per chunk (head_dim % 4 == 0 is enforced on the host).
+    const int chunks = head_dim >> 2;
+    const int total = n_tokens * n_heads * chunks;
 
-    for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    for (int c = blockIdx.x * blockDim.x + threadIdx.x;
-         idx < total;
+         c < total;
-         idx += gridDim.x * blockDim.x) {
+         c += gridDim.x * blockDim.x) {
+        const int chunk = c % chunks;
+        const int tmp = c / chunks;
+        const int head = tmp % n_heads;
+        const int token = tmp / n_heads;
 
-        int pair = idx % half_dim;
+        const int x_off = (tmp * head_dim) + (chunk << 2);
-        int tmp = idx / half_dim;
+        const int f_off = ((token * chunks) + chunk) << 2;
-        int head = tmp % n_heads;
-        tmp /= n_heads;
-        int token = tmp;
 
-        int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
+        const float4 f = *reinterpret_cast<const float4*>(freqs_cis + f_off);
-        int cs_offset = (token * half_dim + pair) * 2;
+        const uint2 xr = *reinterpret_cast<const uint2*>(x + x_off);
+        __nv_bfloat162 p0 = *reinterpret_cast<const __nv_bfloat162*>(&xr.x);
+        __nv_bfloat162 p1 = *reinterpret_cast<const __nv_bfloat162*>(&xr.y);
 
-        __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
+        const float e0 = __bfloat162float(__low2bfloat16(p0));
-        float x_even = __bfloat162float(__low2bfloat16(x_pair));
+        const float o0 = __bfloat162float(__high2bfloat16(p0));
-        float x_odd  = __bfloat162float(__high2bfloat16(x_pair));
+        const float e1 = __bfloat162float(__low2bfloat16(p1));
+        const float o1 = __bfloat162float(__high2bfloat16(p1));
 
-        float c = freqs_cis[cs_offset];
+        __nv_bfloat162 r0 = __floats2bfloat162_rn(
-        float s = freqs_cis[cs_offset + 1];
+            e0 * f.x - o0 * f.y, e0 * f.y + o0 * f.x);
+        __nv_bfloat162 r1 = __floats2bfloat162_rn(
+            e1 * f.z - o1 * f.w, e1 * f.w + o1 * f.z);
 
-        float out_even = x_even * c - x_odd * s;
+        uint2 oraw;
-        float out_odd  = x_even * s + x_odd * c;
+        *reinterpret_cast<__nv_bfloat162*>(&oraw.x) = r0;
-
+        *reinterpret_cast<__nv_bfloat162*>(&oraw.y) = r1;
-        __nv_bfloat162 out_pair = __floats2bfloat162_rn(out_even, out_odd);
+        *reinterpret_cast<uint2*>(out + x_off) = oraw;
-        *reinterpret_cast<__nv_bfloat162*>(out + x_offset) = out_pair;
     }
 }
 
@@ -64,7 +71,7 @@ torch::Tensor rotary_emb(
     int n_heads = x.size(x.dim() - 2);
     int head_dim = x.size(x.dim() - 1);
 
-    TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
+    TORCH_CHECK(head_dim % 4 == 0, "head_dim must be a multiple of 4");
     TORCH_CHECK(freqs_cis.numel() == (int64_t)n_tokens * head_dim,
                 "freqs_cis token or rotary dimension mismatch");
     TORCH_CHECK(freqs_cis.size(-2) == head_dim / 2, "freqs_cis dim/2 mismatch");
@@ -72,10 +79,9 @@ torch::Tensor rotary_emb(
 
     auto out = torch::empty_like(x);
 
-    int half_dim = head_dim / 2;
+    int work = n_tokens * n_heads * (head_dim / 4);
-    int total = n_tokens * n_heads * half_dim;
     int block = 256;
-    int grid = std::min((total + block - 1) / block, 1024);
+    int grid = std::min((work + block - 1) / block, 2048);
 
     rotary_emb_kernel<<<grid, block, 0, stream>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),

csrc/tests/attn_paged_test.cu

@@ -1,5 +1,5 @@
 // Compile:
-//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
+//   nvcc -I csrc/kernels -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
 //        --extra-device-vectorization -Xcompiler -fopenmp \
 //        csrc/tests/attn_paged_test.cu \
 //        -o /tmp/test_paged && /tmp/test_paged
@@ -7,7 +7,7 @@
 #include <cstring>
 #include <vector>
 #include "test_utils.cuh"
-#include "../kernels/attention/dispatchers.cuh"
+#include "attention/dispatchers.cuh"
 
 using namespace astrai::attention;
 

csrc/tests/attn_test.cu

@@ -1,13 +1,13 @@
 /*
 Pure-C test — uses shared dispatcher.  Combines the decode (split-KV) and
 prefill (split-Q) correctness checks + benchmarks into one binary.
-nvcc -I csrc -arch=sm_89 -O3 \
+nvcc -I csrc/kernels -arch=sm_89 -O3 \
     --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
     -Xcompiler -fopenmp csrc/tests/attn_test.cu -o test && ./test
 */
 
 #include "test_utils.cuh"
-#include "../kernels/attention/dispatchers.cuh"
+#include "attention/dispatchers.cuh"
 
 using namespace astrai::attention;
 

csrc/tests/fp8_test.cu

@@ -6,7 +6,7 @@ Part 1 exercises one bf16 -> fp8 -> mma.sync m16n8k32 instruction pair
 Part 2 checks launch_fp8_gemm across all four operand layouts, both K
 tiles, and ragged shapes against an fp32 CPU reference.
 
-nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
+nvcc -I csrc/kernels -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
     && /tmp/fp8_test
 */
 
@@ -21,8 +21,8 @@ nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test
 #include <type_traits>
 #include <vector>
 
-#include "../kernels/common/mma.cuh"
+#include "common/mma.cuh"
-#include "../kernels/fp8/gemm.cuh"
+#include "fp8/gemm.cuh"
 
 using namespace astrai::fp8;
 

docs/developer/architecture.md

@@ -746,13 +746,14 @@ classDiagram
             +on_epoch_end(context)
             +on_batch_begin(context)
             +on_batch_end(context)
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
+            +after_optimizer_step(context)
             +on_error(context)
         }
 
         class GradientClippingCallback {
             +Optional[float] max_grad_norm
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
         }
 
         class GradientCheckpointingCallback {
@@ -767,7 +768,7 @@ classDiagram
             +bool weight_only
             +Callable save_extra_fn
             -_save_checkpoint(context)
-            +on_batch_end(context)
+            +after_optimizer_step(context)
             +on_train_end(context)
             +on_error(context)
             +save_extra(context) dict
@@ -779,7 +780,7 @@ classDiagram
             +IO file
             +tqdm progress_bar
             +on_epoch_begin(context)
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
             +on_epoch_end(context)
         }
 
@@ -788,7 +789,7 @@ classDiagram
             +int save_interval
             +List[str] metrics
             +int val_step
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
             +on_epoch_end(context)
             +on_train_end(context)
             +on_error(context)

docs/developer/cuda_kernels.md

@@ -10,7 +10,7 @@ AstrAI includes optional custom CUDA kernels for attention, rotary embedding, an
 | `attn_prefill` | `attention/prefill.cu` | GQA prefill attention (split-Q) |
 | `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention |
 | `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
-| `rotary_emb` | `rotary/rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
+| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
 | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
 
 Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
@@ -27,7 +27,7 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac
 
 ### Rotary Embedding Kernel
 
-The `rotary_emb` kernel (`csrc/kernels/rotary/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
+The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
 
 - One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store
 - f32 cos/sin input, bf16 compute and output
@@ -407,7 +407,7 @@ blocks.
 Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment. Example:
 
 ```bash
-nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
+nvcc -I csrc/kernels -arch=sm_89 -O3 --use_fast_math \
      --ptxas-options=-O3,-v --extra-device-vectorization \
      -Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
 ```
@@ -423,7 +423,7 @@ Hardware: NVIDIA L20 (sm_89, 46 GB), CUDA 12.8, driver 570.86.
 
 Reproduce (decode + prefill in `attn_test.cu`, paged in `attn_paged_test.cu`):
 ```bash
-nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
+nvcc -I csrc/kernels -arch=sm_89 -O3 --use_fast_math \
      --ptxas-options=-O3,-v --extra-device-vectorization \
      -Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
 ```
@@ -460,8 +460,7 @@ csrc/
 │   │   ├── prefill.cu                #   → module attn_prefill
 │   │   ├── paged_decode.cu           #   → module attn_paged_decode
 │   │   └── paged_prefill.cu          #   → module attn_paged_prefill
-│   ├── rotary/
+│   ├── rotary_emb.cu                  # rotary embedding (kernel + binding in one file) → module rotary_emb
-│   │   └── rotary_emb.cu             # rotary embedding (kernel + binding in one file) → module rotary_emb
 │   └── fp8/                          # FP8 family (module name fp8_ops)
 │       ├── common.h                  #   FP8Format enum, Fp8GemmTraits, FP8Params / FP8QuantizeParams PODs, layout tags (no torch)
 │       ├── quantize.cuh              #   quantize kernels: vectorized + 32×32-tile transpose (out_layout 0/1/2) (no torch)

docs/developer/internals.md

@@ -118,12 +118,13 @@ on_train_begin
         on_batch_end
 
         if executor.sync_gradients:
-          on_optimizer_step
+          before_optimizer_step
           optimizer.step()
           strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
+          after_optimizer_step
     on_epoch_end
 on_train_end
 ```
@@ -139,8 +140,9 @@ Strategy metrics are detached and converted to Python `float` values before the
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
+| `before_optimizer_step` | Every accumulation window, before `optimizer.step()` | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback` |
+| `on_batch_end` | Every batch | — |
+| `after_optimizer_step` | Every accumulation window, after `optimizer.step()` and `scheduler.step()` | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
 | `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |

docs/guides/training.md

@@ -70,12 +70,13 @@ on_train_begin
         on_batch_end
 
         if executor.sync_gradients:
-          on_optimizer_step
+          before_optimizer_step
           optimizer.step()
           strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
+          after_optimizer_step
     on_epoch_end
 on_train_end
 ```
@@ -87,8 +88,9 @@ on_train_end
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
+| `before_optimizer_step` | Every accumulation window, before `optimizer.step()` | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback` |
+| `on_batch_end` | Every batch | — |
+| `after_optimizer_step` | Every accumulation window, after `optimizer.step()` and `scheduler.step()` | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
 | `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |

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)

tests/trainer/test_callbacks.py

@@ -1,8 +1,12 @@
+from pathlib import Path
+
 import torch
 
 from astrai.model.components.decoder_block import DecoderBlock
+from astrai.serialization import Checkpoint
 from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
 from astrai.trainer.trainer import Trainer
+from tests.helpers import RandomTokenDataset
 
 
 def test_gradient_checkpointing_enable_disable(test_model):
@@ -135,3 +139,34 @@ def test_callback_integration(
     assert "on_train_begin" in callback_calls
     assert "on_batch_end" in callback_calls
     assert "on_epoch_end" in callback_calls
+
+
+def test_checkpoint_captures_completed_optimizer_step(
+    base_test_env, train_config_factory, device
+):
+    """Checkpoint state must include the update represented by its step number."""
+    model = base_test_env["model"]
+    initial_state = {
+        name: tensor.detach().cpu().clone()
+        for name, tensor in model.state_dict().items()
+    }
+    train_config = train_config_factory(
+        model_fn=lambda: model,
+        dataset=RandomTokenDataset(length=2),
+        test_dir=base_test_env["test_dir"],
+        device=device,
+        batch_per_device=2,
+        ckpt_interval=1,
+    )
+
+    Trainer(train_config).train()
+
+    checkpoint = Checkpoint.load(
+        str(Path(base_test_env["test_dir"]) / "epoch_0_step_1")
+    )
+    assert any(
+        not torch.equal(checkpoint.state_dict[name].cpu(), initial_tensor)
+        for name, initial_tensor in initial_state.items()
+    )
+    assert checkpoint.extra["optimizer"]["state"]
+    assert checkpoint.extra["scheduler"]["last_epoch"] == 1