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 "
2026-09-01 16:56:57 +08:00
parent aabf366633
commit b4d702cd14
5 changed files with 918 additions and 94 deletions
@@ -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)