Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/source/kernel-requirements.md
Original file line number Diff line number Diff line change
Expand Up @@ -277,8 +277,8 @@ These requirements can be checked with the ABI checker (see below).

- Use [ABI3/Limited API](https://docs.python.org/3/c-api/stable.html#stable-application-binary-interface)
for compatibility with Python 3.9 and later.
- macOS deployment target 15.0.
- Metal 3.0 (`-std=metal3.0`).
- macOS deployment target 26.0.
- Metal 4.0 (`-std=metal4.0`).

The ABI3 requirement can be checked with the ABI checker (see below).

Expand Down
25 changes: 25 additions & 0 deletions kernels/src/kernels/variants.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,19 @@
CUDA,
XPU,
Backend,
Metal,
ROCm,
_select_backend,
parse_backend,
)
from kernels.compat import has_torch, has_tvm_ffi

# Metal kernels are currently built with `-std=metal4.0`, which requires
# macOS 26 or later. Until the macOS/Metal version is encoded in the build
# variant, reject Metal arch variants on older systems, since they fail to
# load at runtime.
_METAL_MIN_MACOS_VERSION = Version("26.0")


@dataclass(unsafe_hash=True)
class Torch:
Expand Down Expand Up @@ -285,8 +292,12 @@ def resolve_variants(variants: list[Variant], backend: str | None = None) -> tup
cpu = platform.machine()
os = platform.system().lower()

macos_version = None
if os == "darwin":
cpu = "aarch64" if cpu == "arm64" else cpu
mac_release = platform.mac_ver()[0]
if mac_release:
macos_version = parse(mac_release)
elif os == "windows":
cpu = "x86_64" if cpu == "AMD64" else cpu

Expand Down Expand Up @@ -317,6 +328,7 @@ def resolve_variants(variants: list[Variant], backend: str | None = None) -> tup
torch_version=torch_version,
torch_cxx11_abi=torch_cxx11_abi,
tvm_ffi_version=tvm_ffi_version,
macos_version=macos_version,
)


Expand All @@ -328,6 +340,7 @@ def _resolve_variant_for_system(
torch_version: Version | None,
torch_cxx11_abi: bool | None,
tvm_ffi_version: Version | None,
macos_version: Version | None = None,
) -> tuple[list[Variant], list[Decision]]:
"""Resolve the best matching variant given explicit system parameters.

Expand All @@ -341,6 +354,7 @@ def _resolve_variant_for_system(
torch_version,
torch_cxx11_abi,
tvm_ffi_version,
macos_version,
)
trace = _sort_variants(trace)

Expand All @@ -362,6 +376,7 @@ def _check_variants(
torch_version: Version | None,
torch_cxx11_abi: bool | None,
tvm_ffi_version: Version | None,
macos_version: Version | None = None,
) -> list[Decision]:
"""Return only the variants applicable to the current system."""
is_unsupported_free_threaded = _is_unsupported_free_threaded_build()
Expand Down Expand Up @@ -466,6 +481,16 @@ def _check_variants(
)
)
continue
elif isinstance(v.arch.backend, Metal) and (
macos_version is None or macos_version < _METAL_MIN_MACOS_VERSION
):
result.append(
VariantRejected(
variant=v,
reason=f"Metal kernels require macOS {_METAL_MIN_MACOS_VERSION} or later, system macOS version is {macos_version}",
)
)
continue
elif isinstance(v, NoarchVariant):
# Only noarch variants with a matching backend or "universal"
# are applicable.
Expand Down
58 changes: 55 additions & 3 deletions kernels/tests/test_variants.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@
from huggingface_hub import HfApi
from packaging.version import Version

from kernels.backends import CPU, CUDA, ROCm
from kernels.backends import CPU, CUDA, Metal, ROCm
from kernels.variants import (
VariantAccepted,
VariantRejected,
_resolve_variant_for_system,
get_variants,
parse_variant,
Expand Down Expand Up @@ -283,19 +284,70 @@ def test_resolve_cpu_darwin():
def test_resolve_metal_darwin():
result, trace = _resolve_variant_for_system(
variants=RESOLVE_VARIANTS,
selected_backend=CPU(),
selected_backend=Metal(),
cpu="aarch64",
os="darwin",
torch_version=Version("2.10"),
torch_cxx11_abi=None,
tvm_ffi_version=None,
macos_version=Version("26.0"),
)
assert result != []
assert result[0].variant_str == "torch210-cpu-aarch64-darwin"
assert result[0].variant_str == "torch210-metal-aarch64-darwin"
assert result == [vs.variant for vs in trace if isinstance(vs, VariantAccepted)]
assert {vs.variant for vs in trace} == set(RESOLVE_VARIANTS)


RESOLVE_VARIANTS_METAL = [
parse_variant(s)
for s in [
"torch210-cpu-aarch64-darwin",
"torch210-metal-aarch64-darwin",
"torch-metal",
]
]


@pytest.mark.parametrize("macos_version", [Version("15.7"), None])
def test_resolve_metal_darwin_old_macos(macos_version):
# Metal arch kernels are built for macOS 26+, so they must be rejected
# on older systems (falling back to the noarch variant).
result, trace = _resolve_variant_for_system(
variants=RESOLVE_VARIANTS_METAL,
selected_backend=Metal(),
cpu="aarch64",
os="darwin",
torch_version=Version("2.10"),
torch_cxx11_abi=None,
tvm_ffi_version=None,
macos_version=macos_version,
)
assert result != []
assert result[0].variant_str == "torch-metal"
rejected = {vs.variant.variant_str: vs.reason for vs in trace if isinstance(vs, VariantRejected)}
assert "require macOS 26.0" in rejected["torch210-metal-aarch64-darwin"]
assert result == [vs.variant for vs in trace if isinstance(vs, VariantAccepted)]
assert {vs.variant for vs in trace} == set(RESOLVE_VARIANTS_METAL)


def test_resolve_metal_darwin_new_macos():
# On macOS 26+ the Metal arch kernel is accepted and preferred.
result, trace = _resolve_variant_for_system(
variants=RESOLVE_VARIANTS_METAL,
selected_backend=Metal(),
cpu="aarch64",
os="darwin",
torch_version=Version("2.10"),
torch_cxx11_abi=None,
tvm_ffi_version=None,
macos_version=Version("26.1"),
)
assert result != []
assert result[0].variant_str == "torch210-metal-aarch64-darwin"
assert result == [vs.variant for vs in trace if isinstance(vs, VariantAccepted)]
assert {vs.variant for vs in trace} == set(RESOLVE_VARIANTS_METAL)


def test_resolve_noarch_fallback():
# With no matching arch variant, should fall back to torch noarch.
result, trace = _resolve_variant_for_system(
Expand Down
Loading