diff --git a/docs/source/basic-usage.md b/docs/source/basic-usage.md index 29b17586..481205a1 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 00000000..a9346f58 --- /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/deps.py b/kernels/src/kernels/deps.py index 3e0b3c6e..17e0ada3 100644 --- a/kernels/src/kernels/deps.py +++ b/kernels/src/kernels/deps.py @@ -1,11 +1,12 @@ from contextvars import ContextVar from dataclasses import dataclass from types import ModuleType -from typing import Generic, TypeVar +from typing import Generic, Protocol, TypeVar from huggingface_hub.hf_api import HfApi from kernels_data import KernelDependency, KernelVersion +from kernels.archs import _check_arch_incompatibility from kernels.backends import _backend from kernels.hf_hub import RepoInfo from kernels.importer import _import_from_path @@ -69,17 +70,51 @@ def load(self) -> ModuleType: return _import_from_path(self.location.variant_path, repo_info=repo_info, deps=deps) - def validate_dependencies(self) -> None: - """Validate that the dependencies of this kernel are satisfied.""" + def validate( + self: "DepTreeNode[LocalKernel | RemoteKernel]", + validator: "Validator", + ) -> None: + """Validate this kernel and its dependencies with the given validator.""" + validator.validate(tree=self) + + for node in self.deps.values(): + node.validate(validator) + + +class Validator(Protocol): + """Validator for a node in a resolved kernel dependency tree.""" + + def validate(self, *, tree: DepTreeNode[LocalKernel | RemoteKernel]) -> None: ... + + +class DependencyValidator: + """Validate the Python dependencies of a kernel tree node.""" + + def validate(self, *, tree: DepTreeNode[LocalKernel | RemoteKernel]) -> None: validate_dependencies( - self.location.metadata.name.python_name, - self.location.metadata.python_depends, + tree.location.metadata.name.python_name, + tree.location.metadata.python_depends, _backend(), ) - for node in self.deps.values(): - node.validate_dependencies() + +class ArchValidator: + """Validate architecture compatibility for a kernel tree node.""" + + def validate(self, *, tree: DepTreeNode[LocalKernel | RemoteKernel]) -> None: + _check_arch_incompatibility(tree.location.metadata, tree.location.variant_str) + + +@dataclass +class AllValidator: + """Apply multiple validators to a kernel dependency tree.""" + + validators: list[Validator] + + def validate(self, *, tree: DepTreeNode[LocalKernel | RemoteKernel]) -> None: + for validator in self.validators: + validator.validate(tree=tree) LoadCache = dict[KernelDependency, DepTreeNode] diff --git a/kernels/src/kernels/layer/func.py b/kernels/src/kernels/layer/func.py index 597aaca3..881258b3 100644 --- a/kernels/src/kernels/layer/func.py +++ b/kernels/src/kernels/layer/func.py @@ -10,6 +10,7 @@ from kernels.resolver import LockedHubCacheResolver, LockedHubResolver from .._versions import select_revision_or_version +from ..deps import DependencyValidator from ..hf_hub import _get_hf_api from ..load import ( get_kernel, @@ -331,6 +332,7 @@ def load(self) -> Type["nn.Module"]: backend=None, kernel=self.kernel_dep, resolver=resolver, + validator=DependencyValidator(), ) return _get_kernel_func(self, kernel) diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 2e7aae20..fa34b8c8 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -15,6 +15,7 @@ from kernels.resolver import LockedHubCacheResolver, LockedHubResolver from .._versions import select_revision_or_version +from ..deps import DependencyValidator from ..hf_hub import _get_hf_api from ..load import ( get_kernel, @@ -240,6 +241,7 @@ def load(self) -> Type["nn.Module"]: backend=None, kernel=self.kernel_dep, resolver=resolver, + validator=DependencyValidator(), ) return _get_kernel_layer(self, kernel) diff --git a/kernels/src/kernels/load.py b/kernels/src/kernels/load.py index bddb0c01..e37af7ee 100644 --- a/kernels/src/kernels/load.py +++ b/kernels/src/kernels/load.py @@ -7,7 +7,7 @@ from kernels_data import KernelDependency, KernelVersion from kernels._versions import revision_or_version -from kernels.deps import resolve_kernel_tree +from kernels.deps import AllValidator, ArchValidator, DependencyValidator, Validator, resolve_kernel_tree from kernels.hf_hub import _get_hf_api from kernels.locking import ( get_caller_locked_kernel_revision, @@ -55,6 +55,7 @@ def get_kernel_with_resolver( backend: str | None, kernel: KernelDependency, resolver: Resolver | None, + validator: Validator, ) -> ModuleType: """ Load a kernel and its (transitive) dependencies using the given resolver. @@ -69,6 +70,8 @@ def get_kernel_with_resolver( The kernel to load. resolver (`Resolver`, *optional*): The resolver used to resolve the kernel and its (transitive) dependencies. + validator (`Validator`): + The validator to apply to the resolved kernel dependency tree. Returns: `ModuleType`: The imported kernel module. @@ -79,7 +82,7 @@ def get_kernel_with_resolver( kernel=kernel, resolver=resolver, ) - tree.validate_dependencies() + tree.validate(validator) tree_only_local = tree.install(api=api) return tree_only_local.load() @@ -92,6 +95,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. @@ -118,6 +122,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. @@ -146,11 +156,16 @@ def get_kernel( ), ] + validators: list[Validator] = [DependencyValidator()] + if check_arch: + validators.append(ArchValidator()) + return get_kernel_with_resolver( api=api, backend=backend, kernel=KernelDependency(repo_id=repo_id, version=kernel_version), resolver=SequentialResolver(resolvers=resolvers), + validator=AllValidator(validators=validators), ) @@ -204,6 +219,7 @@ def get_local_kernel( # We don't have a name for the kernel, so let's just use the path. kernel=KernelDependency(repo_id=str(repo_path), version=KernelVersion.Version(0)), resolver=SequentialResolver(resolvers), + validator=DependencyValidator(), ) @@ -214,6 +230,7 @@ def has_kernel( version: int | None = None, backend: str | None = None, trust_remote_code: bool | list[str] = False, + check_arch: bool = True, ) -> bool: """ Check whether a kernel build exists for the current environment (framework version and backend). @@ -237,6 +254,12 @@ def has_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: `bool`: `True` if a kernel is available for the current environment. @@ -255,7 +278,7 @@ def has_kernel( ] try: - resolve_kernel_tree( + tree = resolve_kernel_tree( api=api, backend=backend, kernel=KernelDependency(repo_id=repo_id, version=kernel_version), @@ -264,6 +287,12 @@ def has_kernel( except FileNotFoundError: return False + if check_arch: + try: + tree.validate(ArchValidator()) + except RuntimeError: + return False + return True @@ -307,6 +336,7 @@ def load_kernel( backend=backend, kernel=kernel_dep, resolver=resolver, + validator=DependencyValidator(), ) @@ -347,4 +377,5 @@ def get_locked_kernel( backend=None, kernel=kernel_dep, resolver=resolver, + validator=DependencyValidator(), ) diff --git a/kernels/src/kernels/resolver.py b/kernels/src/kernels/resolver.py index 0c6227d6..2d8d994e 100644 --- a/kernels/src/kernels/resolver.py +++ b/kernels/src/kernels/resolver.py @@ -36,6 +36,10 @@ class LocalKernel: metadata: Metadata origin: "RemoteKernel | None" = None + @property + def variant_str(self) -> str: + return self.variant_path.name + def install(self, *, api: HfApi) -> "LocalKernel": # Local kernels are already installed, so we just return self. return self @@ -50,6 +54,10 @@ class RemoteKernel: metadata: Metadata variant: Variant + @property + def variant_str(self) -> str: + return self.variant.variant_str + def install(self, *, api: HfApi) -> LocalKernel: allow_patterns = [f"build/{self.variant.variant_str}/*"] diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py new file mode 100644 index 00000000..6d8a04c9 --- /dev/null +++ b/kernels/tests/test_archs.py @@ -0,0 +1,195 @@ +import json +from pathlib import Path + +import pytest +import torch +from kernels_data import Metadata + +import kernels.deps as deps_module +from kernels import get_kernel, has_kernel, install_kernel +from kernels.archs import _supports_cuda_capability +from kernels.deps import ArchValidator, DepTreeNode +from kernels.resolver import LocalKernel + + +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") + ) + + +def make_tree(metadata: Metadata, variant: str = "test-variant") -> DepTreeNode: + return DepTreeNode( + location=LocalKernel(Path(variant), metadata), + deps={}, + ) + + +@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: + ArchValidator().validate(tree=make_tree(make_metadata("cuda", ["8.0", "9.0a"]))) + 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"]): + ArchValidator().validate(tree=make_tree(make_metadata("cuda", archs))) + + +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. + ArchValidator().validate(tree=make_tree(make_metadata("cuda", None))) + ArchValidator().validate(tree=make_tree(make_metadata("cuda", []))) + + +def test_rocm_arch_check(fake_rocm_device): + ArchValidator().validate(tree=make_tree(make_metadata("rocm", ["gfx90a", "gfx942"]))) + ArchValidator().validate(tree=make_tree(make_metadata("rocm", None))) + with pytest.raises(RuntimeError) as exc_info: + ArchValidator().validate(tree=make_tree(make_metadata("rocm", ["gfx942"]))) + 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) + ArchValidator().validate(tree=make_tree(make_metadata("cuda", ["8.0"]))) + + +def test_check_skipped_for_backends_without_archs(fake_cuda_device): + # Archs of other backends cannot be checked against the current device. + ArchValidator().validate(tree=make_tree(make_metadata("cpu", None))) + ArchValidator().validate(tree=make_tree(make_metadata("metal", ["applegpu_g13"]))) + + +def test_arch_validator_checks_entire_dependency_tree(monkeypatch): + tree = DepTreeNode( + location=LocalKernel(Path("root-variant"), make_metadata("cuda", ["8.0"])), + deps={ + "test/dependency": DepTreeNode( + location=LocalKernel(Path("dependency-variant"), make_metadata("cuda", ["9.0"])), + deps={}, + ) + }, + ) + validated = [] + + monkeypatch.setattr( + deps_module, + "_check_arch_incompatibility", + lambda metadata, variant: validated.append((metadata.backend.archs, variant)), + ) + + tree.validate(ArchValidator()) + + assert validated == [ + (["8.0"], "root-variant"), + (["9.0"], "dependency-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"): + ArchValidator().validate(tree=make_tree(metadata)) + + +@pytest.mark.cuda_only +def test_get_kernel_rejects_unsupported_capability(monkeypatch): + variant_path = install_kernel("kernels-community/relu", version=1) + metadata = Metadata.read_from_file(variant_path / "metadata.json") + if not metadata.backend.archs: + pytest.skip("kernel build does not declare archs") + + # Fake a capability *below* every declared arch. + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (1, 0)) + + 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 diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py index f93db419..1b8e7488 100755 --- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py @@ -1,8 +1,9 @@ import os from pathlib import Path -from kernels.load import get_kernel_with_resolver +from kernels.deps import DependencyValidator from kernels.hf_hub import _get_hf_api +from kernels.load import get_kernel_with_resolver from kernels.resolver import KernelPathsResolver, RepoPathsResolver, SequentialResolver from kernels_data import KernelDependency, KernelPaths, KernelVersion @@ -29,4 +30,5 @@ backend=None, kernel=kernel, resolver=SequentialResolver(resolvers=resolvers), + validator=DependencyValidator(), )