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:
@@ -27,6 +27,20 @@ from astrai.extension.backend import (
|
||||
attn_backend,
|
||||
get_backend,
|
||||
)
|
||||
from astrai.extension.dispatch import (
|
||||
CallContext,
|
||||
ExplicitSelectionError,
|
||||
ImplRecord,
|
||||
Resolution,
|
||||
Spec,
|
||||
explain,
|
||||
explain_plan,
|
||||
op_backend,
|
||||
register_env_alias,
|
||||
register_family,
|
||||
resolve,
|
||||
resolve_plan,
|
||||
)
|
||||
from astrai.extension.loader import KERNEL_NAMES, is_available
|
||||
from astrai.extension.ops import (
|
||||
TensorLayout,
|
||||
@@ -52,4 +66,16 @@ __all__ = [
|
||||
"is_available",
|
||||
"KERNEL_NAMES",
|
||||
"apply_rotary_emb",
|
||||
"CallContext",
|
||||
"ExplicitSelectionError",
|
||||
"ImplRecord",
|
||||
"Resolution",
|
||||
"Spec",
|
||||
"explain",
|
||||
"explain_plan",
|
||||
"op_backend",
|
||||
"register_env_alias",
|
||||
"register_family",
|
||||
"resolve",
|
||||
"resolve_plan",
|
||||
]
|
||||
|
||||
@@ -21,10 +21,15 @@ Usage — mirroring ``torch.nn.attention.sdpa_kernel``:
|
||||
...
|
||||
|
||||
Thread-safe via ``contextvars`` — each scheduler thread gets its own
|
||||
active backend. Backend resolution follows a strict precedence:
|
||||
active backend. Backend resolution is a thin facade over the generic
|
||||
operator dispatcher (``astrai.extension.dispatch``): the three backends
|
||||
are registered as the "attention" family and the decision table lives in
|
||||
``_attention_records``. Resolution follows a strict precedence:
|
||||
|
||||
1. explicit ``attn_backend(...)`` context (wins over everything),
|
||||
2. the process-wide ``ASTR_BACKEND`` environment override,
|
||||
2. the process-wide ``ASTR_BACKEND`` environment override
|
||||
(or an ``ASTR_OPS`` ``attention=`` entry, which wins over the legacy
|
||||
variable),
|
||||
3. an implicit default picked from the available backends
|
||||
(cuda > flash > torch).
|
||||
|
||||
@@ -40,12 +45,9 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
|
||||
(blhd). The backend returns ``[batch, seq_len, n_heads * head_dim]``.
|
||||
"""
|
||||
|
||||
import contextvars
|
||||
import enum
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
|
||||
@@ -54,6 +56,18 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.extension.dispatch import (
|
||||
CallContext,
|
||||
ImplRecord,
|
||||
Spec,
|
||||
env_selection,
|
||||
get_override,
|
||||
register_env_alias,
|
||||
register_family,
|
||||
resolve as _dispatch_resolve,
|
||||
reset_override,
|
||||
set_override,
|
||||
)
|
||||
from astrai.extension.loader import is_available
|
||||
from astrai.extension.ops.attention import (
|
||||
attn_paged_decode,
|
||||
@@ -72,15 +86,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_default_backend_lock = threading.Lock()
|
||||
_env_backend_name: Optional[str] = None
|
||||
_env_backend: Optional["AttentionBackend"] = None
|
||||
_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
|
||||
contextvars.ContextVar("attn_backend", default=None)
|
||||
)
|
||||
|
||||
# Backends are stateless — one canonical instance per class, created lazily
|
||||
# and reused everywhere (resolution, fallback, context managers).
|
||||
_singletons: Dict[type, "AttentionBackend"] = {}
|
||||
|
||||
|
||||
@@ -157,28 +162,6 @@ def _resolve_default_backend() -> "AttentionBackend":
|
||||
return _priority_backends()[0]
|
||||
|
||||
|
||||
def _environment_backend() -> Optional["AttentionBackend"]:
|
||||
"""Resolve the process-wide ``ASTR_BACKEND`` override, if configured."""
|
||||
global _env_backend, _env_backend_name
|
||||
name = os.environ.get("ASTR_BACKEND", "").strip().lower()
|
||||
if not name:
|
||||
return None
|
||||
if name != _env_backend_name:
|
||||
with _default_backend_lock:
|
||||
if name != _env_backend_name:
|
||||
try:
|
||||
_env_backend = _resolve_backend(name)
|
||||
except (ValueError, RuntimeError):
|
||||
_env_backend = None
|
||||
logger.warning(
|
||||
"ASTR_BACKEND=%r is not a registered attention backend; "
|
||||
"falling back to default resolution",
|
||||
name,
|
||||
)
|
||||
_env_backend_name = name
|
||||
return _env_backend
|
||||
|
||||
|
||||
def _resolve_backend(
|
||||
backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
|
||||
) -> "AttentionBackend":
|
||||
@@ -204,18 +187,44 @@ def _resolve_backend(
|
||||
return _resolve_default_backend()
|
||||
|
||||
|
||||
_ENV_WARNED: set = set()
|
||||
|
||||
|
||||
def _environment_backend() -> Optional["AttentionBackend"]:
|
||||
"""Resolve the process-wide env override (``ASTR_OPS`` or the legacy
|
||||
``ASTR_BACKEND``) to a backend instance, if it names a registered one.
|
||||
|
||||
Invalid names warn once and are ignored, falling back to default
|
||||
resolution — the override is soft, never fatal.
|
||||
"""
|
||||
name = env_selection("attention")
|
||||
if name is None:
|
||||
return None
|
||||
try:
|
||||
return _resolve_backend(name)
|
||||
except (ValueError, RuntimeError):
|
||||
message = (
|
||||
f"ASTR_BACKEND/ASTR_OPS value {name!r} is not a registered "
|
||||
f"attention backend; falling back to default resolution"
|
||||
)
|
||||
if message not in _ENV_WARNED:
|
||||
_ENV_WARNED.add(message)
|
||||
logger.warning(message)
|
||||
return None
|
||||
|
||||
|
||||
def get_backend(
|
||||
use_default: bool = True,
|
||||
) -> Optional["AttentionBackend"]:
|
||||
"""Resolve the active backend: explicit context > env > default.
|
||||
|
||||
An ``attn_backend(...)`` context is the caller's explicit choice and
|
||||
always wins. ``ASTR_BACKEND`` is a process-wide override consulted
|
||||
only when no context is set. Pass ``use_default=False`` at request
|
||||
submission to retain only an environment override or the caller's
|
||||
:func:`attn_backend` value.
|
||||
always wins. ``ASTR_BACKEND`` (or ``ASTR_OPS``) is a process-wide
|
||||
override consulted only when no context is set. Pass
|
||||
``use_default=False`` at request submission to retain only an
|
||||
environment override or the caller's :func:`attn_backend` value.
|
||||
"""
|
||||
context_backend = _current_backend.get()
|
||||
context_backend = get_override("attention")
|
||||
if context_backend is not None:
|
||||
return context_backend
|
||||
env_backend = _environment_backend()
|
||||
@@ -241,11 +250,11 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
|
||||
...
|
||||
"""
|
||||
instance = _resolve_backend(backend)
|
||||
token = _current_backend.set(instance)
|
||||
token = set_override("attention", instance)
|
||||
try:
|
||||
yield instance
|
||||
finally:
|
||||
_current_backend.reset(token)
|
||||
reset_override(token)
|
||||
|
||||
|
||||
def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
|
||||
@@ -260,6 +269,28 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
|
||||
)
|
||||
|
||||
|
||||
def _context_from_call(
|
||||
q: Tensor,
|
||||
kv_cache: Optional["KVCache"],
|
||||
attn_mask: Optional[Tensor],
|
||||
is_causal: bool,
|
||||
fwd: Optional[str],
|
||||
) -> CallContext:
|
||||
"""Snapshot the axes the attention decision table depends on."""
|
||||
return CallContext(
|
||||
family="attention",
|
||||
fwd=fwd,
|
||||
dtype=q.dtype,
|
||||
device_cuda=q.is_cuda,
|
||||
ndim=q.dim(),
|
||||
head_dim=q.size(-1) if q.dim() >= 1 else None,
|
||||
has_cache=kv_cache is not None,
|
||||
has_mask=attn_mask is not None,
|
||||
grad_enabled=torch.is_grad_enabled(),
|
||||
raw=(q, kv_cache, attn_mask, is_causal, fwd),
|
||||
)
|
||||
|
||||
|
||||
def attention(
|
||||
q: Tensor,
|
||||
k: Tensor,
|
||||
@@ -300,37 +331,12 @@ def attention(
|
||||
Returns:
|
||||
[batch, q_len, n_heads * head_dim]
|
||||
"""
|
||||
if backend is not None:
|
||||
selected = _resolve_backend(backend)
|
||||
explicit = True
|
||||
else:
|
||||
context_backend = _current_backend.get()
|
||||
explicit = context_backend is not None
|
||||
# Resolve through the same chain as inference: explicit context >
|
||||
# ASTR_BACKEND env > default. Training calls (fwd=None, no cache)
|
||||
# land on the CUDA backend and fall back by capability below —
|
||||
# flash when it can handle the call, else torch SDPA.
|
||||
selected = get_backend()
|
||||
assert selected is not None
|
||||
|
||||
if not selected.supports_call(q, kv_cache, attn_mask, is_causal, fwd):
|
||||
if explicit:
|
||||
raise RuntimeError(
|
||||
f"Explicitly-set backend {type(selected).__name__} cannot "
|
||||
f"handle this attention call (shape={q.shape}, "
|
||||
f"dtype={q.dtype}, kv_cache={'none' if kv_cache is None else 'present'}, "
|
||||
f"attn_mask={'none' if attn_mask is None else 'present'}). "
|
||||
f"Remove the attn_backend() context or switch to a compatible backend."
|
||||
)
|
||||
selected = next(
|
||||
(
|
||||
candidate
|
||||
for candidate in _priority_backends()
|
||||
if candidate.supports_call(q, kv_cache, attn_mask, is_causal, fwd)
|
||||
),
|
||||
_instance(TorchNativeBackend),
|
||||
)
|
||||
return selected.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
|
||||
ctx = _context_from_call(q, kv_cache, attn_mask, is_causal, fwd)
|
||||
explicit = _resolve_backend(backend) if backend is not None else None
|
||||
resolution = _dispatch_resolve("attention", ctx, explicit=explicit)
|
||||
return resolution.record.obj.forward(
|
||||
q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd
|
||||
)
|
||||
|
||||
|
||||
class AttentionBackend(ABC):
|
||||
@@ -362,11 +368,11 @@ class AttentionBackend(ABC):
|
||||
"""
|
||||
|
||||
def __enter__(self) -> "AttentionBackend":
|
||||
self._token = _current_backend.set(self)
|
||||
self._token = set_override("attention", self)
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> None:
|
||||
_current_backend.reset(self._token)
|
||||
reset_override(self._token)
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
@@ -817,3 +823,76 @@ class FlashAttnBackend(AttentionBackend):
|
||||
causal=True,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# Family registration over the generic dispatcher: the "attention" decision
|
||||
# table. The Specs mirror each backend's ``supports_call`` exactly (a unit
|
||||
# test asserts they never drift). The provider is re-evaluated per
|
||||
# resolution, so monkeypatching ``flash_attn_available`` (plus clearing
|
||||
# ``_priority_backends``) is honored, as before.
|
||||
|
||||
_CLASS_TO_NAME: Dict[type, str] = {
|
||||
CudaBackend: ATTN_BACKEND.CUDA.value,
|
||||
FlashAttnBackend: ATTN_BACKEND.FLASH.value,
|
||||
TorchNativeBackend: ATTN_BACKEND.TORCH_NATIVE.value,
|
||||
}
|
||||
|
||||
_SPEC_CUDA = (
|
||||
Spec.fwd_in("prefill", "decode")
|
||||
& Spec.has_cache()
|
||||
& Spec.ndim(3)
|
||||
& Spec.dtype_in(torch.bfloat16)
|
||||
& Spec.head_dim_in(*CudaBackend.HEAD_DIMS)
|
||||
& Spec.of(
|
||||
lambda ctx: is_available(f"attn_paged_{ctx.fwd}"), "paged kernels loaded"
|
||||
)
|
||||
)
|
||||
|
||||
_SPEC_FLASH = (
|
||||
Spec.dtype_in(torch.float16, torch.bfloat16)
|
||||
& Spec.of(lambda ctx: flash_attn_available(), "flash-attn available")
|
||||
& (
|
||||
(
|
||||
Spec.inference()
|
||||
& Spec.ndim(3)
|
||||
& Spec.of(
|
||||
lambda ctx: _flash_attn is not None
|
||||
and hasattr(_flash_attn, "flash_attn_varlen_func"),
|
||||
"varlen api present",
|
||||
)
|
||||
)
|
||||
| (Spec.training() & Spec.mask_free())
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _attention_records() -> list:
|
||||
specs = {
|
||||
CudaBackend: _SPEC_CUDA,
|
||||
FlashAttnBackend: _SPEC_FLASH,
|
||||
TorchNativeBackend: Spec.always(),
|
||||
}
|
||||
return [
|
||||
ImplRecord(
|
||||
family="attention",
|
||||
name=_CLASS_TO_NAME[type(backend)],
|
||||
obj=backend,
|
||||
spec=specs[type(backend)],
|
||||
priority=position,
|
||||
)
|
||||
for position, backend in enumerate(_priority_backends())
|
||||
]
|
||||
|
||||
|
||||
def _reference_record() -> ImplRecord:
|
||||
return ImplRecord(
|
||||
family="attention",
|
||||
name=ATTN_BACKEND.TORCH_NATIVE.value,
|
||||
obj=_instance(TorchNativeBackend),
|
||||
spec=Spec.always(),
|
||||
priority=999,
|
||||
)
|
||||
|
||||
|
||||
register_family("attention", _attention_records, _reference_record)
|
||||
register_env_alias("attention", "ASTR_BACKEND")
|
||||
|
||||
@@ -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,17 @@ freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.extension.dispatch import (
|
||||
CallContext,
|
||||
ImplRecord,
|
||||
Spec,
|
||||
register_family,
|
||||
resolve,
|
||||
)
|
||||
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 = Spec.cuda_device() & Spec.dtype_in(torch.bfloat16) & Spec.no_grad()
|
||||
|
||||
|
||||
def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
||||
@@ -33,6 +36,29 @@ 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", _rotary_records, lambda: _rotary_records()[-1])
|
||||
|
||||
|
||||
def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
||||
"""Apply rotary embedding to x.
|
||||
|
||||
@@ -43,11 +69,11 @@ 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)
|
||||
ctx = CallContext(
|
||||
family="rotary",
|
||||
dtype=x.dtype,
|
||||
device_cuda=x.is_cuda,
|
||||
grad_enabled=torch.is_grad_enabled(),
|
||||
raw=(x, freqs_cis),
|
||||
)
|
||||
return resolve("rotary", ctx).record.obj(x, freqs_cis)
|
||||
|
||||
@@ -0,0 +1,388 @@
|
||||
"""Operator dispatch: one selection mechanism for all op families.
|
||||
|
||||
A family registers an ordered list of ``ImplRecord`` rows (name, impl
|
||||
object, capability ``Spec``, machine-level ``available``). Resolution:
|
||||
explicit/context selection (strict — raises when incapable) > ``ASTR_OPS``
|
||||
env entry (soft — falls through) > first capable row > family fallback.
|
||||
``Spec`` is a composable predicate over ``CallContext``; the rows are the
|
||||
family's decision table, printable via ``explain``.
|
||||
|
||||
Records flagged ``faithful=False`` change numerics (e.g. fp8) and are only
|
||||
reachable through an explicit selection, never the implicit chain.
|
||||
"""
|
||||
|
||||
import contextvars
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CallContext:
|
||||
"""Axes a decision may depend on, snapshotted per call.
|
||||
|
||||
``raw``/``raw_kwargs`` keep the original call so unregistered handles
|
||||
can still be probed through their own capability method.
|
||||
"""
|
||||
|
||||
family: str
|
||||
fwd: Optional[str] = None
|
||||
dtype: Optional[torch.dtype] = None
|
||||
device_cuda: bool = False
|
||||
ndim: Optional[int] = None
|
||||
head_dim: Optional[int] = None
|
||||
has_cache: bool = False
|
||||
has_mask: bool = False
|
||||
grad_enabled: bool = True
|
||||
raw: Tuple[Any, ...] = ()
|
||||
raw_kwargs: Mapping[str, Any] = field(default_factory=dict)
|
||||
|
||||
def evolve(self, **changes) -> "CallContext":
|
||||
return replace(self, **changes)
|
||||
|
||||
|
||||
class Spec:
|
||||
"""Composable, self-describing predicate over ``CallContext``."""
|
||||
|
||||
__slots__ = ("_fn", "_desc")
|
||||
|
||||
def __init__(self, fn: Callable[[CallContext], bool], desc: str):
|
||||
self._fn = fn
|
||||
self._desc = desc
|
||||
|
||||
def matches(self, ctx: CallContext) -> bool:
|
||||
return bool(self._fn(ctx))
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return self._desc
|
||||
|
||||
def __and__(self, other: "Spec") -> "Spec":
|
||||
return Spec(
|
||||
lambda ctx: self._fn(ctx) and other._fn(ctx),
|
||||
f"({self._desc} and {other._desc})",
|
||||
)
|
||||
|
||||
def __or__(self, other: "Spec") -> "Spec":
|
||||
return Spec(
|
||||
lambda ctx: self._fn(ctx) or other._fn(ctx),
|
||||
f"({self._desc} or {other._desc})",
|
||||
)
|
||||
|
||||
def __invert__(self) -> "Spec":
|
||||
return Spec(lambda ctx: not self._fn(ctx), f"not({self._desc})")
|
||||
|
||||
@classmethod
|
||||
def always(cls) -> "Spec":
|
||||
return cls(lambda ctx: True, "always")
|
||||
|
||||
@classmethod
|
||||
def of(cls, fn: Callable[[CallContext], bool], desc: str) -> "Spec":
|
||||
return cls(fn, desc)
|
||||
|
||||
@classmethod
|
||||
def dtype_in(cls, *dtypes: torch.dtype) -> "Spec":
|
||||
names = [str(getattr(d, "name", d)) for d in dtypes]
|
||||
return cls(
|
||||
lambda ctx: ctx.dtype in dtypes, f"dtype∈{{{','.join(names)}}}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def ndim(cls, n: int) -> "Spec":
|
||||
return cls(lambda ctx: ctx.ndim == n, f"ndim=={n}")
|
||||
|
||||
@classmethod
|
||||
def head_dim_in(cls, *dims: int) -> "Spec":
|
||||
return cls(
|
||||
lambda ctx: ctx.head_dim in dims,
|
||||
f"head_dim∈{{{','.join(map(str, dims))}}}",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fwd_in(cls, *modes: Optional[str]) -> "Spec":
|
||||
names = [str(m) for m in modes]
|
||||
return cls(lambda ctx: ctx.fwd in modes, f"fwd∈{{{','.join(names)}}}")
|
||||
|
||||
@classmethod
|
||||
def inference(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.fwd is not None, "inference")
|
||||
|
||||
@classmethod
|
||||
def training(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.fwd is None, "training")
|
||||
|
||||
@classmethod
|
||||
def no_grad(cls) -> "Spec":
|
||||
return cls(lambda ctx: not ctx.grad_enabled, "no_grad")
|
||||
|
||||
@classmethod
|
||||
def has_cache(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.has_cache, "has_cache")
|
||||
|
||||
@classmethod
|
||||
def mask_free(cls) -> "Spec":
|
||||
return cls(lambda ctx: not ctx.has_mask, "mask_free")
|
||||
|
||||
@classmethod
|
||||
def cuda_device(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.device_cuda, "cuda_device")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImplRecord:
|
||||
"""One decision-table row: an implementation plus its capability."""
|
||||
|
||||
family: str
|
||||
name: str
|
||||
obj: Any
|
||||
spec: Spec
|
||||
available: Callable[[], bool] = lambda: True
|
||||
priority: int = 100
|
||||
faithful: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpFamily:
|
||||
name: str
|
||||
provider: Callable[[], List[ImplRecord]]
|
||||
fallback: Callable[[], ImplRecord]
|
||||
|
||||
|
||||
_FAMILIES: Dict[str, OpFamily] = {}
|
||||
_ENV_ALIASES: Dict[str, str] = {}
|
||||
|
||||
_current_overrides: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
|
||||
"astrai_op_overrides", default={}
|
||||
)
|
||||
|
||||
_env_lock = threading.Lock()
|
||||
_env_cache: Dict[tuple, Optional[Dict[str, str]]] = {}
|
||||
_warned: set = set()
|
||||
|
||||
|
||||
def register_family(
|
||||
name: str,
|
||||
provider: Callable[[], List[ImplRecord]],
|
||||
fallback: Callable[[], ImplRecord],
|
||||
) -> None:
|
||||
"""Register (or replace) a family; ``provider`` is re-evaluated per
|
||||
resolution so availability changes (tests, late imports) are honored."""
|
||||
_FAMILIES[name] = OpFamily(name, provider, fallback)
|
||||
|
||||
|
||||
def register_env_alias(family: str, varname: str) -> None:
|
||||
"""Legacy single-value env var for a family (e.g. attention →
|
||||
ASTR_BACKEND); an ASTR_OPS entry wins when both are set."""
|
||||
_ENV_ALIASES[family] = varname
|
||||
|
||||
|
||||
def _family(name: str) -> OpFamily:
|
||||
fam = _FAMILIES.get(name)
|
||||
if fam is None:
|
||||
raise KeyError(f"no operator family registered under {name!r}")
|
||||
return fam
|
||||
|
||||
|
||||
def _warn_once(message: str) -> None:
|
||||
if message not in _warned:
|
||||
_warned.add(message)
|
||||
logger.warning(message)
|
||||
|
||||
|
||||
def set_override(family: str, handle: Any) -> contextvars.Token:
|
||||
overrides = dict(_current_overrides.get())
|
||||
overrides[family] = handle
|
||||
return _current_overrides.set(overrides)
|
||||
|
||||
|
||||
def reset_override(token: contextvars.Token) -> None:
|
||||
_current_overrides.reset(token)
|
||||
|
||||
|
||||
def get_override(family: str) -> Optional[Any]:
|
||||
return _current_overrides.get().get(family)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def op_backend(**handles: Any):
|
||||
"""Select implementations per family for the enclosed scope::
|
||||
|
||||
with op_backend(attention="torch_native", rotary="torch"):
|
||||
engine.generate(...)
|
||||
|
||||
String handles are validated eagerly against the family's currently
|
||||
available implementations; object handles pass through unchecked.
|
||||
"""
|
||||
for family, handle in handles.items():
|
||||
if isinstance(handle, str):
|
||||
fam = _FAMILIES.get(family)
|
||||
if fam is None:
|
||||
raise ValueError(f"unknown operator family {family!r}")
|
||||
if _record_for_handle(fam, handle) is None:
|
||||
raise ValueError(f"Unknown {family} implementation: {handle!r}")
|
||||
tokens = [set_override(f, h) for f, h in handles.items()]
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for token in reversed(tokens):
|
||||
reset_override(token)
|
||||
|
||||
|
||||
def env_overrides() -> Dict[str, str]:
|
||||
"""Merged ASTR_OPS + legacy-alias selections (family or "profile").
|
||||
|
||||
Cached per distinct env content; unknown families / malformed entries
|
||||
warn once and are dropped (soft override, never fatal).
|
||||
"""
|
||||
with _env_lock:
|
||||
merged: Dict[str, str] = {}
|
||||
raw = os.environ.get("ASTR_OPS", "").strip()
|
||||
if raw:
|
||||
key = ("ASTR_OPS", raw)
|
||||
if key not in _env_cache:
|
||||
parsed: Dict[str, str] = {}
|
||||
for item in raw.split(","):
|
||||
key_part, sep, value = item.strip().partition("=")
|
||||
key_part, value = key_part.strip(), value.strip()
|
||||
if not sep or not key_part or not value:
|
||||
_warn_once(f"ASTR_OPS: ignoring malformed entry {item!r}")
|
||||
continue
|
||||
parsed[key_part] = value
|
||||
_env_cache[key] = parsed or None
|
||||
merged.update(_env_cache[key] or {})
|
||||
for fam, varname in _ENV_ALIASES.items():
|
||||
raw = os.environ.get(varname, "").strip()
|
||||
if raw:
|
||||
key = (varname, raw)
|
||||
if key not in _env_cache:
|
||||
_env_cache[key] = {fam: raw.lower()}
|
||||
merged.setdefault(fam, _env_cache[key][fam])
|
||||
for fam in [f for f in merged if f not in _FAMILIES and f != "profile"]:
|
||||
_warn_once(f"ASTR_OPS: unknown operator family {fam!r}; dropping it")
|
||||
merged.pop(fam)
|
||||
return merged
|
||||
|
||||
|
||||
def env_selection(family: str) -> Optional[str]:
|
||||
return env_overrides().get(family)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Resolution:
|
||||
record: ImplRecord
|
||||
origin: str
|
||||
|
||||
|
||||
class ExplicitSelectionError(RuntimeError):
|
||||
"""An explicitly selected implementation cannot handle the call."""
|
||||
|
||||
|
||||
def _record_for_handle(fam: OpFamily, handle: Any) -> Optional[ImplRecord]:
|
||||
records = sorted(fam.provider(), key=lambda r: r.priority)
|
||||
if isinstance(handle, str):
|
||||
return next((r for r in records if r.name == handle), None)
|
||||
return next((r for r in records if r.obj is handle), None)
|
||||
|
||||
|
||||
def _adhoc_record(family: str, handle: Any) -> ImplRecord:
|
||||
"""Wrap an unregistered object; capability probes its own method."""
|
||||
supports = getattr(handle, "supports_call", None)
|
||||
if supports is not None:
|
||||
spec = Spec.of(
|
||||
lambda ctx: bool(supports(*ctx.raw, **ctx.raw_kwargs)),
|
||||
f"{type(handle).__name__}.supports_call",
|
||||
)
|
||||
else:
|
||||
spec = Spec.always()
|
||||
return ImplRecord(family, type(handle).__name__, handle, spec)
|
||||
|
||||
|
||||
def resolve(
|
||||
family: str, ctx: CallContext, explicit: Optional[Any] = None
|
||||
) -> Resolution:
|
||||
"""Resolve one family for one call (explicit-strict / implicit-loose)."""
|
||||
fam = _family(family)
|
||||
|
||||
handle: Optional[Any] = None
|
||||
origin = "chain"
|
||||
if explicit is not None:
|
||||
handle, origin = explicit, "explicit"
|
||||
elif get_override(family) is not None:
|
||||
handle, origin = get_override(family), "context"
|
||||
else:
|
||||
env_name = env_selection(family)
|
||||
if env_name is not None:
|
||||
handle, origin = env_name, "env"
|
||||
|
||||
if handle is not None:
|
||||
record = _record_for_handle(fam, handle)
|
||||
if record is None and not isinstance(handle, str):
|
||||
record = _adhoc_record(family, handle)
|
||||
if record is None:
|
||||
if origin in ("explicit", "context"):
|
||||
raise ValueError(f"Unknown {family} implementation: {handle!r}")
|
||||
_warn_once(f"ASTR_OPS: {family}={handle!r} is not registered; ignoring")
|
||||
else:
|
||||
if record.available() and record.spec.matches(ctx):
|
||||
return Resolution(record, origin)
|
||||
if origin in ("explicit", "context"):
|
||||
raise ExplicitSelectionError(
|
||||
f"Explicitly-set backend {type(record.obj).__name__} cannot "
|
||||
f"handle this {family} call; required: {record.spec.description}"
|
||||
)
|
||||
|
||||
if handle is None and env_overrides().get("profile") == "reference":
|
||||
return Resolution(fam.fallback(), "profile")
|
||||
|
||||
for record in sorted(fam.provider(), key=lambda r: r.priority):
|
||||
if record.available() and record.faithful and record.spec.matches(ctx):
|
||||
return Resolution(record, "chain")
|
||||
return Resolution(fam.fallback(), "fallback")
|
||||
|
||||
|
||||
def resolve_plan(ctxs: Mapping[str, CallContext]) -> Dict[str, Resolution]:
|
||||
"""Resolve several families at once (one decision snapshot)."""
|
||||
return {family: resolve(family, ctx) for family, ctx in ctxs.items()}
|
||||
|
||||
|
||||
def _describe_ctx(ctx: CallContext) -> str:
|
||||
parts = [f"fwd={ctx.fwd}", f"dtype={ctx.dtype}"]
|
||||
if ctx.ndim is not None:
|
||||
parts.append(f"ndim={ctx.ndim}")
|
||||
if ctx.head_dim is not None:
|
||||
parts.append(f"head_dim={ctx.head_dim}")
|
||||
parts += [f"cache={ctx.has_cache}", f"mask={ctx.has_mask}", f"grad={ctx.grad_enabled}"]
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def explain(family: str, ctx: CallContext, explicit: Optional[Any] = None) -> str:
|
||||
"""Human-readable decision trace for one family."""
|
||||
fam = _family(family)
|
||||
records = sorted(fam.provider(), key=lambda r: r.priority)
|
||||
lines = [f"[{family}] {_describe_ctx(ctx)}"]
|
||||
for record in records:
|
||||
if not record.available():
|
||||
lines.append(f" {record.name}: SKIP unavailable")
|
||||
elif not record.faithful:
|
||||
lines.append(f" {record.name}: SKIP not faithful (explicit-only)")
|
||||
elif record.spec.matches(ctx):
|
||||
lines.append(f" {record.name}: MATCH ({record.spec.description})")
|
||||
else:
|
||||
lines.append(f" {record.name}: reject ({record.spec.description})")
|
||||
try:
|
||||
resolution = resolve(family, ctx, explicit)
|
||||
lines.append(f" => {resolution.record.name} (origin={resolution.origin})")
|
||||
except (ExplicitSelectionError, ValueError) as exc:
|
||||
lines.append(f" => ERROR: {exc}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def explain_plan(ctxs: Mapping[str, CallContext]) -> str:
|
||||
return "\n".join(explain(family, ctx) for family, ctx in ctxs.items())
|
||||
@@ -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