feat: arch/cc compatibility check.

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/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(