feat: implement capability check in main

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

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."
+    )

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]

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)
 

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)