Add support for kernel versions (#209)

* Add support for kernel versions

This change adds support for kernel versions branches as
implemented by `kernel-builder`:

* `kernels upload` will upload the branch corresponding to the
  kernel's version.
* `get_kernel`, `LayerRepository`, etc. accept version branches
  through the `version` argument.
* Kernel locking (e.g. `kernels lock`) support reading version
  branches from `pyproject.yaml` and correctly locking them.

All the kernel functions and layers already supported a `version`
argument for version specifiers (based on tags). These are now
overloaded -- if `version` is a `str`, it is parsed as a version
specifier. If it is an `int`, it is used as a version branch.

This overloading is used to gradually move to version branches,
deprecating version specifiers. For now, using version specifiers
will emit a warning.

* Documtenation updates

* Some fixes in the docs about bumping the kernel version

* Add basic `kernel versions` subcommand

* Describe migration

* Warn when trying to use without versions

* Fixes

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

* black

* Remove unnecessary `_get_available_versions` call

* Add files for old locking tests

* Remove spurious quotation

* Revert build variant regex

Get the full name by using the path in `kernels versions`.

* `kernels versions`: sort versions

---------

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

Author: danieldk

README.md

@@ -45,7 +45,7 @@ import torch
 from kernels import get_kernel
 
 # Download optimized kernels from the Hugging Face hub
-activation = get_kernel("kernels-community/activation")
+activation = get_kernel("kernels-community/activation", version=1)
 
 # Random tensor
 x = torch.randn((10, 10), dtype=torch.float16, device="cuda")

builder/tests/run-kernels/pyproject.toml

@@ -1,2 +1,2 @@
 [tool.kernels.dependencies]
-"kernels-test/versions" = ">=0.1.0,<0.2.0"
+"kernels-test/versions" = 1

docs/source/_toctree.yml

@@ -15,6 +15,9 @@
       title: Environment Variables
     - local: faq
       title: FAQ
+    - local: migration
+      title: Migrating from older versions
+  title: Usage Guide
   title: Using kernels
 - sections:
     - local: builder/writing-kernels

docs/source/basic-usage.md

@@ -9,7 +9,7 @@ import torch
 from kernels import get_kernel
 
 # Download optimized kernels from the Hugging Face hub
-activation = get_kernel("kernels-community/activation")
+activation = get_kernel("kernels-community/activation", version=1)
 
 # Create a random tensor
 x = torch.randn((10, 10), dtype=torch.float16, device="cuda")
@@ -21,30 +21,22 @@ activation.gelu_fast(y, x)
 print(y)
 ```
 
-### Using version bounds
+This fetches version `1` of the kernel `kernels-community/activation`.
+Kernels are versioned using a major version number. Using `version=1` will
+get the latest kernel build from the `v1` branch.
 
-Kernels are versioned using tags of the form `v<major>.<minor>.<patch>`.
-You can specify which version to download using Python version specifiers:
-
-```python
-import torch
-from kernels import get_kernel
-
-activation = get_kernel("kernels-community/activation", version=">=0.0.4,<0.1.0")
-```
-
-This will get the latest kernel tagged `v0.0.z` where `z` is at least 4. It
-is strongly recommended to specify a version bound, since a kernel author
-might push incompatible changes to the `main` branch.
+Kernels within a version branch must never break the API or remove builds
+for older PyTorch versions. This ensures that your code will continue to work.
 
 ## Checking Kernel Availability
 
-You can check if a specific kernel is available for your environment:
+You can check if a particular version of a kernel supports the environment
+that the program is running on:
 
 ```python
 from kernels import has_kernel
 
 # Check if kernel is available for current environment
-is_available = has_kernel("kernels-community/activation")
+is_available = has_kernel("kernels-community/activation", version=1)
 print(f"Kernel available: {is_available}")
 ```

docs/source/kernel-requirements.md

@@ -37,7 +37,14 @@ must be available for that combination.
 ## Kernel metadata
 
 The build variant directory can optionally contain a `metadata.json` file.
-Currently the only purpose of the metadata is to specify the kernel Python dependencies.
+Currently the metadata specifies the kernel's version and Python dependencies,
+for example:
+
+```json
+{
+  "python-depends": ["einops"],
+  "version": 1
+```
 
 ### Python dependencies
 
@@ -75,7 +82,8 @@ You can specify both general and backend-specific dependencies:
   "python-depends": ["einops"],
   "python-depends-backends": {
     "cuda": ["nvidia-cutlass-dsl"]
-  }
+  },
+  "version": 1
 }
 ```
 
@@ -84,21 +92,31 @@ You can specify both general and backend-specific dependencies:
 The following dependencies are currently allowed:
 
 **General dependencies:**
+
 - `einops`
 
 **Backend-specific dependencies:**
+
 - CUDA: `nvidia-cutlass-dsl`
 - XPU: `onednn`
 
 Dependencies are validated based on the backend being used. When a kernel is loaded, only the dependencies relevant to the active backend are checked.
 
 ## Versioning
 
-Kernels are versioned on the Hub using Git tags. Version tags must be of
-the form `v<major>.<minor>.<patch>`. Versions are used by [locking](./locking.md)
-to resolve the version constraints.
+Kernels are versioned using a major version. The kernel revisions of a
+version are stored in a branch of the form `v<version>`. Each build
+variant will also have the kernel version in `metadata.json`.
+
+The version **must** be bumped in the following cases:
+
+- The kernel API is changed in an incompatible way.
+- The API is extended in a compatible way, but not all build variants
+  receive the extension (e.g. because they are for older Torch versions
+  that are not supported by `kernel-builder` anymore).
 
-We recommend using [semver](https://semver.org/) to version kernels.
+In both cases, build variants that are not updated must be removed from
+the new version's branch.
 
 ## Native Python module
 

docs/source/layers.md

@@ -159,10 +159,12 @@ kernel_layer_mapping = {
         "cuda": LayerRepository(
             repo_id="kernels-community/activation",
             layer_name="SiluAndMul",
+            version=1,
         ),
         "rocm": LayerRepository(
             repo_id="kernels-community/activation",
             layer_name="SiluAndMul",
+            version=1,
         )
     }
 }
@@ -202,33 +204,6 @@ kernel_layer_mapping = {
 }
 ```
 
-### Using version bounds
-
-Kernels are versioned using tags of the form `v<major>.<minor>.<patch>`.
-You can specify which version of the kernel to download using Python version
-specifiers:
-
-```python
-kernel_layer_mapping = {
-    "SiluAndMul": {
-        "cuda": LayerRepository(
-            repo_id="kernels-community/activation",
-            layer_name="SiluAndMul",
-            version=">=0.0.4,<0.1.0",
-        ),
-        "rocm": LayerRepository(
-            repo_id="kernels-community/activation",
-            layer_name="SiluAndMul",
-            version=">=0.0.4,<0.1.0",
-        )
-    }
-}
-```
-
-This will get the layer from latest kernel tagged `v0.0.z` where `z` is at
-least 4. It is strongly recommended to specify a version bound, since a
-kernel author might push incompatible changes to the `main` branch.
-
 ### Registering kernels for specific modes
 
 You might want to register two different kernels for a particular layer,
@@ -242,10 +217,12 @@ kernel_layer_mapping = {
           Mode.INFERENCE: LayerRepository(
               repo_id="kernels-community/activation-inference-optimized",
               layer_name="SiluAndMul",
+              version=1,
           ),
           Mode.TRAINING | Mode.TORCH_COMPILE: LayerRepository(
               repo_id="kernels-community/activation-training-optimized",
               layer_name="SiluAndMul",
+              version=1,
           ),
       }
     }
@@ -273,14 +250,17 @@ kernel_layer_mapping = {
             Mode.FALLBACK: LayerRepository(
                 repo_id="kernels-community/activation",
                 layer_name="SiluAndMul",
+                version=1,
             ),
             Mode.INFERENCE: LayerRepository(
                 repo_id="kernels-community/activation-inference-optimized",
                 layer_name="SiluAndMul",
+                version=1,
             ),
             Mode.TRAINING: LayerRepository(
                 repo_id="kernels-community/activation-training-optimized",
                 layer_name="SiluAndMul",
+                version=1,
             ),
         }
     }
@@ -310,6 +290,7 @@ kernel_layer_mapping = {
         ): LayerRepository(
             repo_id="kernels-community/activation",
             layer_name="SiluAndMul",
+            version=1,
         ),
         Device(
             type="cuda",
@@ -319,6 +300,7 @@ kernel_layer_mapping = {
         ): LayerRepository(
             repo_id="kernels-community/activation-hopper",
             layer_name="SiluAndMul",
+            version=1,
         ),
     }
 }

docs/source/locking.md

@@ -10,7 +10,7 @@ requires = ["kernels", "setuptools"]
 build-backend = "setuptools.build_meta"
 
 [tool.kernels.dependencies]
-"kernels-community/activation" = ">=0.0.1"
+"kernels-community/activation" = 1
 ```
 
 Then run `kernels lock .` in the project directory. This generates a `kernels.lock` file with

docs/source/migration.md

@@ -0,0 +1,68 @@
+# Migrating from older versions
+
+## 0.12
+
+### Adopting kernel versions
+
+Before `kernels` 0.12, kernels could be pulled from a repository
+without specifying a version. This is deprecated in kernels 0.12
+and will become an error in kernels 0.14. Instead, use of a kernel
+should always specify a version (except for local kernels).
+
+Kernels only use a major version. The kernel maintainer is responsible
+for never breaking a kernel within a major version and should bump up
+the major version if the kernel API changes and/or when support for
+older Torch versions is removed.
+
+You can find the versions that are supported by a kernel using the
+`kernels versions command`. For example:
+
+```bash
+$ kernels versions kernels-community/activation
+Version 1: torch210-cxx11-cu126-x86_64-linux, torch210-cxx11-cu128-x86_64-linux, torch210-cxx11-cu130-x86_64-linux, torch27-cxx11-cu118-x86_64-linux, torch27-cxx11-cu126-x86_64-linux, torch27-cxx11-cu128-aarch64-linux, torch27-cxx11-cu128-x86_64-linux ✅, torch28-cxx11-cu126-aarch64-linux, torch28-cxx11-cu126-x86_64-linux, torch28-cxx11-cu128-aarch64-linux, torch28-cxx11-cu128-x86_64-linux, torch28-cxx11-cu129-aarch64-linux, torch28-cxx11-cu129-x86_64-linux, torch29-cxx11-cu126-aarch64-linux, torch29-cxx11-cu126-x86_64-linux, torch29-cxx11-cu128-aarch64-linux, torch29-cxx11-cu128-x86_64-linux, torch29-cxx11-cu130-aarch64-linux, torch29-cxx11-cu130-x86_64-linux
+```
+
+The command lists all available versions (here only version 1) with
+all the variants that are supported. A check mark is printed after
+the variant that is compatible with your current environment.
+
+Code that uses a kernel can be updated as follows:
+
+```python
+# Old:
+activation = get_kernel("kernels-community/activation")
+activation = get_kernel("kernels-community/activation", version=">=0.0.2 && <0.1.0")
+
+# New:
+activation = get_kernel("kernels-community/activation", version=1)
+
+# Old:
+kernel_layer_mapping = {
+    "SiluAndMul": {
+        "cuda": LayerRepository(
+            repo_id="kernels-community/activation",
+            layer_name="SiluAndMul",
+        ),
+    }
+}
+kernel_layer_mapping = {
+    "SiluAndMul": {
+        "cuda": LayerRepository(
+            repo_id="kernels-community/activation",
+            layer_name="SiluAndMul",
+            version=">=0.0.2 && <0.1.0",
+        ),
+    }
+}
+
+# New:
+kernel_layer_mapping = {
+    "SiluAndMul": {
+        "cuda": LayerRepository(
+            repo_id="kernels-community/activation",
+            layer_name="SiluAndMul",
+            version=1,
+        ),
+    }
+}
+```

examples/basic.py

@@ -5,7 +5,7 @@
 print("Starting examples/basic.py demo")
 
 # Download optimized kernels from the Hugging Face hub
-activation = get_kernel("kernels-community/activation")
+activation = get_kernel("kernels-community/activation", version=1)
 
 print("Activation kernel fetched")
 

kernels/src/kernels/__init__.py

@@ -1,6 +1,5 @@
 import importlib.metadata
 
-
 __version__ = importlib.metadata.version("kernels")
 
 from kernels.layer import (