kernels 0.14.0.dev0__tar.gz → 0.14.1__tar.gz

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

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: kernels
-Version: 0.14.0.dev0
+Version: 0.14.1
 Summary: Download compute kernels
 Author-email: Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>
 License: Apache-2.0
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: huggingface-hub>=1.10.0
+Requires-Dist: kernels-data>=0.14.0.dev1
 Requires-Dist: packaging>=20.0
 Requires-Dist: pyyaml>=6
 Requires-Dist: tomli>=2.0; python_version < "3.11"

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

@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.14.0.dev0"
+version = "0.14.1"
 description = "Download compute kernels"
 authors = [
   { name = "Daniel de Kok", email = "daniel@huggingface.co" },
@@ -11,6 +11,7 @@ readme = "README.md"
 requires-python = ">= 3.10"
 dependencies = [
   "huggingface-hub>=1.10.0",
+  "kernels-data>=0.14.0.dev1",
   "packaging>=20.0",
   "pyyaml>=6",
   "tomli>=2.0; python_version<'3.11'",

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

@@ -2,6 +2,8 @@ import importlib.metadata
 
 __version__ = importlib.metadata.version("kernels")
 
+from kernels_data import Metadata
+
 from kernels._windows import _add_additional_dll_paths
 from kernels.benchmark import Benchmark
 from kernels.layer import (
@@ -22,6 +24,8 @@ from kernels.layer import (
     use_kernel_mapping,
 )
 from kernels.utils import (
+    LoadedKernel,
+    RepoInfo,
     get_kernel,
     get_loaded_kernels,
     get_local_kernel,
@@ -40,11 +44,14 @@ __all__ = [
     "Device",
     "FuncRepository",
     "LayerRepository",
+    "LoadedKernel",
     "LocalFuncRepository",
     "LocalLayerRepository",
     "LockedFuncRepository",
     "LockedLayerRepository",
+    "Metadata",
     "Mode",
+    "RepoInfo",
     "get_kernel",
     "get_loaded_kernels",
     "get_local_kernel",

{kernels-0.14.0.dev0 → kernels-0.14.1}/src/kernels/layer/func.py

@@ -64,12 +64,14 @@ class FuncRepository:
         func_name: str,
         revision: str | None = None,
         version: int | None = None,
+        trust_remote_code: bool | list[str] = False,
     ):
         if revision is not None and version is not None:
             raise ValueError("Either a revision or a version must be specified, not both.")
 
         self._repo_id = repo_id
         self.func_name = func_name
+        self._trust_remote_code = trust_remote_code
 
         # We are going to resolve these lazily, since we do not want
         # to do a network request for every registered FuncRepository.
@@ -85,7 +87,9 @@ class FuncRepository:
         )
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_kernel(self._repo_id, revision=self._resolve_revision())
+        kernel = get_kernel(
+            self._repo_id, revision=self._resolve_revision(), trust_remote_code=self._trust_remote_code
+        )
         return _get_kernel_func(self, kernel)
 
     def __eq__(self, other):
@@ -95,10 +99,11 @@ class FuncRepository:
             and self._repo_id == other._repo_id
             and self._revision == other._revision
             and self._version == other._version
+            and self._trust_remote_code == other._trust_remote_code
         )
 
     def __hash__(self):
-        return hash((self.func_name, self._repo_id, self._revision, self._version))
+        return hash((self.func_name, self._repo_id, self._revision, self._version, self._trust_remote_code))
 
     def __str__(self) -> str:
         return f"`{self._repo_id}` (revision: {self._resolve_revision()}), function `{self.func_name}`"
@@ -111,8 +116,6 @@ class LocalFuncRepository:
     Args:
         repo_path (`Path`):
             The local repository containing the layer.
-        package_name (`str`):
-            Package name of the kernel.
         func_name (`str`):
             The name of the function within the kernel repository.
 
@@ -125,7 +128,6 @@ class LocalFuncRepository:
         # Reference a specific layer by revision
         layer_repo = LocalFuncRepository(
             repo_path=Path("/home/daniel/kernels/activation"),
-            package_name="activation",
             func_name="silu_and_mul",
         )
         ```
@@ -135,15 +137,13 @@ class LocalFuncRepository:
         self,
         repo_path: Path,
         *,
-        package_name: str,
         func_name: str,
     ):
         self._repo_path = repo_path
-        self._package_name = package_name
         self.func_name = func_name
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_local_kernel(self._repo_path, self._package_name)
+        kernel = get_local_kernel(self._repo_path)
         return _get_kernel_func(self, kernel)
 
     def __eq__(self, other):
@@ -151,14 +151,13 @@ class LocalFuncRepository:
             isinstance(other, LocalFuncRepository)
             and self.func_name == other.func_name
             and self._repo_path == other._repo_path
-            and self._package_name == other._package_name
         )
 
     def __hash__(self):
-        return hash((self.func_name, self._repo_path, self._package_name))
+        return hash((self.func_name, self._repo_path))
 
     def __str__(self) -> str:
-        return f"`{self._repo_path}` (package: {self._package_name}), layer `{self.func_name}`"
+        return f"`{self._repo_path}` (layer `{self.func_name}`"
 
 
 def use_kernel_func_from_hub(func_name: str):
@@ -230,6 +229,7 @@ class LockedFuncRepository:
         *,
         lockfile: Path | None = None,
         func_name: str,
+        trust_remote_code: bool | list[str] = False,
     ):
         """
         Construct a function repository.
@@ -239,10 +239,13 @@ class LockedFuncRepository:
             lockfile (`Path`, *optional*): Path to the lockfile. If not provided,
                 the lockfile will be inferred from the caller's context.
             func_name (`str`): The name of the function within the kernel repository.
+            trust_remote_code (`bool`, *optional*, defaults to `False`):
+                Whether to allow loading kernels from untrusted organisations.
         """
         self._repo_id = repo_id
         self._lockfile = lockfile
         self.func_name = func_name
+        self._trust_remote_code = trust_remote_code
         self._revision = self._resolve_revision()
 
     def _resolve_revision(self) -> str:
@@ -258,7 +261,7 @@ class LockedFuncRepository:
         return locked_sha
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_kernel(repo_id=self._repo_id, revision=self._revision)
+        kernel = get_kernel(repo_id=self._repo_id, revision=self._revision, trust_remote_code=self._trust_remote_code)
         return _get_kernel_func(self, kernel)
 
     def __eq__(self, other):
@@ -267,10 +270,11 @@ class LockedFuncRepository:
             and self.func_name == other.func_name
             and self._repo_id == other._repo_id
             and self._revision == other._revision
+            and self._trust_remote_code == other._trust_remote_code
         )
 
     def __hash__(self):
-        return hash((self.func_name, self._repo_id, self._revision))
+        return hash((self.func_name, self._repo_id, self._revision, self._trust_remote_code))
 
     def __str__(self) -> str:
         return f"`{self._repo_id}` (revision: {self._revision}), function `{self.func_name}`"

{kernels-0.14.0.dev0 → kernels-0.14.1}/src/kernels/layer/kernelize.py

@@ -139,12 +139,12 @@ def register_kernel_mapping(
             "MultiHeadAttention": {
                 "cuda": {
                     Mode.TRAINING: LayerRepository(
-                        repo_id="username/training-kernels",
+                        repo_id="kernels-community/training-kernels",
                         layer_name="TrainingAttention",
                         version=1,
                     ),
                     Mode.INFERENCE: LayerRepository(
-                        repo_id="username/inference-kernels",
+                        repo_id="kernels-community/inference-kernels",
                         layer_name="FastAttention",
                         version=1,
                     ),
@@ -206,7 +206,7 @@ def kernelize(
         import torch
         import torch.nn as nn
 
-        from kernels import kernelize, Mode, register_kernel_mapping, LayerRepository
+        from kernels import kernelize, Mode, use_kernel_mapping, LayerRepository
         from kernels import use_kernel_forward_from_hub
 
         @use_kernel_forward_from_hub("SiluAndMul")
@@ -220,10 +220,10 @@ def kernelize(
                 "cuda": LayerRepository(
                     repo_id="kernels-community/activation",
                     layer_name="SiluAndMul",
+                    version=1,
                 )
             }
         }
-        register_kernel_mapping(mapping)
 
         # Create and kernelize a model
         model = nn.Sequential(
@@ -232,7 +232,8 @@ def kernelize(
         )
 
         # Kernelize for inference
-        kernelized_model = kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE)
+        with use_kernel_mapping(mapping):
+            kernelized_model = kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE)
         ```
     """
 

{kernels-0.14.0.dev0 → kernels-0.14.1}/src/kernels/layer/layer.py

@@ -42,6 +42,10 @@ class LayerRepository:
             The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`.
         version (`int`, *optional*):
             The kernel version to download. Cannot be used together with `revision`.
+        trust_remote_code (`bool | list[str]`, *optional*, defaults to `False`):
+            Whether to allow loading kernels from untrusted organisations. A list
+            of signing identities can be provided for future verification support;
+            until then it warns and falls back to the default trust check.
 
     Example:
         ```python
@@ -63,12 +67,14 @@ class LayerRepository:
         layer_name: str,
         revision: str | None = None,
         version: int | None = None,
+        trust_remote_code: bool | list[str] = False,
     ):
         if revision is not None and version is not None:
             raise ValueError("Either a revision or a version must be specified, not both.")
 
         self._repo_id = repo_id
         self.layer_name = layer_name
+        self._trust_remote_code = trust_remote_code
 
         # We are going to resolve these lazily, since we do not want
         # to do a network request for every registered LayerRepository.
@@ -84,7 +90,9 @@ class LayerRepository:
         )
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_kernel(self._repo_id, revision=self._resolve_revision())
+        kernel = get_kernel(
+            self._repo_id, revision=self._resolve_revision(), trust_remote_code=self._trust_remote_code
+        )
         return _get_kernel_layer(self, kernel)
 
     def __eq__(self, other):
@@ -94,10 +102,11 @@ class LayerRepository:
             and self._repo_id == other._repo_id
             and self._revision == other._revision
             and self._version == other._version
+            and self._trust_remote_code == other._trust_remote_code
         )
 
     def __hash__(self):
-        return hash((self.layer_name, self._repo_id, self._revision, self._version))
+        return hash((self.layer_name, self._repo_id, self._revision, self._version, self._trust_remote_code))
 
     def __str__(self) -> str:
         return f"`{self._repo_id}` (revision: {self._resolve_revision()}), layer `{self.layer_name}`"
@@ -110,8 +119,6 @@ class LocalLayerRepository:
     Args:
         repo_path (`Path`):
             The local repository containing the layer.
-        package_name (`str`):
-            Package name of the kernel.
         layer_name (`str`):
             The name of the layer within the kernel repository.
 
@@ -124,7 +131,6 @@ class LocalLayerRepository:
         # Reference a specific layer by revision
         layer_repo = LocalLayerRepository(
             repo_path=Path("/home/daniel/kernels/activation"),
-            package_name="activation",
             layer_name="SiluAndMul",
         )
         ```
@@ -134,15 +140,13 @@ class LocalLayerRepository:
         self,
         repo_path: Path,
         *,
-        package_name: str,
         layer_name: str,
     ):
         self._repo_path = repo_path
-        self._package_name = package_name
         self.layer_name = layer_name
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_local_kernel(self._repo_path, self._package_name)
+        kernel = get_local_kernel(self._repo_path)
         return _get_kernel_layer(self, kernel)
 
     def __eq__(self, other):
@@ -150,14 +154,13 @@ class LocalLayerRepository:
             isinstance(other, LocalLayerRepository)
             and self.layer_name == other.layer_name
             and self._repo_path == other._repo_path
-            and self._package_name == other._package_name
         )
 
     def __hash__(self):
-        return hash((self.layer_name, self._repo_path, self._package_name))
+        return hash((self.layer_name, self._repo_path))
 
     def __str__(self) -> str:
-        return f"`{self._repo_path}` (package: {self._package_name}), layer `{self.layer_name}`"
+        return f"`{self._repo_path}` (layer `{self.layer_name}`"
 
 
 class LockedLayerRepository:
@@ -174,6 +177,7 @@ class LockedLayerRepository:
         *,
         lockfile: Path | None = None,
         layer_name: str,
+        trust_remote_code: bool | list[str] = False,
     ):
         """
         Construct a layer repository.
@@ -184,6 +188,7 @@ class LockedLayerRepository:
         self._repo_id = repo_id
         self._lockfile = lockfile
         self.layer_name = layer_name
+        self._trust_remote_code = trust_remote_code
         self._revision = self._resolve_revision()
 
     def _resolve_revision(self) -> str:
@@ -199,7 +204,7 @@ class LockedLayerRepository:
         return locked_sha
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_kernel(repo_id=self._repo_id, revision=self._revision)
+        kernel = get_kernel(repo_id=self._repo_id, revision=self._revision, trust_remote_code=self._trust_remote_code)
         return _get_kernel_layer(self, kernel)
 
     def __eq__(self, other):
@@ -208,10 +213,11 @@ class LockedLayerRepository:
             and self.layer_name == other.layer_name
             and self._repo_id == other._repo_id
             and self._revision == other._revision
+            and self._trust_remote_code == other._trust_remote_code
         )
 
     def __hash__(self):
-        return hash((self.layer_name, self._repo_id, self._revision))
+        return hash((self.layer_name, self._repo_id, self._revision, self._trust_remote_code))
 
     def __str__(self) -> str:
         return f"`{self._repo_id}` (revision: {self._revision}), layer `{self.layer_name}`"