From e25af4d5377df93b180705529a367d02a878ef27 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 18 Aug 2026 19:44:04 +0200 Subject: [PATCH 1/7] feat: arch/cc compatibility check./ --- docs/source/basic-usage.md | 40 +++++++++++ kernels/src/kernels/archs.py | 98 +++++++++++++++++++++++++ kernels/src/kernels/utils.py | 46 ++++++++++-- kernels/tests/test_archs.py | 134 +++++++++++++++++++++++++++++++++++ 4 files changed, 311 insertions(+), 7 deletions(-) create mode 100644 kernels/src/kernels/archs.py create mode 100644 kernels/tests/test_archs.py 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..c1d1c4ce0 --- /dev/null +++ b/kernels/src/kernels/archs.py @@ -0,0 +1,98 @@ +import re + +from kernels_data import Backend, Metadata + +from kernels.compat import has_torch + + +class UnsupportedArchError(RuntimeError): + """The kernel build does not support the architecture of the current device.""" + + +# CUDA archs are compute capabilities like `9.0`, optionally with an +# architecture-specific (`9.0a`) or family-specific (`10.0f`) suffix. +_CUDA_ARCH_REGEX = re.compile(r"(\d+)\.(\d+)([af]?)") + + +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. + """ + m = _CUDA_ARCH_REGEX.fullmatch(arch) + if m is None: + return None + + arch_major, arch_minor, suffix = int(m.group(1)), int(m.group(2)), m.group(3) + major, minor = capability + + if suffix == "a": + # Architecture-specific builds only run on that exact capability. + return (major, minor) == (arch_major, arch_minor) + + # Base and 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 + + +def _cuda_archs_support_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 _arch_incompatibility(metadata: Metadata) -> str | None: + """Check a kernel build against the architecture of the current device.""" + archs = metadata.backend.archs + if not archs or not has_torch: + return None + + import torch + + backend_type = metadata.backend.backend_type + if backend_type == Backend.CUDA: + if torch.version.cuda is None or not torch.cuda.is_available(): + return None + major, minor = torch.cuda.get_device_capability() + if _cuda_archs_support_capability(archs, (major, minor)): + return None + return ( + 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 None + 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 in archs: + return None + return ( + f"ROCm arch {gcn_arch} of the current device is not supported " + f"by the architectures of the build: {', '.join(archs)}" + ) + + return None + + +def check_arch_compatibility(metadata: Metadata, variant_str: str) -> None: + """Raise `UnsupportedArchError` when the build does not support the current device.""" + reason = _arch_incompatibility(metadata) + if reason is not None: + raise UnsupportedArchError( + f"Kernel '{metadata.name}' variant '{variant_str}' does not support " + f"the 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." + ) + + +def is_arch_compatible(metadata: Metadata) -> bool: + """Check whether the build supports the architecture of the current device.""" + return _arch_incompatibility(metadata) is None diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index c54a59227..0a9fe0c96 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_compatibility, is_arch_compatible 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_compatibility(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,26 @@ 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 + + return is_arch_compatible(Metadata.read_from_file(metadata_path)) def get_kernel_variants( diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py new file mode 100644 index 000000000..8f44d8d48 --- /dev/null +++ b/kernels/tests/test_archs.py @@ -0,0 +1,134 @@ +import json + +import pytest +import torch +from kernels_data import Metadata + +from kernels import get_kernel, has_kernel, install_kernel +from kernels.archs import ( + UnsupportedArchError, + _arch_incompatibility, + _cuda_archs_support_capability, + check_arch_compatibility, + is_arch_compatible, +) + + +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), + # 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_cuda_archs_support_capability(archs, capability, supported): + assert _cuda_archs_support_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): + metadata = make_metadata("cuda", ["8.0", "9.0a"]) + assert not is_arch_compatible(metadata) + with pytest.raises(UnsupportedArchError, match=r"CUDA capability 10\.0.*8\.0, 9\.0a"): + check_arch_compatibility(metadata, "torch-stable-abi29-cu128-x86_64-linux") + + +def test_cuda_compatible_arch_is_accepted(fake_cuda_device): + for archs in (["8.0", "10.0"], ["10.0a"], ["10.0f"], None): + metadata = make_metadata("cuda", archs) + assert is_arch_compatible(metadata) + check_arch_compatibility(metadata, "torch-stable-abi29-cu128-x86_64-linux") + + +def test_rocm_arch_check(fake_rocm_device): + assert is_arch_compatible(make_metadata("rocm", ["gfx90a", "gfx942"])) + assert is_arch_compatible(make_metadata("rocm", None)) + metadata = make_metadata("rocm", ["gfx942"]) + assert not is_arch_compatible(metadata) + with pytest.raises(UnsupportedArchError, match=r"ROCm arch gfx90a.*gfx942"): + check_arch_compatibility(metadata, "torch28-cxx11-rocm64-x86_64-linux") + + +def test_check_skipped_without_device(monkeypatch): + monkeypatch.setattr(torch.version, "cuda", "12.8", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert _arch_incompatibility(make_metadata("cuda", ["8.0"])) is None + + +def test_check_skipped_for_backends_without_archs(fake_cuda_device): + # Archs of other backends cannot be checked against the current device. + assert is_arch_compatible(make_metadata("cpu", None)) + assert is_arch_compatible(make_metadata("metal", ["applegpu_g13"])) + + +@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(UnsupportedArchError): + 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 From 8e8726ea1832844682f018ec1f5f3c15745e5599 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 18 Aug 2026 20:04:01 +0200 Subject: [PATCH 2/7] get rid of unneeded stuff. --- kernels/src/kernels/archs.py | 47 ++++++++++++++---------------------- kernels/src/kernels/utils.py | 15 +++++++++--- kernels/tests/test_archs.py | 38 ++++++++++++----------------- 3 files changed, 45 insertions(+), 55 deletions(-) diff --git a/kernels/src/kernels/archs.py b/kernels/src/kernels/archs.py index c1d1c4ce0..2a93ba792 100644 --- a/kernels/src/kernels/archs.py +++ b/kernels/src/kernels/archs.py @@ -1,17 +1,25 @@ -import re - from kernels_data import Backend, Metadata from kernels.compat import has_torch -class UnsupportedArchError(RuntimeError): - """The kernel build does not support the architecture of the current device.""" +def _parse_cuda_arch(arch: str) -> tuple[int, int, str] | None: + """Parse a CUDA arch into its capability and suffix. + CUDA archs are compute capabilities like `9.0`, optionally with an + architecture-specific (`9.0a`) or family-specific (`10.0f`) suffix. + Returns `None` when the arch string is not in this format. + """ + suffix = "" + if arch.endswith(("a", "f")): + suffix = arch[-1] + arch = arch[:-1] -# CUDA archs are compute capabilities like `9.0`, optionally with an -# architecture-specific (`9.0a`) or family-specific (`10.0f`) suffix. -_CUDA_ARCH_REGEX = re.compile(r"(\d+)\.(\d+)([af]?)") + 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 def _cuda_arch_supports(arch: str, capability: tuple[int, int]) -> bool | None: @@ -19,11 +27,11 @@ def _cuda_arch_supports(arch: str, capability: tuple[int, int]) -> bool | None: Returns `None` when the arch string is not in a known format. """ - m = _CUDA_ARCH_REGEX.fullmatch(arch) - if m is None: + parsed = _parse_cuda_arch(arch) + if parsed is None: return None - arch_major, arch_minor, suffix = int(m.group(1)), int(m.group(2)), m.group(3) + arch_major, arch_minor, suffix = parsed major, minor = capability if suffix == "a": @@ -77,22 +85,3 @@ def _arch_incompatibility(metadata: Metadata) -> str | None: ) return None - - -def check_arch_compatibility(metadata: Metadata, variant_str: str) -> None: - """Raise `UnsupportedArchError` when the build does not support the current device.""" - reason = _arch_incompatibility(metadata) - if reason is not None: - raise UnsupportedArchError( - f"Kernel '{metadata.name}' variant '{variant_str}' does not support " - f"the 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." - ) - - -def is_arch_compatible(metadata: Metadata) -> bool: - """Check whether the build supports the architecture of the current device.""" - return _arch_incompatibility(metadata) is None diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index 0a9fe0c96..e8fd959ad 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -19,7 +19,7 @@ from kernels._system import glibc_version from kernels._versions import select_revision_or_version -from kernels.archs import check_arch_compatibility, is_arch_compatible +from kernels.archs import _arch_incompatibility from kernels.backends import _backend, _select_backend from kernels.compat import has_torch, has_tvm_ffi from kernels.deps import validate_dependencies @@ -524,7 +524,16 @@ def get_kernel( ) if check_arch: metadata = Metadata.read_from_file(variant_path / "metadata.json") - check_arch_compatibility(metadata, variant_path.name) + reason = _arch_incompatibility(metadata) + if reason is not None: + raise RuntimeError( + f"Kernel '{metadata.name}' variant '{variant_path.name}' does not " + f"support the 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." + ) return _import_from_path(variant_path, repo_info=repo_info) @@ -623,7 +632,7 @@ def has_kernel( except EntryNotFoundError: return False - return is_arch_compatible(Metadata.read_from_file(metadata_path)) + return _arch_incompatibility(Metadata.read_from_file(metadata_path)) is None def get_kernel_variants( diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py index 8f44d8d48..27d1cabfe 100644 --- a/kernels/tests/test_archs.py +++ b/kernels/tests/test_archs.py @@ -5,13 +5,7 @@ from kernels_data import Metadata from kernels import get_kernel, has_kernel, install_kernel -from kernels.archs import ( - UnsupportedArchError, - _arch_incompatibility, - _cuda_archs_support_capability, - check_arch_compatibility, - is_arch_compatible, -) +from kernels.archs import _arch_incompatibility, _cuda_archs_support_capability def make_metadata(backend_type: str, archs: list[str] | None) -> Metadata: @@ -81,26 +75,24 @@ class FakeProperties: def test_cuda_incompatible_arch_is_rejected(fake_cuda_device): - metadata = make_metadata("cuda", ["8.0", "9.0a"]) - assert not is_arch_compatible(metadata) - with pytest.raises(UnsupportedArchError, match=r"CUDA capability 10\.0.*8\.0, 9\.0a"): - check_arch_compatibility(metadata, "torch-stable-abi29-cu128-x86_64-linux") + reason = _arch_incompatibility(make_metadata("cuda", ["8.0", "9.0a"])) + assert reason is not None + assert "CUDA capability 10.0" in reason + assert "8.0, 9.0a" in reason def test_cuda_compatible_arch_is_accepted(fake_cuda_device): for archs in (["8.0", "10.0"], ["10.0a"], ["10.0f"], None): - metadata = make_metadata("cuda", archs) - assert is_arch_compatible(metadata) - check_arch_compatibility(metadata, "torch-stable-abi29-cu128-x86_64-linux") + assert _arch_incompatibility(make_metadata("cuda", archs)) is None def test_rocm_arch_check(fake_rocm_device): - assert is_arch_compatible(make_metadata("rocm", ["gfx90a", "gfx942"])) - assert is_arch_compatible(make_metadata("rocm", None)) - metadata = make_metadata("rocm", ["gfx942"]) - assert not is_arch_compatible(metadata) - with pytest.raises(UnsupportedArchError, match=r"ROCm arch gfx90a.*gfx942"): - check_arch_compatibility(metadata, "torch28-cxx11-rocm64-x86_64-linux") + assert _arch_incompatibility(make_metadata("rocm", ["gfx90a", "gfx942"])) is None + assert _arch_incompatibility(make_metadata("rocm", None)) is None + reason = _arch_incompatibility(make_metadata("rocm", ["gfx942"])) + assert reason is not None + assert "ROCm arch gfx90a" in reason + assert "gfx942" in reason def test_check_skipped_without_device(monkeypatch): @@ -111,8 +103,8 @@ def test_check_skipped_without_device(monkeypatch): def test_check_skipped_for_backends_without_archs(fake_cuda_device): # Archs of other backends cannot be checked against the current device. - assert is_arch_compatible(make_metadata("cpu", None)) - assert is_arch_compatible(make_metadata("metal", ["applegpu_g13"])) + assert _arch_incompatibility(make_metadata("cpu", None)) is None + assert _arch_incompatibility(make_metadata("metal", ["applegpu_g13"])) is None @pytest.mark.cuda_only @@ -124,7 +116,7 @@ def test_get_kernel_rejects_unsupported_capability(monkeypatch): monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (99, 9)) - with pytest.raises(UnsupportedArchError): + 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) From 63234c4e54b2d951a0f882cddce6c75a50534cb1 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 18 Aug 2026 20:11:02 +0200 Subject: [PATCH 3/7] add a minimal test to mimic remy's issue: https://github.com/huggingface/kernels/issues/707 --- kernels/tests/test_archs.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py index 27d1cabfe..0645f2346 100644 --- a/kernels/tests/test_archs.py +++ b/kernels/tests/test_archs.py @@ -107,6 +107,15 @@ def test_check_skipped_for_backends_without_archs(fake_cuda_device): assert _arch_incompatibility(make_metadata("metal", ["applegpu_g13"])) is None +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"]) + assert _arch_incompatibility(metadata) is not None + + @pytest.mark.cuda_only def test_get_kernel_rejects_unsupported_capability(monkeypatch): variant_path = install_kernel("kernels-community/relu", revision="v1") From 487024b337f779db07332ddb833c4f616ea95bb7 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 19 Aug 2026 15:15:48 +0530 Subject: [PATCH 4/7] ptx --- kernels/src/kernels/archs.py | 32 ++++++++++++++++++++++++-------- kernels/tests/test_archs.py | 12 ++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/kernels/src/kernels/archs.py b/kernels/src/kernels/archs.py index 2a93ba792..26903a8f3 100644 --- a/kernels/src/kernels/archs.py +++ b/kernels/src/kernels/archs.py @@ -3,13 +3,19 @@ from kernels.compat import has_torch -def _parse_cuda_arch(arch: str) -> tuple[int, int, str] | None: - """Parse a CUDA arch into its capability and suffix. +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. - Returns `None` when the arch string is not in this format. + 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] @@ -19,7 +25,7 @@ def _parse_cuda_arch(arch: str) -> tuple[int, int, str] | None: if not sep or not arch_major.isdigit() or not arch_minor.isdigit(): return None - return int(arch_major), int(arch_minor), suffix + return int(arch_major), int(arch_minor), suffix, ptx def _cuda_arch_supports(arch: str, capability: tuple[int, int]) -> bool | None: @@ -31,15 +37,25 @@ def _cuda_arch_supports(arch: str, capability: tuple[int, int]) -> bool | None: if parsed is None: return None - arch_major, arch_minor, suffix = parsed + 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) - # Base and family-specific builds run on capabilities of the same - # generation with the same or a newer minor version. + 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 diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py index 0645f2346..bdff9db2e 100644 --- a/kernels/tests/test_archs.py +++ b/kernels/tests/test_archs.py @@ -43,6 +43,18 @@ def make_metadata(backend_type: str, archs: list[str] | None) -> Metadata: (["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 From 5f1e7e9a78158e000af4c6aa7d165416401268e5 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 19 Aug 2026 15:29:42 +0530 Subject: [PATCH 5/7] descriptive name. --- kernels/src/kernels/archs.py | 4 ++-- kernels/tests/test_archs.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/kernels/src/kernels/archs.py b/kernels/src/kernels/archs.py index 26903a8f3..21803db7d 100644 --- a/kernels/src/kernels/archs.py +++ b/kernels/src/kernels/archs.py @@ -59,7 +59,7 @@ def _cuda_arch_supports(arch: str, capability: tuple[int, int]) -> bool | None: return major == arch_major and minor >= arch_minor -def _cuda_archs_support_capability(archs: list[str], capability: tuple[int, int]) -> bool: +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 @@ -81,7 +81,7 @@ def _arch_incompatibility(metadata: Metadata) -> str | None: if torch.version.cuda is None or not torch.cuda.is_available(): return None major, minor = torch.cuda.get_device_capability() - if _cuda_archs_support_capability(archs, (major, minor)): + if _supports_cuda_capability(archs, (major, minor)): return None return ( f"CUDA capability {major}.{minor} of the current device is not " diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py index bdff9db2e..50065c68e 100644 --- a/kernels/tests/test_archs.py +++ b/kernels/tests/test_archs.py @@ -5,7 +5,7 @@ from kernels_data import Metadata from kernels import get_kernel, has_kernel, install_kernel -from kernels.archs import _arch_incompatibility, _cuda_archs_support_capability +from kernels.archs import _arch_incompatibility, _supports_cuda_capability def make_metadata(backend_type: str, archs: list[str] | None) -> Metadata: @@ -62,8 +62,8 @@ def make_metadata(backend_type: str, archs: list[str] | None) -> Metadata: (["garbage"], (9, 0), True), ], ) -def test_cuda_archs_support_capability(archs, capability, supported): - assert _cuda_archs_support_capability(archs, capability) == supported +def test_supports_cuda_capability(archs, capability, supported): + assert _supports_cuda_capability(archs, capability) == supported @pytest.fixture From b120618deacf2e233ef7fb567339cd6da8bc22da Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 19 Aug 2026 15:38:22 +0530 Subject: [PATCH 6/7] raise and remove unpythonic bits. --- kernels/src/kernels/archs.py | 50 ++++++++++++++++++++++-------------- kernels/src/kernels/utils.py | 20 ++++++--------- kernels/tests/test_archs.py | 34 ++++++++++++------------ 3 files changed, 57 insertions(+), 47 deletions(-) diff --git a/kernels/src/kernels/archs.py b/kernels/src/kernels/archs.py index 21803db7d..a9346f588 100644 --- a/kernels/src/kernels/archs.py +++ b/kernels/src/kernels/archs.py @@ -68,36 +68,48 @@ def _supports_cuda_capability(archs: list[str], capability: tuple[int, int]) -> return any(supports) -def _arch_incompatibility(metadata: Metadata) -> str | None: - """Check a kernel build against the architecture of the current device.""" +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 None + 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 None + return major, minor = torch.cuda.get_device_capability() - if _supports_cuda_capability(archs, (major, minor)): - return None - return ( - f"CUDA capability {major}.{minor} of the current device is not " - f"supported by the architectures of the build: {', '.join(archs)}" - ) + 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 None + 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 in archs: - return None - return ( - f"ROCm arch {gcn_arch} of the current device is not supported " - f"by the architectures of the build: {', '.join(archs)}" - ) - - return None + 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 e8fd959ad..03ff9ac1e 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -19,7 +19,7 @@ from kernels._system import glibc_version from kernels._versions import select_revision_or_version -from kernels.archs import _arch_incompatibility +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 @@ -524,16 +524,7 @@ def get_kernel( ) if check_arch: metadata = Metadata.read_from_file(variant_path / "metadata.json") - reason = _arch_incompatibility(metadata) - if reason is not None: - raise RuntimeError( - f"Kernel '{metadata.name}' variant '{variant_path.name}' does not " - f"support the 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." - ) + _check_arch_incompatibility(metadata, variant_path.name) return _import_from_path(variant_path, repo_info=repo_info) @@ -632,7 +623,12 @@ def has_kernel( except EntryNotFoundError: return False - return _arch_incompatibility(Metadata.read_from_file(metadata_path)) is None + 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 index 50065c68e..733f19449 100644 --- a/kernels/tests/test_archs.py +++ b/kernels/tests/test_archs.py @@ -5,7 +5,7 @@ from kernels_data import Metadata from kernels import get_kernel, has_kernel, install_kernel -from kernels.archs import _arch_incompatibility, _supports_cuda_capability +from kernels.archs import _check_arch_incompatibility, _supports_cuda_capability def make_metadata(backend_type: str, archs: list[str] | None) -> Metadata: @@ -87,36 +87,37 @@ class FakeProperties: def test_cuda_incompatible_arch_is_rejected(fake_cuda_device): - reason = _arch_incompatibility(make_metadata("cuda", ["8.0", "9.0a"])) - assert reason is not None - assert "CUDA capability 10.0" in reason - assert "8.0, 9.0a" in reason + 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"], None): - assert _arch_incompatibility(make_metadata("cuda", archs)) is None + _check_arch_incompatibility(make_metadata("cuda", archs), "test-variant") def test_rocm_arch_check(fake_rocm_device): - assert _arch_incompatibility(make_metadata("rocm", ["gfx90a", "gfx942"])) is None - assert _arch_incompatibility(make_metadata("rocm", None)) is None - reason = _arch_incompatibility(make_metadata("rocm", ["gfx942"])) - assert reason is not None - assert "ROCm arch gfx90a" in reason - assert "gfx942" in reason + _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) - assert _arch_incompatibility(make_metadata("cuda", ["8.0"])) is None + _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. - assert _arch_incompatibility(make_metadata("cpu", None)) is None - assert _arch_incompatibility(make_metadata("metal", ["applegpu_g13"])) is None + _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): @@ -125,7 +126,8 @@ def test_issue_707_fa3_on_b200(fake_cuda_device): # 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"]) - assert _arch_incompatibility(metadata) is not None + with pytest.raises(RuntimeError, match="does not support the current device"): + _check_arch_incompatibility(metadata, "test-variant") @pytest.mark.cuda_only From ddeee385d2a00be4dc31c69c0f156f11a3bd7d41 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 19 Aug 2026 15:58:09 +0530 Subject: [PATCH 7/7] test --- kernels/tests/test_archs.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/kernels/tests/test_archs.py b/kernels/tests/test_archs.py index 733f19449..fa57843d3 100644 --- a/kernels/tests/test_archs.py +++ b/kernels/tests/test_archs.py @@ -95,10 +95,18 @@ def test_cuda_incompatible_arch_is_rejected(fake_cuda_device): def test_cuda_compatible_arch_is_accepted(fake_cuda_device): - for archs in (["8.0", "10.0"], ["10.0a"], ["10.0f"], None): + 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")