kernels 0.12.3__tar.gz → 0.14.0.dev0__tar.gz

{kernels-0.12.3 → kernels-0.14.0.dev0}/PKG-INFO

@@ -1,20 +1,22 @@
 Metadata-Version: 2.4
 Name: kernels
-Version: 0.12.3
+Version: 0.14.0.dev0
 Summary: Download compute kernels
 Author-email: Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>
 License: Apache-2.0
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: huggingface_hub<2.0,>=0.26.0
+Requires-Dist: huggingface-hub>=1.10.0
 Requires-Dist: packaging>=20.0
 Requires-Dist: pyyaml>=6
 Requires-Dist: tomli>=2.0; python_version < "3.11"
+Requires-Dist: tomlkit>=0.13.3
 Provides-Extra: abi-check
-Requires-Dist: kernel-abi-check<0.13.0,>=0.12.0; extra == "abi-check"
+Requires-Dist: kernel-abi-check<0.7.0,>=0.6.2; extra == "abi-check"
 Provides-Extra: benchmark
+Requires-Dist: matplotlib>=3.7.0; extra == "benchmark"
 Requires-Dist: numpy>=2.0.2; extra == "benchmark"
-Requires-Dist: requests>=2.32.5; extra == "benchmark"
+Requires-Dist: tabulate>=0.9.0; extra == "benchmark"
 Requires-Dist: torch; extra == "benchmark"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"

{kernels-0.12.3 → kernels-0.14.0.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.12.3"
+version = "0.14.0.dev0"
 description = "Download compute kernels"
 authors = [
   { name = "Daniel de Kok", email = "daniel@huggingface.co" },
@@ -8,12 +8,13 @@ authors = [
 ]
 license = { text = "Apache-2.0" }
 readme = "README.md"
-requires-python = ">= 3.9"
+requires-python = ">= 3.10"
 dependencies = [
-  "huggingface_hub>=0.26.0,<2.0",
+  "huggingface-hub>=1.10.0",
   "packaging>=20.0",
   "pyyaml>=6",
   "tomli>=2.0; python_version<'3.11'",
+  "tomlkit>=0.13.3",
 ]
 
 [build-system]
@@ -28,15 +29,17 @@ dev = [
   # Whatever version is compatible with pytest.
   "pytest-benchmark",
   "torch>=2.5",
+  "apache-tvm-ffi>=0.1.9,<0.2.0",
   "types-pyyaml",
-  "types-requests",
+  "types-tabulate",
 ]
 
 [project.optional-dependencies]
-abi-check = ["kernel-abi-check>=0.12.0,<0.13.0"]
+abi-check = ["kernel-abi-check>=0.6.2,<0.7.0"]
 benchmark = [
+  "matplotlib>=3.7.0",
   "numpy>=2.0.2",
-  "requests>=2.32.5",
+  "tabulate>=0.9.0",
   "torch",
 ]
 torch = ["torch"]
@@ -53,11 +56,10 @@ kernels = "kernels.cli:main"
 [tool.setuptools.package-data]
 kernels = ["python_depends.json"]
 
-[tool.isort]
-profile = "black"
-line_length = 119
-
 [tool.ruff]
+# If the version is changed, apply the change in the Nix overlay
+# as well.
+required-version = "==0.15.10"
 exclude = [
   ".eggs",
   ".git",
@@ -82,4 +84,6 @@ line-length = 119
 # Ignored rules:
 # "E501" -> line length violation
 lint.ignore = ["E501"]
-lint.select = ["E", "F", "W"]
+lint.select = ["E", "F", "I", "W"]
+
+[tool.ruff.format]

{kernels-0.12.3 → kernels-0.14.0.dev0}/src/kernels/__init__.py

@@ -2,6 +2,8 @@ import importlib.metadata
 
 __version__ = importlib.metadata.version("kernels")
 
+from kernels._windows import _add_additional_dll_paths
+from kernels.benchmark import Benchmark
 from kernels.layer import (
     CUDAProperties,
     Device,
@@ -21,16 +23,13 @@ from kernels.layer import (
 )
 from kernels.utils import (
     get_kernel,
+    get_loaded_kernels,
     get_local_kernel,
     get_locked_kernel,
     has_kernel,
     install_kernel,
     load_kernel,
 )
-from kernels.benchmark import Benchmark
-
-
-from kernels._windows import _add_additional_dll_paths
 
 _add_additional_dll_paths()
 
@@ -47,6 +46,7 @@ __all__ = [
     "LockedLayerRepository",
     "Mode",
     "get_kernel",
+    "get_loaded_kernels",
     "get_local_kernel",
     "get_locked_kernel",
     "has_kernel",

kernels-0.14.0.dev0/src/kernels/_versions.py

@@ -0,0 +1,72 @@
+import logging
+import warnings
+
+from huggingface_hub.hf_api import GitRefInfo
+
+logger = logging.getLogger(__name__)
+
+
+def _get_available_versions(repo_id: str) -> dict[int, GitRefInfo]:
+    """Get kernel versions that are available in the repository."""
+    from kernels.utils import _get_hf_api
+
+    refs = _get_hf_api().list_repo_refs(repo_id=repo_id, repo_type="kernel")
+
+    versions = {}
+    for branch in refs.branches:
+        if not branch.name.startswith("v"):
+            continue
+        try:
+            versions[int(branch.name[1:])] = branch
+        except ValueError:
+            continue
+
+    return versions
+
+
+def resolve_version_spec_as_ref(repo_id: str, version_spec: int) -> GitRefInfo:
+    """
+    Get the ref for a kernel with the given version.
+    """
+    versions = _get_available_versions(repo_id)
+
+    ref = versions.get(version_spec, None)
+    if ref is None:
+        raise ValueError(
+            f"Version {version_spec} not found, available versions: {', '.join(sorted(str(v) for v in versions.keys()))}"
+        )
+
+    latest_version = max(versions.keys())
+    if version_spec < latest_version:
+        logger.warning(
+            "You are using version %d of '%s', but version %d is available.",
+            version_spec,
+            repo_id,
+            latest_version,
+        )
+
+    return ref
+
+
+def select_revision_or_version(
+    repo_id: str,
+    *,
+    revision: str | None,
+    version: int | None,
+) -> str:
+    if revision is not None and version is not None:
+        raise ValueError("Only one of `revision` or `version` must be specified.")
+
+    if revision is not None:
+        return revision
+    elif version is not None:
+        return resolve_version_spec_as_ref(repo_id, version).target_commit
+
+    warnings.warn(
+        "Future versions of `kernels` (>=0.15) will require specifying a kernel version or revision. "
+        "See: https://huggingface.co/docs/kernels/migration",
+        FutureWarning,
+        stacklevel=2,
+    )
+
+    return "main"

kernels-0.14.0.dev0/src/kernels/backends.py

@@ -0,0 +1,276 @@
+import ctypes
+import ctypes.util
+import re
+import warnings
+from dataclasses import dataclass
+from typing import ClassVar, Optional, Protocol, runtime_checkable
+
+from huggingface_hub.dataclasses import strict
+from packaging.version import Version
+
+from kernels.compat import has_torch
+
+
+@runtime_checkable
+class Backend(Protocol):
+    @property
+    def name(self) -> str:
+        """
+        Short name of the backend, e.g. "cuda", "rocm", "cpu", etc.
+        """
+        ...
+
+    @property
+    def variant_str(self) -> str:
+        """
+        The name of the backend as used in a build variant, e.g. `cu128`
+        for CUDA 12.8.
+        """
+        ...
+
+
+@dataclass(unsafe_hash=True)
+class CANN:
+    _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"cann(\d+)(\d+)")
+
+    version: Version
+
+    @property
+    def name(self) -> str:
+        return "cann"
+
+    @property
+    def variant_str(self) -> str:
+        return f"cann{self.version.major}{self.version.minor}"
+
+    @staticmethod
+    def parse(s: str) -> "CANN":
+        m = CANN._VARIANT_REGEX.fullmatch(s)
+        if not m:
+            raise ValueError(f"Invalid CANN variant string: {s!r}")
+        return CANN(version=Version(f"{m.group(1)}.{m.group(2)}"))
+
+
+@strict
+@dataclass(unsafe_hash=True)
+class CPU:
+    @property
+    def name(self) -> str:
+        return "cpu"
+
+    @property
+    def variant_str(self) -> str:
+        return "cpu"
+
+    @staticmethod
+    def parse(s: str) -> "CPU":
+        if s != "cpu":
+            raise ValueError(f"Invalid CPU variant string: {s!r}")
+        return CPU()
+
+
+@dataclass(unsafe_hash=True)
+class CUDA:
+    _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"cu(\d+)(\d+)")
+
+    version: Version
+
+    @property
+    def name(self) -> str:
+        return "cuda"
+
+    @property
+    def variant_str(self) -> str:
+        return f"cu{self.version.major}{self.version.minor}"
+
+    @staticmethod
+    def parse(s: str) -> "CUDA":
+        m = CUDA._VARIANT_REGEX.fullmatch(s)
+        if not m:
+            raise ValueError(f"Invalid CUDA variant string: {s!r}")
+        return CUDA(version=Version(f"{m.group(1)}.{m.group(2)}"))
+
+
+@strict
+@dataclass(unsafe_hash=True)
+class Metal:
+    @property
+    def name(self) -> str:
+        return "metal"
+
+    @property
+    def variant_str(self) -> str:
+        return "metal"
+
+    @staticmethod
+    def parse(s: str) -> "Metal":
+        if s != "metal":
+            raise ValueError(f"Invalid Metal variant string: {s!r}")
+        return Metal()
+
+
+@strict
+@dataclass(unsafe_hash=True)
+class Neuron:
+    @property
+    def name(self) -> str:
+        return "neuron"
+
+    @property
+    def variant_str(self) -> str:
+        return "neuron"
+
+    @staticmethod
+    def parse(s: str) -> "Neuron":
+        if s != "neuron":
+            raise ValueError(f"Invalid Neuron variant string: {s!r}")
+        return Neuron()
+
+
+@dataclass(unsafe_hash=True)
+class ROCm:
+    _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"rocm(\d+)(\d+)")
+
+    version: Version
+
+    @property
+    def name(self) -> str:
+        return "rocm"
+
+    @property
+    def variant_str(self) -> str:
+        return f"rocm{self.version.major}{self.version.minor}"
+
+    @staticmethod
+    def parse(s: str) -> "ROCm":
+        m = ROCm._VARIANT_REGEX.fullmatch(s)
+        if not m:
+            raise ValueError(f"Invalid ROCm variant string: {s!r}")
+        return ROCm(version=Version(f"{m.group(1)}.{m.group(2)}"))
+
+
+@dataclass(unsafe_hash=True)
+class XPU:
+    _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"xpu(\d+)(\d+)")
+
+    version: Version
+
+    @property
+    def name(self) -> str:
+        return "xpu"
+
+    @property
+    def variant_str(self) -> str:
+        return f"xpu{self.version.major}{self.version.minor}"
+
+    @staticmethod
+    def parse(s: str) -> "XPU":
+        m = XPU._VARIANT_REGEX.fullmatch(s)
+        if not m:
+            raise ValueError(f"Invalid XPU variant string: {s!r}")
+        return XPU(version=Version(f"{m.group(1)}.{m.group(2)}"))
+
+
+def parse_backend(s: str) -> Backend:
+    """Parse a backend variant string (e.g. 'cu128', 'rocm61', 'cpu') into a Backend."""
+    if s == "cpu":
+        return CPU.parse(s)
+    elif s == "metal":
+        return Metal.parse(s)
+    elif s == "neuron":
+        return Neuron.parse(s)
+    elif s.startswith("cu"):
+        return CUDA.parse(s)
+    elif s.startswith("rocm"):
+        return ROCm.parse(s)
+    elif s.startswith("xpu"):
+        return XPU.parse(s)
+    elif s.startswith("cann"):
+        return CANN.parse(s)
+    else:
+        raise ValueError(f"Unknown backend variant string: {s!r}")
+
+
+def _backend() -> Backend:
+    if has_torch:
+        import torch
+
+        if hasattr(torch, "neuron"):
+            # Needs to be sorted before specific Torch builds, since Neuron
+            # extension can be loaded into e.g. CUDA Torch builds.
+            return Neuron()
+        elif torch.version.cuda is not None:
+            cuda_version = Version(torch.version.cuda)
+            return CUDA(version=cuda_version)
+        elif torch.version.hip is not None:
+            rocm_version = Version(torch.version.hip.split("-")[0])
+            return ROCm(version=rocm_version)
+        elif torch.backends.mps.is_available():
+            return Metal()
+        elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+            version = f"{torch.version.xpu[0:4]}.{torch.version.xpu[5:6]}"
+            return XPU(version=Version(version))
+        elif _get_torch_privateuse_backend_name() == "npu":
+            from torch_npu.utils.collect_env import get_cann_version  # type: ignore[import-not-found]
+
+            cann_major, cann_minor = get_cann_version()[0], get_cann_version()[2]
+            return CANN(version=Version(f"{cann_major}.{cann_minor}"))
+        else:
+            return CPU()
+    else:
+        cuda = _get_cuda()
+        if cuda is not None:
+            return cuda
+
+        return CPU()
+
+
+def _get_torch_privateuse_backend_name() -> str | None:
+    import torch
+
+    if hasattr(torch._C, "_get_privateuse1_backend_name"):
+        return torch._C._get_privateuse1_backend_name()
+    return None
+
+
+def _select_backend(backend: str | None) -> Backend:
+    if backend is None:
+        return _backend()
+
+    supported = _supported_backends()
+    if backend in supported:
+        return supported[backend]
+
+    raise ValueError(f"Invalid backend '{backend}', system supported backends: {', '.join(sorted(supported.keys()))}")
+
+
+def _supported_backends() -> dict[str, Backend]:
+    backend = _backend()
+    return {"cpu": CPU(), backend.name: backend}
+
+
+def _get_cuda() -> Optional[CUDA]:
+    """
+    Get CUDA runtime library information.
+    """
+    lib_name = ctypes.util.find_library("cudart")
+    if lib_name is None:
+        return None
+
+    try:
+        libcudart = ctypes.CDLL(lib_name)
+    except OSError:
+        return None
+
+    runtime_version = ctypes.c_int(0)
+    result = libcudart.cudaRuntimeGetVersion(ctypes.byref(runtime_version))
+    if result != 0:
+        warnings.warn("System has CUDA runtime library, but cannot get runtime version.")
+        return None
+
+    # cudaRuntimeGetVersion encodes the version as (major * 1000 + minor * 10).
+    version_int = runtime_version.value
+    major = version_int // 1000
+    minor = (version_int % 1000) // 10
+
+    return CUDA(version=Version(f"{major}.{minor}"))

kernels-0.14.0.dev0/src/kernels/benchmark.py

@@ -0,0 +1,38 @@
+from typing import Any
+
+
+class Benchmark:
+    """Base class for kernel benchmarks.
+
+    Subclass this to create a benchmark script with automatic timing,
+    verification, and reproducibility support. The kernel is loaded
+    automatically from the repo_id specified in the CLI command.
+
+    Example:
+        class MyBenchmark(Benchmark):
+            seed = 42
+
+            def setup(self):
+                self.x = torch.randn(128, 1024, device=self.device, dtype=torch.float16)
+                self.out = torch.empty(128, 512, device=self.device, dtype=torch.float16)
+
+            def benchmark_silu(self):
+                self.kernel.silu_and_mul(self.out, self.x)
+
+            def verify_silu(self) -> torch.Tensor:
+                # Return reference tensor; runner compares with self.out
+                return torch.nn.functional.silu(self.x[..., :512]) * self.x[..., 512:]
+
+    Run with: kernels benchmark <repo_id>
+    """
+
+    seed: int | None = None  # Optional: seed for reproducibility
+    device: str = "cpu"  # Set automatically by runner
+
+    def __init__(self) -> None:
+        self.kernel: Any = None
+        self.out: Any = None  # Output tensor, set by setup methods
+
+    def setup(self) -> None:
+        """Override to set up tensors as instance attributes."""
+        pass

kernels-0.14.0.dev0/src/kernels/benchmarks/activation.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn.functional as F
+
+from kernels.benchmark import Benchmark
+
+
+class SiluAndMulBenchmark(Benchmark):
+    seed: int = 42
+
+    # Workload: small
+    def setup_small(self):
+        self.x = torch.randn(8, 1024, 2048, device=self.device, dtype=torch.float16)
+        self.out = torch.empty(8, 1024, 1024, device=self.device, dtype=torch.float16)
+
+    def benchmark_small(self):
+        self.kernel.silu_and_mul(self.out, self.x)
+
+    def verify_small(self) -> torch.Tensor:
+        d = self.x.shape[-1] // 2
+        return F.silu(self.x[..., :d]) * self.x[..., d:]
+
+    # Workload: medium
+    def setup_medium(self):
+        self.x = torch.randn(8, 2048, 4096, device=self.device, dtype=torch.float16)
+        self.out = torch.empty(8, 2048, 2048, device=self.device, dtype=torch.float16)
+
+    def benchmark_medium(self):
+        self.kernel.silu_and_mul(self.out, self.x)
+
+    def verify_medium(self) -> torch.Tensor:
+        d = self.x.shape[-1] // 2
+        return F.silu(self.x[..., :d]) * self.x[..., d:]
+
+    # Workload: large
+    def setup_large(self):
+        self.x = torch.randn(8, 4096, 8192, device=self.device, dtype=torch.float16)
+        self.out = torch.empty(8, 4096, 4096, device=self.device, dtype=torch.float16)
+
+    def benchmark_large(self):
+        self.kernel.silu_and_mul(self.out, self.x)
+        self.kernel.silu_and_mul(self.out, self.x)
+
+    def verify_large(self) -> torch.Tensor:
+        d = self.x.shape[-1] // 2
+        return F.silu(self.x[..., :d]) * self.x[..., d:]
+
+
+class GeluAndMulBenchmark(Benchmark):
+    seed: int = 42
+
+    # Workload: small
+    def setup_small(self):
+        self.x = torch.randn(8, 1024, 2048, device=self.device, dtype=torch.float16)
+        self.out = torch.empty(8, 1024, 1024, device=self.device, dtype=torch.float16)
+
+    def benchmark_small(self):
+        self.kernel.gelu_and_mul(self.out, self.x)
+
+    def verify_small(self) -> torch.Tensor:
+        d = self.x.shape[-1] // 2
+        return F.gelu(self.x[..., :d]) * self.x[..., d:]
+
+    # Workload: medium
+    def setup_medium(self):
+        self.x = torch.randn(8, 2048, 4096, device=self.device, dtype=torch.float16)
+        self.out = torch.empty(8, 2048, 2048, device=self.device, dtype=torch.float16)
+
+    def benchmark_medium(self):
+        self.kernel.gelu_and_mul(self.out, self.x)
+
+    def verify_medium(self) -> torch.Tensor:
+        d = self.x.shape[-1] // 2
+        return F.gelu(self.x[..., :d]) * self.x[..., d:]
+
+    # Workload: large
+    def setup_large(self):
+        self.x = torch.randn(8, 4096, 8192, device=self.device, dtype=torch.float16)
+        self.out = torch.empty(8, 4096, 4096, device=self.device, dtype=torch.float16)
+
+    def benchmark_large(self):
+        self.kernel.gelu_and_mul(self.out, self.x)
+
+    def verify_large(self) -> torch.Tensor:
+        d = self.x.shape[-1] // 2
+        return F.gelu(self.x[..., :d]) * self.x[..., d:]