ViperEkura
acbe57a0f9
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.
310 lines
9.7 KiB
Python
310 lines
9.7 KiB
Python
"""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)
|