Export __version__ for doc-builder

Also document public functions better.

.github/workflows/build_documentation.yaml

@@ -0,0 +1,19 @@
+name: Build documentation
+
+on:
+  push:
+    paths:
+      - "docs/source/**"
+    branches:
+      - main
+      - doc-builder*
+      - v*-release
+
+jobs:
+  build:
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main
+    with:
+      commit_sha: ${{ github.sha }}
+      package: kernels
+    secrets:
+      hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}

.github/workflows/build_pr_documentation.yaml

@@ -0,0 +1,18 @@
+name: Build PR Documentation
+
+on:
+  pull_request:
+    paths:
+      - "docs/source/**"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main
+    with:
+      commit_sha: ${{ github.event.pull_request.head.sha }}
+      pr_number: ${{ github.event.number }}
+      package: kernels

.github/workflows/publish.yml

@@ -1,119 +0,0 @@
-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/source/_toctree.yml

@@ -0,0 +1,26 @@
+- sections:
+    - local: index
+      title: Introduction
+    - local: installation
+      title: Installation
+  title: Getting started
+- sections:
+    - local: basic_usage
+      title: Basic Usage
+    - local: layers
+      title: Using Layers
+    - local: locking
+      title: Locking Kernel Versions
+    - local: env
+      title: Environment Variables
+  title: Usage Guide
+- sections:
+    - local: api/kernels
+      title: Kernels
+    - local: api/layers
+      title: Layers
+  title: API Reference
+- sections:
+    - local: kernel_requirements
+      title: Kernel Requirements
+  title: Developer Guide

docs/source/api/kernels.md

@@ -0,0 +1,21 @@
+# Kernels API Reference
+
+## Main Functions
+
+### get_kernel
+
+[[autodoc]] kernels.get_kernel
+
+### has_kernel
+
+[[autodoc]] kernels.has_kernel
+
+## Loading locked kernels
+
+### load_kernel
+
+[[autodoc]] kernels.load_kernel
+
+### get_locked_kernel
+
+[[autodoc]] kernels.get_locked_kernel

docs/source/api/layers.md

@@ -0,0 +1,31 @@
+# Layers API Reference
+
+## Making layers kernel-aware
+
+### use_kernel_forward_from_hub
+
+[[autodoc]] kernels.use_kernel_forward_from_hub
+
+### replace_kernel_forward_from_hub
+
+[[autodoc]] kernels.replace_kernel_forward_from_hub
+
+## Registering kernel mappings
+
+### use_kernel_mapping
+
+[[autodoc]] kernels.use_kernel_mapping
+
+### register_kernel_mapping
+
+[[autodoc]] kernels.register_kernel_mapping
+
+## Classes
+
+### LayerRepository
+
+[[autodoc]] kernels.LayerRepository
+
+### Device
+
+[[autodoc]] kernels.Device

docs/source/basic_usage.md

@@ -0,0 +1,34 @@
+# Basic Usage
+
+## Loading Kernels
+
+Here is how you would use the [activation](https://huggingface.co/kernels-community/activation) kernels from the Hugging Face Hub:
+
+```python
+import torch
+from kernels import get_kernel
+
+# Download optimized kernels from the Hugging Face hub
+activation = get_kernel("kernels-community/activation")
+
+# Create a random tensor
+x = torch.randn((10, 10), dtype=torch.float16, device="cuda")
+
+# Run the kernel
+y = torch.empty_like(x)
+activation.gelu_fast(y, x)
+
+print(y)
+```
+
+## Checking Kernel Availability
+
+You can check if a specific kernel is available for your environment:
+
+```python
+from kernels import has_kernel
+
+# Check if kernel is available for current environment
+is_available = has_kernel("kernels-community/activation")
+print(f"Kernel available: {is_available}")
+```

docs/source/index.md

@@ -0,0 +1,20 @@
+# Kernels
+
+<div align="center">
+<img src="https://github.com/user-attachments/assets/64a652f3-0cd3-4829-b3c1-df13f7933569" width="450" height="450" alt="kernel-builder logo">
+</div>
+
+The Kernel Hub allows Python libraries and applications to load compute
+kernels directly from the [Hub](https://hf.co/). To support this kind
+of dynamic loading, Hub kernels differ from traditional Python kernel
+packages in that they are made to be:
+
+- **Portable**: a kernel can be loaded from paths outside `PYTHONPATH`.
+- **Unique**: multiple versions of the same kernel can be loaded in the
+  same Python process.
+- **Compatible**: kernels must support all recent versions of Python and
+  the different PyTorch build configurations (various CUDA versions
+  and C++ ABIs). Furthermore, older C library versions must be supported.
+
+You can [search for kernels](https://huggingface.co/models?other=kernel) on
+the Hub.

docs/source/installation.md

@@ -0,0 +1,16 @@
+# Installation
+
+Install the `kernels` package with `pip` (requires `torch>=2.5` and CUDA):
+
+```bash
+pip install kernels
+```
+
+# Using kernels in a Docker container
+
+build and run the reference [examples/basic.py](examples/basic.py) in a Docker container with the following commands:
+
+```bash
+docker build --platform linux/amd64 -t kernels-reference -f docker/Dockerfile.reference .
+docker run --gpus all -it --rm -e HF_TOKEN=$HF_TOKEN kernels-reference
+```

docs/source/kernel_requirements.md

@@ -1,11 +1,8 @@
 # Kernel requirements
 
-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.
-
+Kernels on the Hub must fulfill the requirements outlined on this page.
 You can use [kernel-builder](https://github.com/huggingface/kernel-builder/)
-to build compliant kernels.
+to build conforming kernels.
 
 ## Directory layout
 
@@ -13,21 +10,34 @@ 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`.
+For example `build/torch26-cxx98-cu118-x86_64-linux`. The currently
+recommended build variants are:
 
-Each variant directory must contain a single directory with the same name
+- `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
 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
@@ -109,12 +119,9 @@ requirements:
 - The `forward` method has a signature that is compatible with the
   `forward` method that it is extending.
 
-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).
+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).
 
 This is an example of a pure layer:
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.5.0.dev0"
+version = "0.4.4"
 description = "Download compute kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },
@@ -31,7 +31,9 @@ dev = [
 ]
 
 [project.optional-dependencies]
-torch = ["torch"]
+docs = [
+  "hf-doc-builder",
+]
 
 [project.scripts]
 kernels = "kernels.cli:main"
@@ -39,7 +41,6 @@ kernels = "kernels.cli:main"
 [project.entry-points."egg_info.writers"]
 "kernels.lock" = "kernels.lockfile:write_egg_lockfile"
 
-
 [tool.ruff]
 exclude = [
   ".eggs",

src/kernels/__init__.py

@@ -1,3 +1,5 @@
+import importlib.metadata
+
 from kernels.layer import (
     Device,
     LayerRepository,
@@ -27,3 +29,5 @@ __all__ = [
     "LayerRepository",
     "Device",
 ]
+
+__version__ = importlib.metadata.version("kernels")

src/kernels/layer.py

@@ -64,11 +64,18 @@ def use_kernel_mapping(
     inherit_mapping: bool = True,
 ):
     """
-    Context manager that sets a mapping for a duration of the context.
+    Context manager that sets a kernel mapping for the duration of the context.
 
-    When `inherit_mapping` is set to `True` the current mapping will be
-    extended by `mapping` inside the context. If it is `False`, only
-    `mapping` is used inside the context.
+    Args:
+        mapping (`Dict[str, Dict[Union[Device, str], LayerRepository]]`):
+            A mapping between layer names and their corresponding kernel repositories.
+        inherit_mapping (`bool`, *optional*, defaults to `True`):
+            The current mapping will be extended by `mapping` when set to `True`.
+            When set to `False`, the current mapping will be replaced by `mapping`
+            for the duration of the context.
+
+    Returns:
+        `ContextManager`: Context manager that sets up the mapping.
     """
 
     class ContextManager:
@@ -87,27 +94,31 @@ def use_kernel_mapping(
 
 
 def register_kernel_mapping(
-    mapping: Dict[str, Dict[Union[Device, str], LayerRepository]]
+    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.
+    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
+    Args:
+        mapping (`Dict[str, Dict[Union[Device, str], LayerRepository]]`):
+            A mapping between layer names and their corresponding kernel repositories.
 
-    kernel_layer_mapping = {
-      "LlamaRMSNorm": {
-          "cuda": LayerRepository(
-              repo_id="kernels-community/activation",
-              layer_name="RmsNorm",
-              revision="layers",
-          ),
-      },
-    }
-    register_kernel_mapping(kernel_layer_mapping)
-    ```
+    Example:
+        ```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():
@@ -125,8 +136,18 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
     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
+    [`register_layer_mapping`]. The device type is inferred from the first
     argument to `forward`.
+
+    Args:
+        cls (`nn.Module`):
+            The layer class to replace the forward function of.
+        layer_name (`str`):
+            The name to assign to the layer.
+        use_fallback (`bool`, *optional*, defaults to `True`):
+            Whether to use the fallback forward function if no kernel mapping
+            is found. If set to `False`, a `ValueError` will be raised if no kernel
+            mapping is found.
     """
 
     fallback_forward = cls.forward
@@ -138,8 +159,6 @@ 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(
@@ -167,14 +186,7 @@ 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:
-            # 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:
+            if needs_backward and not getattr(layer, "has_backward", True):
                 return fallback_forward(self, x, *args, **kwargs)
             return layer.forward(self, x, *args, **kwargs)
 
@@ -194,15 +206,8 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
 
         cached_layer[repo] = layer
 
-        # 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:
+        if needs_backward and not getattr(layer, "has_backward", True):
             return fallback_forward(self, x, *args, **kwargs)
-
         return layer.forward(self, x, *args, **kwargs)
 
     cls.forward = forward
@@ -211,11 +216,34 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
 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
+    [`register_layer_mapping`]. The device type is inferred from the first
     argument to `forward`.
+
+    Args:
+        layer_name (`str`):
+            The name to assign to the layer.
+        use_fallback (`bool`, *optional*, defaults to `True`):
+            Whether to use the fallback forward function if no kernel mapping
+            is found. If set to `False`, a `ValueError` will be raised if no kernel
+            mapping is found.
+
+    Example:
+        ```python
+        from kernels import use_kernel_forward_from_hub
+
+        @use_kernel_forward_from_hub(layer_name="LlamaRMSNorm")
+        class LlamaRMSNorm(nn.Module):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+
+            def forward(self, x):
+                # Original forward implementation
+                pass
+        ```
     """
 
     def decorator(cls):
@@ -261,8 +289,7 @@ 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
-    # verify if : difference ⊄ {"can_torch_compile", "has_backward"}
-    if not difference <= {"can_torch_compile", "has_backward"}:
+    if difference != set() and difference != {"has_backward"}:
         raise TypeError("Layer must not contain additional members.")
 
     # Check whether the forward signatures are similar.
@@ -279,19 +306,3 @@ 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

src/kernels/utils.py

@@ -157,6 +157,27 @@ def install_kernel_all_variants(
 
 
 def get_kernel(repo_id: str, revision: str = "main") -> ModuleType:
+    """
+    Load a kernel from the kernel hub.
+
+    This function downloads a kernel to the local Hugging Face Hub cache
+    directory (if it was not downloaded before) and then loads the kernel.
+
+    Args:
+        repo_id (`str`): The Hub repository containing the kernel.
+        revision (`str`, *optional*, defaults to `"main"`): The specific
+            revision (branch, tag, or commit) to download.
+
+    Returns:
+        `ModuleType`: The imported kernel module.
+
+    Example:
+        ```python
+        from kernels import get_kernel
+        kernel = get_kernel("username/my-kernel")
+        result = kernel.kernel_function(input_data)
+        ```
+    """
     package_name, package_path = install_kernel(repo_id, revision=revision)
     return import_from_path(package_name, package_path / package_name / "__init__.py")
 
@@ -164,7 +185,20 @@ def get_kernel(repo_id: str, revision: str = "main") -> ModuleType:
 def has_kernel(repo_id: str, revision: str = "main") -> bool:
     """
     Check whether a kernel build exists for the current environment
-    (Torch version and compute framework).
+
+    This function checks whether there exists a kernel build for the current
+    environment (Torch version, compute framework and architecture).
+
+    Args:
+        repo_id (`str`):
+            The Hub repository containing the kernel.
+        revision (`str`, *optional*, defaults to `"main"`):
+            The kernel revision.
+
+    Returns:
+        `bool`:
+            `True` if a compatible kernel build exists for the current environment,
+            `False` otherwise.
     """
     package_name = package_name_from_repo_id(repo_id)
     variant = build_variant()
@@ -186,10 +220,25 @@ def has_kernel(repo_id: str, revision: str = "main") -> bool:
 
 def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
     """
-    Get a pre-downloaded, locked kernel.
+    Loads a pre-downloaded, locked kernel module from the local cache.
 
-    If `lockfile` is not specified, the lockfile will be loaded from the
-    caller's package metadata.
+    This function retrieves a kernel that was locked at a specific revision with
+    `kernels lock <project>` and then downloaded with `kernels download <project>`.
+
+    This function will fail if the kernel was not locked or downloaded. If you want
+    the kernel to be downloaded when it is not in the cache, use [`get_locked_kernel`]
+    instead.
+
+    Args:
+        repo_id (`str`):
+            The Hub repository containing the kernel.
+        lockfile (`Optional[Path]`, *optional*, defaults to `None`):
+            Path to a lockfile containing the commit SHA for the kernel. If `None`,
+            the lock information is automatically retrieved from the metadata of the
+            calling package.
+
+    Returns:
+        `ModuleType`: The imported kernel module corresponding to the locked version.
     """
     if lockfile is None:
         locked_sha = _get_caller_locked_kernel(repo_id)
@@ -234,7 +283,27 @@ def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
 
 
 def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleType:
-    """Get a kernel using a lock file."""
+    """
+    Loads a locked kernel module.
+
+    This function retrieves a kernel that was locked at a specific revision with
+    `kernels lock <project>`.
+
+    This function will download the locked kernel when it is not available in the
+    cache. If you want loading to fail if the kernel is not in the cache, use
+    [`load_kernel`] instead.
+
+    Args:
+        repo_id (`str`):
+            The Hub repository containing the kernel.
+        lockfile (`Optional[Path]`, *optional*, defaults to `None`):
+            Path to a lockfile containing the commit SHA for the kernel. If `None`,
+            the lock information is automatically retrieved from the metadata of the
+            calling package.
+
+    Returns:
+        `ModuleType`: The imported kernel module corresponding to the locked version.
+    """
     locked_sha = _get_caller_locked_kernel(repo_id)
 
     if locked_sha is None:

tests/test_layer.py

@@ -19,12 +19,6 @@ 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",
@@ -49,11 +43,6 @@ 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
@@ -112,29 +101,8 @@ 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",
-        "SiluAndMulNoCompile",
-    }
+    assert set(_KERNEL_MAPPING.get().keys()) == {"SiluAndMul", "SiluAndMulStringDevice"}
 
     extra_mapping1 = {
         "TestKernel": {
@@ -150,7 +118,6 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
-            "SiluAndMulNoCompile",
             "TestKernel",
         }
 
@@ -168,7 +135,6 @@ def test_mapping_contexts():
             assert set(_KERNEL_MAPPING.get().keys()) == {
                 "SiluAndMul",
                 "SiluAndMulStringDevice",
-                "SiluAndMulNoCompile",
                 "TestKernel",
             }
             assert (
@@ -179,7 +145,6 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
-            "SiluAndMulNoCompile",
             "TestKernel",
         }
         assert (
@@ -199,7 +164,6 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
-            "SiluAndMulNoCompile",
             "TestKernel",
         }
         assert (
@@ -210,7 +174,6 @@ def test_mapping_contexts():
     assert set(_KERNEL_MAPPING.get().keys()) == {
         "SiluAndMul",
         "SiluAndMulStringDevice",
-        "SiluAndMulNoCompile",
     }