Add publishing workflow (#81)

* Allow layers to opt in to `torch.compile`

This change allows a layer to set the `can_torch_compile` class
variable to indicate that the layer is compatible with `torch.compile`.
When enabled, the layer does not fall back to the original
implementation when `torch.compile` is used.

* Comment fixes

Co-authored-by: Mohamed Mekkouri <93391238+MekkCyber@users.noreply.github.com>

---------

Co-authored-by: Mohamed Mekkouri <93391238+MekkCyber@users.noreply.github.com>

.github/workflows/publish.yml

@@ -0,0 +1,119 @@
+name: Publish Python 🐍 distribution 📦 to PyPI and TestPyPI
+
+on: push
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.9"
+      - name: Install pypa/build
+        run: >-
+          python3 -m
+          pip install
+          build
+          --user
+      - name: Build a binary wheel and a source tarball
+        run: python3 -m build
+      - name: Store the distribution packages
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    if: startsWith(github.ref, 'refs/tags/') # only publish to PyPI on tag pushes
+    needs:
+      - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/<package-name> # Replace <package-name> with your PyPI project name
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    name: >-
+      Sign the Python 🐍 distribution 📦 with Sigstore
+      and upload them to GitHub Release
+    needs:
+      - publish-to-pypi
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write # IMPORTANT: mandatory for making GitHub Releases
+      id-token: write # IMPORTANT: mandatory for sigstore
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Sign the dists with Sigstore
+        uses: sigstore/gh-action-sigstore-python@v3.0.0
+        with:
+          inputs: >-
+            ./dist/*.tar.gz
+            ./dist/*.whl
+      - name: Create GitHub Release
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        run: >-
+          gh release create
+          "$GITHUB_REF_NAME"
+          --repo "$GITHUB_REPOSITORY"
+          --notes ""
+      - name: Upload artifact signatures to GitHub Release
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        # Upload to GitHub Release using the `gh` CLI.
+        # `dist/` contains the built packages, and the
+        # sigstore-produced signatures and certificates.
+        run: >-
+          gh release upload
+          "$GITHUB_REF_NAME" dist/**
+          --repo "$GITHUB_REPOSITORY"
+
+  publish-to-testpypi:
+    name: Publish Python 🐍 distribution 📦 to TestPyPI
+    needs:
+      - build
+    runs-on: ubuntu-latest
+
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/<package-name>
+
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Publish distribution 📦 to TestPyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/

docs/kernel-requirements.md

@@ -1,8 +1,11 @@
 # Kernel requirements
 
-Kernels on the Hub must fulfill the requirements outlined on this page.
+Kernels on the Hub must fulfill the requirements outlined on this page. By
+ensuring kernels are compliant, they can be used on a wide range of Linux
+systems and Torch builds.
+
 You can use [kernel-builder](https://github.com/huggingface/kernel-builder/)
-to build conforming kernels.
+to build compliant kernels.
 
 ## Directory layout
 
@@ -10,34 +13,21 @@ A kernel repository on the Hub must contain a `build` directory. This
 directory contains build variants of a kernel in the form of directories
 following the template
 `<framework><version>-cxx<abiver>-<cu><cudaver>-<arch>-<os>`.
-For example `build/torch26-cxx98-cu118-x86_64-linux`. The currently
-recommended build variants are:
+For example `build/torch26-cxx98-cu118-x86_64-linux`.
 
-- `torch25-cxx11-cu118-x86_64-linux`
-- `torch25-cxx11-cu121-x86_64-linux`
-- `torch25-cxx11-cu124-x86_64-linux`
-- `torch25-cxx98-cu118-x86_64-linux`
-- `torch25-cxx98-cu121-x86_64-linux`
-- `torch25-cxx98-cu124-x86_64-linux`
-- `torch26-cxx11-cu118-x86_64-linux`
-- `torch26-cxx11-cu124-x86_64-linux`
-- `torch26-cxx11-cu126-x86_64-linux`
-- `torch26-cxx98-cu118-x86_64-linux`
-- `torch26-cxx98-cu124-x86_64-linux`
-- `torch26-cxx98-cu126-x86_64-linux`
-
-This list will be updated as new PyTorch versions are released. Kernels
-that are in pure Python (e.g. Triton kernels) only need to provide a
-single build variant:
-
-- `torch-universal`
-
-Each variant directory should contain a single directory with the same name
+Each variant directory must contain a single directory with the same name
 as the repository (replacing `-` by `_`). For instance, kernels in the
 `kernels-community/activation` repository have a directories like
 `build/<variant>/activation`. This directory
 must be a Python package with an `__init__.py` file.
 
+## Build variants
+
+A kernel can be compliant for a specific compute framework (e.g. CUDA) or
+architecture (e.g. x86_64). For compliance with a compute framework and
+architecture combination, all the variants from the [build variant list](https://github.com/huggingface/kernel-builder/blob/main/docs/build-variants.md)
+must be available for that combination.
+
 ## Versioning
 
 Kernels are versioned on the Hub using Git tags. Version tags must be of
@@ -119,9 +109,12 @@ requirements:
 - The `forward` method has a signature that is compatible with the
   `forward` method that it is extending.
 
-The only exception to the _no class variables rule_ is addition of a
-`has_backward` class variable. This variable is used to indicate whether
-the layer has a backward pass implemented (`True` when absent).
+There are two exceptions to the _no class variables rule_:
+
+1. The `has_backward` variable can be used to indicate whether the layer has
+   a backward pass implemented (`True` when absent).
+2. The `can_torch_compile` variable can be used to indicate whether the layer
+   supports `torch.compile` (`False` when absent).
 
 This is an example of a pure layer:
 

docs/locking.md

@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 "kernels-community/activation" = ">=0.0.1"
 ```
 
-Then run `kernel lock .` in the project directory. This generates a `kernels.lock` file with
+Then run `kernels lock .` in the project directory. This generates a `kernels.lock` file with
 the locked revisions. The locked revision will be used when loading a kernel with
 `get_locked_kernel`:
 
@@ -28,7 +28,7 @@ to `kernels` after doing an (editable or regular) installation of your project.
 
 ## Pre-downloading locked kernels
 
-Locked kernels can be pre-downloaded by running `kernel download .` in your
+Locked kernels can be pre-downloaded by running `kernels download .` in your
 project directory. This will download the kernels to your local Hugging Face
 Hub cache.
 

pyproject.toml

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

src/kernels/layer.py

@@ -138,6 +138,8 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
             return fallback_forward(self, x, *args, **kwargs)
 
         needs_backward = self.training
+        is_compiling = _is_torchdynamo_compiling()
+
         kernel = _KERNEL_MAPPING.get().get(layer_name)
         if kernel is None:
             warnings.warn(
@@ -165,7 +167,14 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
         # Short-circuit if we already loaded the layer.
         layer = cached_layer.get(repo, None)
         if layer is not None:
-            if needs_backward and not getattr(layer, "has_backward", True):
+            # Switch to fallback when the layer does not support:
+            # compilation/compile when needed.
+            # backward when needed
+            needs_fallback = needs_backward and not getattr(layer, "has_backward", True)
+            needs_fallback |= is_compiling and not getattr(
+                layer, "can_torch_compile", False
+            )
+            if needs_fallback:
                 return fallback_forward(self, x, *args, **kwargs)
             return layer.forward(self, x, *args, **kwargs)
 
@@ -185,8 +194,15 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
 
         cached_layer[repo] = layer
 
-        if needs_backward and not getattr(layer, "has_backward", True):
+        # Switch to fallback when the layer does not support
+        # compilation/compile when needed.
+        needs_fallback = needs_backward and not getattr(layer, "has_backward", True)
+        needs_fallback |= is_compiling and not getattr(
+            layer, "can_torch_compile", False
+        )
+        if needs_fallback:
             return fallback_forward(self, x, *args, **kwargs)
+
         return layer.forward(self, x, *args, **kwargs)
 
     cls.forward = forward
@@ -245,7 +261,8 @@ def _validate_layer(*, check_cls, cls):
     torch_module_members = {name for name, _ in inspect.getmembers(nn.Module)}
     cls_members = {name for name, _ in inspect.getmembers(cls)}
     difference = cls_members - torch_module_members
-    if difference != set() and difference != {"has_backward"}:
+    # verify if : difference ⊄ {"can_torch_compile", "has_backward"}
+    if not difference <= {"can_torch_compile", "has_backward"}:
         raise TypeError("Layer must not contain additional members.")
 
     # Check whether the forward signatures are similar.
@@ -262,3 +279,19 @@ def _validate_layer(*, check_cls, cls):
             raise TypeError(
                 f"Forward signature does not match: different kind of arguments ({param} ({param.kind}) and {ref_param} ({ref_param.kind})"
             )
+
+
+def _is_torchdynamo_compiling():
+    # Importing torch._dynamo causes issues with PyTorch profiler (https://github.com/pytorch/pytorch/issues/130622)
+    # hence rather relying on `torch.compiler.is_compiling()` when possible (torch>=2.3)
+    try:
+        import torch
+
+        return torch.compiler.is_compiling()
+    except Exception:
+        try:
+            import torch._dynamo as dynamo  # noqa: F401
+
+            return dynamo.is_compiling()
+        except Exception:
+            return False

tests/test_layer.py

@@ -19,6 +19,12 @@ kernel_layer_mapping = {
             revision="layers",
         )
     },
+    "SiluAndMulNoCompile": {
+        "cuda": LayerRepository(
+            repo_id="kernels-test/op-without-fake-test",
+            layer_name="SiluAndMul",
+        )
+    },
     "SiluAndMulStringDevice": {
         "cuda": LayerRepository(
             repo_id="kernels-community/activation",
@@ -43,6 +49,11 @@ class SiluAndMul(nn.Module):
         return F.silu(input[..., :d]) * input[..., d:]
 
 
+@use_kernel_forward_from_hub("SiluAndMulNoCompile")
+class SiluAndMulNoCompileKernel(SiluAndMul):
+    pass
+
+
 @use_kernel_forward_from_hub("SiluAndMul")
 class SiluAndMulWithKernel(SiluAndMul):
     pass
@@ -101,8 +112,29 @@ def test_layer_fallback_works():
     SiluAndMulWithKernelFallback()
 
 
+@pytest.mark.parametrize("cls", [SiluAndMulWithKernel, SiluAndMulNoCompileKernel])
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+def test_torch_compile_layer(cls, device):
+    silu_and_mul = SiluAndMul()
+
+    X = torch.randn((32, 64), dtype=torch.float32, device=device)
+    Y = silu_and_mul(X)
+
+    silu_and_mul_with_kernel = cls()
+    silu_and_mul_with_kernel.eval()
+    silu_and_mul_compiled = torch.compile(silu_and_mul_with_kernel)
+
+    Y_compiled = silu_and_mul_compiled(X)
+
+    torch.testing.assert_close(Y_compiled, Y)
+
+
 def test_mapping_contexts():
-    assert set(_KERNEL_MAPPING.get().keys()) == {"SiluAndMul", "SiluAndMulStringDevice"}
+    assert set(_KERNEL_MAPPING.get().keys()) == {
+        "SiluAndMul",
+        "SiluAndMulStringDevice",
+        "SiluAndMulNoCompile",
+    }
 
     extra_mapping1 = {
         "TestKernel": {
@@ -118,6 +150,7 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
+            "SiluAndMulNoCompile",
             "TestKernel",
         }
 
@@ -135,6 +168,7 @@ def test_mapping_contexts():
             assert set(_KERNEL_MAPPING.get().keys()) == {
                 "SiluAndMul",
                 "SiluAndMulStringDevice",
+                "SiluAndMulNoCompile",
                 "TestKernel",
             }
             assert (
@@ -145,6 +179,7 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
+            "SiluAndMulNoCompile",
             "TestKernel",
         }
         assert (
@@ -164,6 +199,7 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
+            "SiluAndMulNoCompile",
             "TestKernel",
         }
         assert (
@@ -174,6 +210,7 @@ def test_mapping_contexts():
     assert set(_KERNEL_MAPPING.get().keys()) == {
         "SiluAndMul",
         "SiluAndMulStringDevice",
+        "SiluAndMulNoCompile",
     }