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.
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user