kernels 0.5.0.dev0__tar.gz → 0.5.1__tar.gz

{kernels-0.5.0.dev0/src/kernels.egg-info → kernels-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernels
-Version: 0.5.0.dev0
+Version: 0.5.1
 Summary: Download compute kernels
 Author-email: OlivierDehaene <olivier@huggingface.co>, Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>, Nicolas Patry <nicolas@huggingface.co>
 License: Apache-2.0

{kernels-0.5.0.dev0 → kernels-0.5.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.5.0.dev0"
+version = "0.5.1"
 description = "Download compute kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },

{kernels-0.5.0.dev0 → kernels-0.5.1}/src/kernels/utils.py

@@ -54,11 +54,29 @@ def build_variant() -> str:
     return f"torch{torch_version.major}{torch_version.minor}-{cxxabi}-{compute_framework}-{cpu}-{os}"
 
 
-def universal_build_variant() -> str:
+def build_variant_noarch() -> str:
+    import torch
+
+    if torch.version.cuda is not None:
+        return "torch-cuda"
+    elif torch.version.hip is not None:
+        return "torch-rocm"
+    elif torch.backends.mps.is_available():
+        return "torch-metal"
+    else:
+        return "torch-cpu"
+
+
+def build_variant_universal() -> str:
     # Once we support other frameworks, detection goes here.
     return "torch-universal"
 
 
+def build_variants() -> List[str]:
+    """Return compatible build variants in preferred order."""
+    return [build_variant(), build_variant_noarch(), build_variant_universal()]
+
+
 def import_from_path(module_name: str, file_path: Path) -> ModuleType:
     # We cannot use the module name as-is, after adding it to `sys.modules`,
     # it would also be used for other imports. So, we make a module name that
@@ -89,25 +107,32 @@ def install_kernel(
     The output path is validated againt `hash` when set.
     """
     package_name = package_name_from_repo_id(repo_id)
-    variant = build_variant()
-    universal_variant = universal_build_variant()
+    allow_patterns = [f"build/{variant}/*" for variant in build_variants()]
     repo_path = Path(
         snapshot_download(
             repo_id,
-            allow_patterns=[f"build/{variant}/*", f"build/{universal_variant}/*"],
+            allow_patterns=allow_patterns,
             cache_dir=CACHE_DIR,
             revision=revision,
             local_files_only=local_files_only,
         )
     )
 
-    variant_path = repo_path / "build" / variant
-    universal_variant_path = repo_path / "build" / universal_variant
+    variants = build_variants()
+    variant = None
+    variant_path = None
+    for candidate_variant in variants:
+        variant_path = repo_path / "build" / candidate_variant
+        if variant_path.exists():
+            variant = candidate_variant
+            break
 
-    if not variant_path.exists() and universal_variant_path.exists():
-        # Fall back to universal variant.
-        variant = universal_variant
-        variant_path = universal_variant_path
+    if variant is None:
+        raise FileNotFoundError(
+            f"Kernel at path `{repo_path}` does not have one of build variants: {', '.join(variants)}"
+        )
+
+    assert variant_path is not None
 
     if variant_locks is not None:
         variant_lock = variant_locks.get(variant)
@@ -167,21 +192,16 @@ def has_kernel(repo_id: str, revision: str = "main") -> bool:
     (Torch version and compute framework).
     """
     package_name = package_name_from_repo_id(repo_id)
-    variant = build_variant()
-    universal_variant = universal_build_variant()
-
-    if file_exists(
-        repo_id,
-        revision=revision,
-        filename=f"build/{universal_variant}/{package_name}/__init__.py",
-    ):
-        return True
-
-    return file_exists(
-        repo_id,
-        revision=revision,
-        filename=f"build/{variant}/{package_name}/__init__.py",
-    )
+
+    for variant in build_variants():
+        if file_exists(
+            repo_id,
+            revision=revision,
+            filename=f"build/{variant}/{package_name}/__init__.py",
+        ):
+            return True
+
+    return False
 
 
 def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
@@ -204,33 +224,29 @@ def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
 
     package_name = package_name_from_repo_id(repo_id)
 
-    variant = build_variant()
-    universal_variant = universal_build_variant()
-
+    allow_patterns = [f"build/{variant}/*" for variant in build_variants()]
     repo_path = Path(
         snapshot_download(
             repo_id,
-            allow_patterns=[f"build/{variant}/*", f"build/{universal_variant}/*"],
+            allow_patterns=allow_patterns,
             cache_dir=CACHE_DIR,
             revision=locked_sha,
             local_files_only=True,
         )
     )
 
-    variant_path = repo_path / "build" / variant
-    universal_variant_path = repo_path / "build" / universal_variant
-    if not variant_path.exists() and universal_variant_path.exists():
-        # Fall back to universal variant.
-        variant = universal_variant
-        variant_path = universal_variant_path
-
-    module_init_path = variant_path / package_name / "__init__.py"
-    if not os.path.exists(module_init_path):
-        raise FileNotFoundError(
-            f"Locked kernel `{repo_id}` does not have build `{variant}` or was not downloaded with `kernels download <project>`"
-        )
+    for variant in build_variants():
+        variant_path = repo_path / "build" / variant
+        module_init_path = variant_path / package_name / "__init__.py"
+        if module_init_path.exists():
+            module_init_path = variant_path / package_name / "__init__.py"
+            return import_from_path(
+                package_name, variant_path / package_name / "__init__.py"
+            )
 
-    return import_from_path(package_name, variant_path / package_name / "__init__.py")
+    raise FileNotFoundError(
+        f"Locked kernel `{repo_id}` does not have applicable variant or was not downloaded with `kernels download <project>`"
+    )
 
 
 def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleType:

{kernels-0.5.0.dev0 → kernels-0.5.1/src/kernels.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernels
-Version: 0.5.0.dev0
+Version: 0.5.1
 Summary: Download compute kernels
 Author-email: OlivierDehaene <olivier@huggingface.co>, Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>, Nicolas Patry <nicolas@huggingface.co>
 License: Apache-2.0

{kernels-0.5.0.dev0 → kernels-0.5.1}/tests/test_basic.py

@@ -6,7 +6,7 @@ from kernels import get_kernel, has_kernel
 
 @pytest.fixture
 def kernel():
-    return get_kernel("kernels-community/activation")
+    return get_kernel("kernels-community/activation", revision="v0.0.3")
 
 
 @pytest.fixture
@@ -39,7 +39,7 @@ def test_gelu_fast(kernel, device):
 @pytest.mark.parametrize(
     "kernel_exists",
     [
-        ("kernels-community/activation", "main", True),
+        ("kernels-community/activation", "v0.0.3", True),
         ("kernels-community/triton-layer-norm", "main", True),
         # Repo only contains Torch 2.4 kernels (and we don't
         # support/test against this version).
@@ -64,3 +64,10 @@ def test_universal_kernel(universal_kernel):
     out_check = out_check.to(torch.float16)
 
     torch.testing.assert_close(out, out_check, rtol=1e-1, atol=1e-1)
+
+
+def test_noarch_kernel(device):
+    supported_devices = ["cpu", "cuda"]
+    if device not in supported_devices:
+        pytest.skip(f"Device is not one of: {','.join(supported_devices)}")
+    get_kernel("kernels-test/silu-and-mul-noarch")

{kernels-0.5.0.dev0 → kernels-0.5.1}/tests/test_benchmarks.py

@@ -6,7 +6,7 @@ from kernels import get_kernel
 
 @pytest.fixture
 def kernel():
-    return get_kernel("kernels-community/activation")
+    return get_kernel("kernels-community/activation", revision="v0.0.3")
 
 
 @pytest.fixture