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