diff --git a/docs/source/basic-usage.md b/docs/source/basic-usage.md index 29b175865..481205a1d 100644 --- a/docs/source/basic-usage.md +++ b/docs/source/basic-usage.md @@ -65,6 +65,46 @@ for decision in get_kernel_variants("kernels-community/activation", version=1): print(f"{name}: rejected ({decision.reason})") ``` +## Device Architecture Checks + +A kernel build declares the device architectures it was built for (e.g. CUDA +compute capabilities such as `8.0` or `9.0a`, or ROCm archs such as `gfx90a`) +in its metadata. [`~kernels.get_kernel`] validates these architectures against +the current device and raises an error when the device is not supported. This +check can be disabled by passing `check_arch=False` to [`~kernels.get_kernel`]. +Without this check, an incompatible kernel would load fine, but fail at +launch time — some kernels even terminate the whole process rather than +raising an exception. + +[`~kernels.has_kernel`] performs the same validation and returns `False` when +the kernel does not support the current device, so you can select a different +kernel or fall back to another implementation without hard-failing. + +The declared architectures are the union over all of a kernel's components, so +a kernel may support more architectures than it declares. For example, a +kernel could fall back to a Triton implementation on architectures that its +CUDA kernels were not built for. For such kernels, the check can be disabled +with `check_arch=False`: + +```python +from kernels import get_kernel, has_kernel + +# Only checks that a compatible build variant exists, the kernel may not +# support the device architecture. +has_kernel("kernels-community/flash-attn3", version=1, check_arch=False) + +# Skips the device architecture check when loading. +flash_attn3 = get_kernel("kernels-community/flash-attn3", version=1, check_arch=False) +``` + +> [!WARNING] +> When the architecture check is disabled, loading an incompatible kernel +> succeeds, but using it can fail or even terminate the process at run time. + +For kernels that are used through layers, prefer +[registering kernels for specific CUDA capabilities](./layers.md#registering-kernels-for-specific-cuda-capabilities) +instead — `kernels` will then never load an incompatible kernel at all. + ## Inspecting Loaded Kernels [`~kernels.get_loaded_kernels`] returns a snapshot of every kernel that has been loaded diff --git a/kernels/src/kernels/archs.py b/kernels/src/kernels/archs.py new file mode 100644 index 000000000..a9346f588 --- /dev/null +++ b/kernels/src/kernels/archs.py @@ -0,0 +1,115 @@ +from kernels_data import Backend, Metadata + +from kernels.compat import has_torch + + +def _parse_cuda_arch(arch: str) -> tuple[int, int, str, bool] | None: + """Parse a CUDA arch into its capability, suffix, and PTX flag. + + CUDA archs are compute capabilities like `9.0`, optionally with an + architecture-specific (`9.0a`) or family-specific (`10.0f`) suffix + and/or a `+PTX` flag (`9.0+PTX`). Returns `None` when the arch string + is not in this format. + """ + ptx = False + if arch.endswith("+PTX"): + ptx = True + arch = arch[: -len("+PTX")] + + suffix = "" + if arch.endswith(("a", "f")): + suffix = arch[-1] + arch = arch[:-1] + + arch_major, sep, arch_minor = arch.partition(".") + if not sep or not arch_major.isdigit() or not arch_minor.isdigit(): + return None + + return int(arch_major), int(arch_minor), suffix, ptx + + +def _cuda_arch_supports(arch: str, capability: tuple[int, int]) -> bool | None: + """Check whether a single declared CUDA arch supports a compute capability. + + Returns `None` when the arch string is not in a known format. + """ + parsed = _parse_cuda_arch(arch) + if parsed is None: + return None + + arch_major, arch_minor, suffix, ptx = parsed + major, minor = capability + + if suffix == "a": + # Architecture-specific builds only run on that exact capability. + return (major, minor) == (arch_major, arch_minor) + + if suffix == "f": + # Family-specific builds run on capabilities of the same generation + # with the same or a newer minor version. + return major == arch_major and minor >= arch_minor + + if ptx: + # PTX is JIT-compiled for the current device, so the build also runs + # on any newer capability. + return (major, minor) >= (arch_major, arch_minor) + + # Base builds run on capabilities of the same generation with the same + # or a newer minor version. + return major == arch_major and minor >= arch_minor + + +def _supports_cuda_capability(archs: list[str], capability: tuple[int, int]) -> bool: + supports = [_cuda_arch_supports(arch, capability) for arch in archs] + if all(support is None for support in supports): + # None of the arch strings are in a known format (e.g. produced by a + # newer kernel-builder), so compatibility cannot be determined. + return True + return any(supports) + + +def _check_arch_incompatibility(metadata: Metadata, variant: str) -> None: + """Check a kernel build against the architecture of the current device. + + Raises `RuntimeError` when the build does not support the device. + """ + archs = metadata.backend.archs + if not archs or not has_torch: + return + + import torch + + reason = None + backend_type = metadata.backend.backend_type + if backend_type == Backend.CUDA: + if torch.version.cuda is None or not torch.cuda.is_available(): + return + major, minor = torch.cuda.get_device_capability() + if not _supports_cuda_capability(archs, (major, minor)): + reason = ( + f"CUDA capability {major}.{minor} of the current device is not " + f"supported by the architectures of the build: {', '.join(archs)}" + ) + elif backend_type == Backend.ROCm: + if torch.version.hip is None or not torch.cuda.is_available(): + return + gcn_arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName + # Strip feature flags, e.g. `gfx90a:sramecc+:xnack-` -> `gfx90a`. + gcn_arch = gcn_arch.split(":")[0] + if gcn_arch not in archs: + reason = ( + f"ROCm arch {gcn_arch} of the current device is not supported " + f"by the architectures of the build: {', '.join(archs)}" + ) + + if reason is None: + return + + raise RuntimeError( + f"Kernel '{metadata.name}' variant '{variant}' does not support the " + f"current device: {reason}. The declared architectures can be " + "incomplete (e.g. when a kernel has a Triton fallback), pass " + "`check_arch=False` to skip this check. Warning: the kernel may " + "fail or crash the process at run time when it does not support " + "the device." + ) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index c54a59227..03ff9ac1e 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -14,11 +14,12 @@ from types import ModuleType from huggingface_hub import HfApi, constants -from huggingface_hub.errors import LocalEntryNotFoundError +from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError from kernels_data import Metadata from kernels._system import glibc_version from kernels._versions import select_revision_or_version +from kernels.archs import _check_arch_incompatibility from kernels.backends import _backend, _select_backend from kernels.compat import has_torch, has_tvm_ffi from kernels.deps import validate_dependencies @@ -455,6 +456,7 @@ def get_kernel( backend: str | None = None, user_agent: str | dict | None = None, trust_remote_code: bool | list[str] = False, + check_arch: bool = True, ) -> ModuleType: """ Load a kernel from the kernel hub. @@ -481,6 +483,12 @@ def get_kernel( repositories are allowed. A list of strings will be used to verify signing identities in a future release; for now it emits a warning and falls back to the default trust check. + check_arch (`bool`, *optional*, defaults to `True`): + Whether to check that the kernel build supports the architecture + (e.g. CUDA compute capability) of the current device. Kernels can + support more architectures than they declare (e.g. through a + Triton fallback), `check_arch=False` skips the check for such + kernels. Returns: `ModuleType`: The imported kernel module. @@ -514,6 +522,9 @@ def get_kernel( user_agent=user_agent, validate_dependencies=True, ) + if check_arch: + metadata = Metadata.read_from_file(variant_path / "metadata.json") + _check_arch_incompatibility(metadata, variant_path.name) return _import_from_path(variant_path, repo_info=repo_info) @@ -558,6 +569,7 @@ def has_kernel( revision: str | None = None, version: int | None = None, backend: str | None = None, + check_arch: bool = True, ) -> bool: """ Check whether a kernel build exists for the current environment (Torch version and compute framework). @@ -573,6 +585,12 @@ def has_kernel( backend (`str`, *optional*): The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. The backend will be detected automatically if not provided. + check_arch (`bool`, *optional*, defaults to `True`): + Whether to check that the kernel build supports the architecture + (e.g. CUDA compute capability) of the current device. Kernels can + support more architectures than they declare (e.g. through a + Triton fallback), `check_arch=False` skips the check for such + kernels. Returns: `bool`: `True` if a kernel is available for the current environment. @@ -586,12 +604,31 @@ def has_kernel( if variant is None: return False - return api.file_exists( - repo_id, - repo_type="kernel", - revision=revision, - filename=f"build/{variant.variant_str}/metadata.json", - ) + if not check_arch: + return api.file_exists( + repo_id, + repo_type="kernel", + revision=revision, + filename=f"build/{variant.variant_str}/metadata.json", + ) + + try: + metadata_path = api.hf_hub_download( + repo_id, + repo_type="kernel", + revision=revision, + filename=f"build/{variant.variant_str}/metadata.json", + cache_dir=CACHE_DIR, + ) + except EntryNotFoundError: + return False + + try: + _check_arch_incompatibility(Metadata.read_from_file(metadata_path), variant.variant_str) + except RuntimeError: + return False + + return True def get_kernel_variants( diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py new file mode 100644 index 000000000..fa57843d3 --- /dev/null +++ b/kernels/tests/test_archs.py @@ -0,0 +1,157 @@ +import json + +import pytest +import torch +from kernels_data import Metadata + +from kernels import get_kernel, has_kernel, install_kernel +from kernels.archs import _check_arch_incompatibility, _supports_cuda_capability + + +def make_metadata(backend_type: str, archs: list[str] | None) -> Metadata: + return Metadata.from_bytes( + json.dumps( + { + "name": "test-kernel", + "id": "test_id", + "version": 1, + "license": "mit", + "python-depends": [], + "backend": {"type": backend_type, "archs": archs}, + } + ).encode("utf-8") + ) + + +@pytest.mark.parametrize( + "archs,capability,supported", + [ + # Base archs run on the same generation with the same or a newer + # minor capability. + (["8.0"], (8, 0), True), + (["8.0"], (8, 6), True), + (["8.6"], (8, 0), False), + (["8.0"], (7, 5), False), + (["8.0"], (9, 0), False), + (["8.0", "9.0"], (9, 0), True), + # Architecture-specific archs only run on that exact capability. + (["9.0a"], (9, 0), True), + (["9.0a"], (9, 1), False), + (["8.0", "9.0a"], (10, 0), False), + # Family-specific archs run on the same generation with the same or + # a newer minor capability. + (["10.0f"], (10, 0), True), + (["10.0f"], (10, 3), True), + (["10.0f"], (12, 0), False), + # PTX is JIT-compiled for the current device, so builds with `+PTX` + # also run on any newer capability. + (["9.0+PTX"], (9, 0), True), + (["9.0+PTX"], (9, 1), True), + (["9.0+PTX"], (10, 0), True), + (["9.0+PTX"], (12, 1), True), + (["9.0+PTX"], (8, 6), False), + # ...but the `a` and `f` suffixes keep their own semantics. + (["9.0a+PTX"], (9, 0), True), + (["9.0a+PTX"], (10, 0), False), + (["10.0f+PTX"], (10, 3), True), + (["10.0f+PTX"], (12, 0), False), + # Arch strings in an unknown format do not count as a match... + (["garbage", "8.0"], (9, 0), False), + # ...but when no arch string can be parsed, compatibility cannot be + # determined and the build is not rejected. + (["garbage"], (9, 0), True), + ], +) +def test_supports_cuda_capability(archs, capability, supported): + assert _supports_cuda_capability(archs, capability) == supported + + +@pytest.fixture +def fake_cuda_device(monkeypatch): + monkeypatch.setattr(torch.version, "cuda", "12.8", raising=False) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (10, 0)) + + +@pytest.fixture +def fake_rocm_device(monkeypatch): + class FakeProperties: + gcnArchName = "gfx90a:sramecc+:xnack-" + + monkeypatch.setattr(torch.version, "cuda", None, raising=False) + monkeypatch.setattr(torch.version, "hip", "6.4.0", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: FakeProperties()) + + +def test_cuda_incompatible_arch_is_rejected(fake_cuda_device): + with pytest.raises(RuntimeError) as exc_info: + _check_arch_incompatibility(make_metadata("cuda", ["8.0", "9.0a"]), "test-variant") + assert "test-variant" in str(exc_info.value) + assert "CUDA capability 10.0" in str(exc_info.value) + assert "8.0, 9.0a" in str(exc_info.value) + + +def test_cuda_compatible_arch_is_accepted(fake_cuda_device): + for archs in (["8.0", "10.0"], ["10.0a"], ["10.0f"]): + _check_arch_incompatibility(make_metadata("cuda", archs), "test-variant") + + +def test_noarch_build_is_accepted(fake_cuda_device): + # Backends that support archs (e.g. CUDA) can have builds that do not + # declare any (noarch kernels, e.g. pure Triton builds). Such builds are + # never rejected. + _check_arch_incompatibility(make_metadata("cuda", None), "test-variant") + _check_arch_incompatibility(make_metadata("cuda", []), "test-variant") + + +def test_rocm_arch_check(fake_rocm_device): + _check_arch_incompatibility(make_metadata("rocm", ["gfx90a", "gfx942"]), "test-variant") + _check_arch_incompatibility(make_metadata("rocm", None), "test-variant") + with pytest.raises(RuntimeError) as exc_info: + _check_arch_incompatibility(make_metadata("rocm", ["gfx942"]), "test-variant") + assert "ROCm arch gfx90a" in str(exc_info.value) + assert "gfx942" in str(exc_info.value) + + +def test_check_skipped_without_device(monkeypatch): + monkeypatch.setattr(torch.version, "cuda", "12.8", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + _check_arch_incompatibility(make_metadata("cuda", ["8.0"]), "test-variant") + + +def test_check_skipped_for_backends_without_archs(fake_cuda_device): + # Archs of other backends cannot be checked against the current device. + _check_arch_incompatibility(make_metadata("cpu", None), "test-variant") + _check_arch_incompatibility(make_metadata("metal", ["applegpu_g13"]), "test-variant") + + +def test_issue_707_fa3_on_b200(fake_cuda_device): + # https://github.com/huggingface/kernels/issues/707: flash-attn3 only + # declares sm_80/sm_90a archs, but loading it on a B200 (capability 10.0) + # succeeded and the first launch exited the process. The declared archs + # must be rejected for this device. + metadata = make_metadata("cuda", ["8.0", "9.0a"]) + with pytest.raises(RuntimeError, match="does not support the current device"): + _check_arch_incompatibility(metadata, "test-variant") + + +@pytest.mark.cuda_only +def test_get_kernel_rejects_unsupported_capability(monkeypatch): + variant_path = install_kernel("kernels-community/relu", revision="v1") + metadata = Metadata.read_from_file(variant_path / "metadata.json") + if not metadata.backend.archs: + pytest.skip("kernel build does not declare archs") + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (99, 9)) + + with pytest.raises(RuntimeError, match="does not support the current device"): + get_kernel("kernels-community/relu", version=1) + assert not has_kernel("kernels-community/relu", version=1) + + # The opt-out only checks that a compatible build variant exists. + assert has_kernel("kernels-community/relu", version=1, check_arch=False) + kernel = get_kernel("kernels-community/relu", version=1, check_arch=False) + assert kernel is not None