kernels 0.2.0__tar.gz → 0.3.0__tar.gz

{kernels-0.2.0 → kernels-0.3.0}/PKG-INFO

@@ -1,7 +1,7 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.1
 Name: kernels
-Version: 0.2.0
-Summary: Download cuda kernels
+Version: 0.3.0
+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>
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
@@ -57,6 +57,7 @@ the Hub.
 
 ## 📚 Documentation
 
+- [Using layers](docs/layers.md)
 - [Locking kernel versions](docs/locking.md)
 - [Using kernels in a Docker container](docs/docker.md)
 - [Kernel requirements](docs/kernel-requirements.md)

{kernels-0.2.0 → kernels-0.3.0}/README.md

@@ -45,6 +45,7 @@ the Hub.
 
 ## 📚 Documentation
 
+- [Using layers](docs/layers.md)
 - [Locking kernel versions](docs/locking.md)
 - [Using kernels in a Docker container](docs/docker.md)
 - [Kernel requirements](docs/kernel-requirements.md)

{kernels-0.2.0 → kernels-0.3.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "kernels"
-version = "0.2.0"
-description = "Download cuda kernels"
+version = "0.3.0"
+description = "Download compute kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },
   { name = "Daniel de Kok", email = "daniel@huggingface.co" },

kernels-0.3.0/src/kernels/__init__.py

@@ -0,0 +1,23 @@
+from kernels.layer import (
+    Device,
+    LayerRepository,
+    register_kernel_mapping,
+    use_kernel_forward_from_hub,
+)
+from kernels.utils import (
+    get_kernel,
+    get_locked_kernel,
+    install_kernel,
+    load_kernel,
+)
+
+__all__ = [
+    "get_kernel",
+    "get_locked_kernel",
+    "load_kernel",
+    "install_kernel",
+    "use_kernel_forward_from_hub",
+    "register_kernel_mapping",
+    "LayerRepository",
+    "Device",
+]

kernels-0.3.0/src/kernels/layer.py

@@ -0,0 +1,231 @@
+import inspect
+from contextvars import ContextVar
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Callable, Dict, Union
+
+from .utils import get_kernel
+
+if TYPE_CHECKING:
+    from torch import nn
+
+
+@dataclass(frozen=True)
+class Device:
+    type: str
+
+    # In the future we might add compute capabilities, etc.
+
+    def __eq__(self, other):
+        return isinstance(other, Device) and self.type == other.type
+
+    def __hash__(self):
+        return hash(self.type)
+
+
+@dataclass
+class LayerRepository:
+    """
+    Repository and name of a layer.
+    """
+
+    layer_name: str = field(
+        metadata={"help": "The name of the layer in the kernel repository."}
+    )
+    repo_id: str = field(metadata={"help": "The kernel hub repository with the layer."})
+    revision: str = field(
+        default="main", metadata={"help": "The revision of the layer."}
+    )
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, LayerRepository)
+            and self.layer_name == other.layer_name
+            and self.repo_id == other.repo_id
+            and self.revision == other.revision
+        )
+
+    def __hash__(self):
+        return hash((self.layer_name, self.repo_id, self.revision))
+
+
+_KERNEL_MAPPING: ContextVar[Dict[str, Dict[Device, LayerRepository]]] = ContextVar(
+    "_KERNEL_MAPPING", default={}
+)
+
+
+def use_kernel_mapping(mapping: Dict[str, Dict[Union[Device, str], LayerRepository]]):
+    class ContextManager:
+        def __enter__(self):
+            # Mappings always stack on previous mappings.
+            self.token = _KERNEL_MAPPING.set(deepcopy(_KERNEL_MAPPING.get()))
+            register_kernel_mapping(mapping)
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            _KERNEL_MAPPING.reset(self.token)
+
+    return ContextManager()
+
+
+def register_kernel_mapping(
+    mapping: Dict[str, Dict[Union[Device, str], LayerRepository]]
+):
+    """
+    Allows one to register a mapping between a layer name the corresponding kernel to use, depending on the device.
+    This should be use in conjunction with `use_kernel_hub_forward` decorator on the classname.
+    Exemple usage:
+
+    ```python
+    from kernels import LayerRepository, register_kernel_mapping
+
+    kernel_layer_mapping = {
+      "LlamaRMSNorm": {
+          "cuda": LayerRepository(
+              repo_id="kernels-community/activation",
+              layer_name="RmsNorm",
+              revision="layers",
+          ),
+      },
+    }
+    register_kernel_mapping(kernel_layer_mapping)
+    ```
+    """
+    # Merge with existing mappings.
+    for new_kernel, new_device_repos in mapping.items():
+        device_repo = _KERNEL_MAPPING.get().setdefault(new_kernel, {})
+        for new_device, new_repo in new_device_repos.items():
+            if isinstance(new_device, str):
+                device_repo[Device(type=new_device)] = new_repo
+            else:
+                device_repo[new_device] = new_repo
+
+
+def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool = True):
+    """
+    Replace the forward function of a layer using a layer from the kernel hub.
+    This function monkeypatches a layer, replacing the `forward` method
+    of the layer with that of a layer from the hub. The replacement is done
+    when a layer matching `layer_name` and device type is registered through
+    `register_layer_mapping`. The device type is inferred from the first
+    argument to `forward`.
+    """
+
+    fallback_forward = cls.forward
+
+    cached_forward: Dict[LayerRepository, Callable] = {}
+
+    def forward(self, x, **args):
+        kernel = _KERNEL_MAPPING.get().get(layer_name)
+        if kernel is None:
+            if not use_fallback:
+                raise ValueError(f"No layer mapping for `{layer_name}`")
+            return fallback_forward(self, x, **args)
+
+        device = getattr(x, "device", None)
+        if device is None:
+            return fallback_forward(self, x, **args)
+
+        repo = kernel.get(Device(type=device.type))
+        if repo is None:
+            if not use_fallback:
+                raise ValueError(
+                    f"No layer mapping for `{layer_name}` with device type `{device.type}`"
+                )
+            return fallback_forward(self, x, **args)
+
+        # Short-circuit if we already loaded the layer.
+        layer_forward = cached_forward.get(repo, None)
+        if layer_forward is not None:
+            return layer_forward(self, x, **args)
+
+        layer = _get_kernel_layer(
+            repo_id=repo.repo_id,
+            layer_name=repo.layer_name,
+            revision=repo.revision,
+        )
+
+        # We have to validate against the original signature.
+        orig_forward = cls.forward
+        try:
+            cls.forward = fallback_forward
+            _validate_layer(check_cls=cls, cls=layer)
+        finally:
+            cls.forward = orig_forward
+
+        layer_forward = layer.forward
+        cached_forward[repo] = layer_forward
+
+        return layer_forward(self, x, **args)
+
+    cls.forward = forward
+
+
+def use_kernel_forward_from_hub(layer_name: str, *, use_fallback: bool = True):
+    """
+    Replace the forward function of a layer using a layer from the kernel hub.
+    This decorator can be applied to a layer and replaces the forward method
+    of the layer with that of a layer from the hub. The replacement is done
+    when a layer matching `layer_name` and device type is registered through
+    `register_layer_mapping`. The device type is inferred from the first
+    argument to `forward`.
+    """
+
+    def decorator(cls):
+        replace_kernel_forward_from_hub(cls, layer_name, use_fallback=use_fallback)
+        return cls
+
+    return decorator
+
+
+def _get_kernel_layer(*, repo_id: str, layer_name: str, revision: str) -> "nn.Module":
+    """Get a layer from a kernel."""
+
+    kernel = get_kernel(repo_id, revision=revision)
+
+    if getattr(kernel, "layers", None) is None:
+        raise ValueError(
+            f"Kernel `{repo_id}` at revision `{revision}` does not define any layers."
+        )
+
+    layer = getattr(kernel.layers, layer_name, None)
+    if layer is None:
+        raise ValueError(f"Layer `{layer_name}` not found in kernel `{repo_id}`.")
+    return layer
+
+
+def _validate_layer(*, check_cls, cls):
+    # The layer must have at least have the following properties: (1) it
+    # must be stateless; (2) the forward signature should correspond to
+    # the signature it is replacing; (3) forward should not call other
+    # methods.
+
+    from torch import nn
+
+    if not issubclass(cls, nn.Module):
+        raise TypeError(f"Layer `{cls}` is not a Torch layer.")
+
+    # We verify statelessness by checking that the does not have its own
+    # constructor (since the constructor could add member variables)...
+    if cls.__init__ is not nn.Module.__init__:
+        raise TypeError("Layer must not override nn.Module constructor.")
+
+    # ... or predefined member variables.
+    torch_module_members = {name for name, _ in inspect.getmembers(nn.Module)}
+    cls_members = {name for name, _ in inspect.getmembers(cls)}
+    if cls_members - torch_module_members != set():
+        raise TypeError("Layer must not contain additional members.")
+
+    # Check whether the forward signatures are similar.
+    params = inspect.signature(cls.forward).parameters
+    ref_params = inspect.signature(check_cls.forward).parameters
+
+    if len(params) != len(ref_params):
+        raise TypeError(
+            "Forward signature does not match: different number of arguments."
+        )
+
+    for param, ref_param in zip(params.values(), ref_params.values()):
+        if param.kind != ref_param.kind:
+            raise TypeError(
+                f"Forward signature does not match: different kind of arguments ({param} ({param.kind}) and {ref_param} ({ref_param.kind})"
+            )

{kernels-0.2.0 → kernels-0.3.0}/src/kernels/utils.py

@@ -144,9 +144,18 @@ def get_kernel(repo_id: str, revision: str = "main") -> ModuleType:
     return import_from_path(package_name, package_path / package_name / "__init__.py")
 
 
-def load_kernel(repo_id: str) -> ModuleType:
-    """Get a pre-downloaded, locked kernel."""
-    locked_sha = _get_caller_locked_kernel(repo_id)
+def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
+    """
+    Get a pre-downloaded, locked kernel.
+
+    If `lockfile` is not specified, the lockfile will be loaded from the
+    caller's package metadata.
+    """
+    if lockfile is None:
+        locked_sha = _get_caller_locked_kernel(repo_id)
+    else:
+        with open(lockfile, "r") as f:
+            locked_sha = _get_locked_kernel(repo_id, f.read())
 
     if locked_sha is None:
         raise ValueError(
@@ -163,6 +172,7 @@ def load_kernel(repo_id: str) -> ModuleType:
             repo_id,
             allow_patterns=[f"build/{variant}/*", f"build/{universal_variant}/*"],
             cache_dir=CACHE_DIR,
+            revision=locked_sha,
             local_files_only=True,
         )
     )
@@ -200,11 +210,19 @@ def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleTyp
 def _get_caller_locked_kernel(repo_id: str) -> Optional[str]:
     for dist in _get_caller_distributions():
         lock_json = dist.read_text("kernels.lock")
-        if lock_json is not None:
-            for kernel_lock_json in json.loads(lock_json):
-                kernel_lock = KernelLock.from_json(kernel_lock_json)
-                if kernel_lock.repo_id == repo_id:
-                    return kernel_lock.sha
+        if lock_json is None:
+            continue
+        locked_sha = _get_locked_kernel(repo_id, lock_json)
+        if locked_sha is not None:
+            return locked_sha
+    return None
+
+
+def _get_locked_kernel(repo_id: str, lock_json: str) -> Optional[str]:
+    for kernel_lock_json in json.loads(lock_json):
+        kernel_lock = KernelLock.from_json(kernel_lock_json)
+        if kernel_lock.repo_id == repo_id:
+            return kernel_lock.sha
     return None
 
 

{kernels-0.2.0 → kernels-0.3.0}/src/kernels.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.1
 Name: kernels
-Version: 0.2.0
-Summary: Download cuda kernels
+Version: 0.3.0
+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>
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
@@ -57,6 +57,7 @@ the Hub.
 
 ## 📚 Documentation
 
+- [Using layers](docs/layers.md)
 - [Locking kernel versions](docs/locking.md)
 - [Using kernels in a Docker container](docs/docker.md)
 - [Kernel requirements](docs/kernel-requirements.md)

{kernels-0.2.0 → kernels-0.3.0}/src/kernels.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@ pyproject.toml
 src/kernels/__init__.py
 src/kernels/cli.py
 src/kernels/compat.py
+src/kernels/layer.py
 src/kernels/lockfile.py
 src/kernels/utils.py
 src/kernels.egg-info/PKG-INFO
@@ -13,4 +14,5 @@ src/kernels.egg-info/requires.txt
 src/kernels.egg-info/top_level.txt
 tests/test_basic.py
 tests/test_benchmarks.py
-tests/test_hash_validation.py
+tests/test_kernel_locking.py
+tests/test_layer.py

kernels-0.3.0/src/kernels.egg-info/top_level.txt

@@ -0,0 +1 @@
+kernels

kernels-0.2.0/tests/test_hash_validation.py → kernels-0.3.0/tests/test_kernel_locking.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from kernels import load_kernel
 from kernels.cli import download_kernels
 
 
@@ -11,11 +12,13 @@ class DownloadArgs:
     project_dir: Path
 
 
-def test_download_hash_validation():
-    project_dir = Path(__file__).parent / "hash_validation"
-    download_kernels(DownloadArgs(all_variants=False, project_dir=project_dir))
-
-
 def test_download_all_hash_validation():
-    project_dir = Path(__file__).parent / "hash_validation"
+    project_dir = Path(__file__).parent / "kernel_locking"
     download_kernels(DownloadArgs(all_variants=True, project_dir=project_dir))
+
+
+def test_load_locked():
+    project_dir = Path(__file__).parent / "kernel_locking"
+    # Also validates that hashing works correctly.
+    download_kernels(DownloadArgs(all_variants=False, project_dir=project_dir))
+    load_kernel("kernels-community/activation", lockfile=project_dir / "kernels.lock")

kernels-0.3.0/tests/test_layer.py

@@ -0,0 +1,168 @@
+import pytest
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from kernels import (
+    Device,
+    LayerRepository,
+    register_kernel_mapping,
+    use_kernel_forward_from_hub,
+)
+from kernels.layer import _KERNEL_MAPPING, _validate_layer, use_kernel_mapping
+
+kernel_layer_mapping = {
+    "SiluAndMul": {
+        Device(type="cuda"): LayerRepository(
+            repo_id="kernels-community/activation",
+            layer_name="SiluAndMul",
+            revision="layers",
+        )
+    },
+    "SiluAndMulStringDevice": {
+        "cuda": LayerRepository(
+            repo_id="kernels-community/activation",
+            layer_name="SiluAndMul",
+            revision="layers",
+        )
+    },
+}
+
+register_kernel_mapping(kernel_layer_mapping)
+
+
+class SiluAndMul(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Used to check that we called hub kernel.
+        self.n_calls = 0
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        self.n_calls += 1
+        d = input.shape[-1] // 2
+        return F.silu(input[..., :d]) * input[..., d:]
+
+
+@use_kernel_forward_from_hub("SiluAndMul")
+class SiluAndMulWithKernel(SiluAndMul):
+    pass
+
+
+@use_kernel_forward_from_hub("SiluAndMulStringDevice")
+class SiluAndMulStringDevice(SiluAndMul):
+    pass
+
+
+@pytest.mark.parametrize("cls", [SiluAndMulWithKernel, SiluAndMulStringDevice])
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+def test_hub_forward(cls, device):
+    torch.random.manual_seed(0)
+
+    silu_and_mul = SiluAndMul()
+    X = torch.randn((32, 64), device=device)
+    Y = silu_and_mul(X)
+
+    silu_and_mul_with_kernel = cls()
+    Y_kernel = silu_and_mul_with_kernel(X)
+
+    torch.testing.assert_close(Y_kernel, Y)
+
+    assert silu_and_mul.n_calls == 1
+    if device == "cuda":
+        assert silu_and_mul_with_kernel.n_calls == 0
+    else:
+        assert silu_and_mul_with_kernel.n_calls == 1
+
+
+def test_layer_fallback_works():
+    @use_kernel_forward_from_hub("SiluAndMulNonExisting")
+    class SiluAndMulWithKernelFallback(SiluAndMul):
+        pass
+
+    # Check that we don't raise an exception for a non-existing kernel.
+    SiluAndMulWithKernelFallback()
+
+
+def test_mapping_contexts():
+    assert set(_KERNEL_MAPPING.get().keys()) == {"SiluAndMul", "SiluAndMulStringDevice"}
+
+    extra_mapping1 = {
+        "TestKernel": {
+            Device(type="cuda"): LayerRepository(
+                repo_id="kernels-community/activation",
+                layer_name="SiluAndMul",
+                revision="layers",
+            )
+        }
+    }
+
+    with use_kernel_mapping(extra_mapping1):
+        assert set(_KERNEL_MAPPING.get().keys()) == {
+            "SiluAndMul",
+            "SiluAndMulStringDevice",
+            "TestKernel",
+        }
+
+        extra_mapping2 = {
+            "SiluAndMul": {
+                Device(type="cuda"): LayerRepository(
+                    repo_id="kernels-community/non-existing",
+                    layer_name="SiluAndMul",
+                    revision="layers",
+                )
+            }
+        }
+
+        with use_kernel_mapping(extra_mapping2):
+            assert set(_KERNEL_MAPPING.get().keys()) == {
+                "SiluAndMul",
+                "SiluAndMulStringDevice",
+                "TestKernel",
+            }
+            assert (
+                _KERNEL_MAPPING.get()["SiluAndMul"][Device(type="cuda")].repo_id
+                == "kernels-community/non-existing"
+            )
+
+        assert set(_KERNEL_MAPPING.get().keys()) == {
+            "SiluAndMul",
+            "SiluAndMulStringDevice",
+            "TestKernel",
+        }
+        assert (
+            _KERNEL_MAPPING.get()["SiluAndMul"][Device(type="cuda")].repo_id
+            == "kernels-community/activation"
+        )
+
+    assert set(_KERNEL_MAPPING.get().keys()) == {
+        "SiluAndMul",
+        "SiluAndMulStringDevice",
+    }
+
+
+def test_validate_kernel_layer():
+    class BadLayer(nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.foo = 42
+
+    with pytest.raises(TypeError, match="not override"):
+        _validate_layer(cls=BadLayer, check_cls=SiluAndMul)
+
+    class BadLayer2(nn.Module):
+        foo: int = 42
+
+    with pytest.raises(TypeError, match="not contain additional members"):
+        _validate_layer(cls=BadLayer2, check_cls=SiluAndMul)
+
+    class BadLayer3(nn.Module):
+        def forward(self, x: torch.Tensor, foo: int) -> torch.Tensor: ...
+
+    with pytest.raises(TypeError, match="different number of arguments"):
+        _validate_layer(cls=BadLayer3, check_cls=SiluAndMul)
+
+    class BadLayer4(nn.Module):
+        def forward(self, *, x: torch.Tensor) -> torch.Tensor: ...
+
+    with pytest.raises(TypeError, match="different kind of arguments"):
+        _validate_layer(cls=BadLayer4, check_cls=SiluAndMul)

kernels-0.2.0/src/kernels/__init__.py

@@ -1,3 +0,0 @@
-from kernels.utils import get_kernel, get_locked_kernel, install_kernel, load_kernel
-
-__all__ = ["get_kernel", "get_locked_kernel", "load_kernel", "install_kernel"]

kernels-0.2.0/src/kernels.egg-info/top_level.txt

@@ -1,2 +0,0 @@
-hf_kernels
-kernels