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/backend/rotary.py

@@ -1,7 +1,9 @@
-"""Rotary embedding with auto-dispatch to CUDA kernel.
+"""Rotary embedding dispatch (family "rotary").
 
-Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
-CUDA kernel when available, falls back to torch complex multiply otherwise.
+Registered rows: the fused CUDA kernel (bf16 CUDA, inference-only) and the
+torch complex-multiply fallback (autograd-safe).  Selection runs through
+the generic dispatcher, so ``op_backend(rotary=...)`` and
+``ASTR_OPS=rotary=torch`` work exactly like for attention.
 
 Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
 freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
@@ -10,16 +12,22 @@ freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
 import torch
 from torch import Tensor
 
+from astrai.extension.dispatch import (
+    ImplRecord,
+    Spec,
+    axis,
+    register_family,
+    resolve,
+    tensor_axes,
+)
 from astrai.extension.loader import is_available
 from astrai.extension.ops.rotary import rotary_emb as _cuda_rotary
 
-_cache = {"available": None}
-
-
-def _cuda_available() -> bool:
-    if _cache["available"] is None:
-        _cache["available"] = is_available("rotary_emb")
-    return _cache["available"]
+_SPEC_CUDA = (
+    axis("device_cuda").truthy()
+    & axis("dtype").in_(torch.bfloat16)
+    & axis("grad_enabled").eq(False)
+)
 
 
 def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
@@ -33,6 +41,34 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
     return x_out.to(dtype)
 
 
+def _rotary_records() -> list:
+    return [
+        ImplRecord(
+            family="rotary",
+            name="cuda",
+            obj=_cuda_rotary,
+            spec=_SPEC_CUDA,
+            available=lambda: is_available("rotary_emb"),
+            priority=0,
+        ),
+        ImplRecord(
+            family="rotary",
+            name="torch",
+            obj=_torch_apply,
+            spec=Spec.always(),
+            priority=99,
+        ),
+    ]
+
+
+register_family(
+    "rotary",
+    lambda x, freqs_cis: tensor_axes(x),
+    _rotary_records,
+    lambda: _rotary_records()[-1],
+)
+
+
 def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     """Apply rotary embedding to x.
 
@@ -43,11 +79,4 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     Returns:
         [batch, seq_len, n_heads, head_dim] (bf16)
     """
-    if (
-        _cuda_available()
-        and not torch.is_grad_enabled()
-        and x.is_cuda
-        and x.dtype == torch.bfloat16
-    ):
-        return _cuda_rotary(x, freqs_cis)
-    return _torch_apply(x, freqs_cis)
+    return resolve("rotary", x, freqs_cis).record.obj(x, freqs_cis)