Add support for kernel versions

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,25 @@ 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. The kernel version is
+bumped is bumped in the following circumstances:
 
-Kernels are versioned using tags of the form `v<major>.<minor>.<patch>`.
-You can specify which version to download using Python version specifiers:
+* The kernel API changes in an incompatible way.
+* Support for an older PyTorch version is dropped.
 
-```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.
+In this way, you can ensure 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

@@ -40,21 +40,30 @@ 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, for example:
+Currently the metadata specifies the kernel's versin and Python dependencies,
+for example:
 
 ```json
-{ "python-depends": ["nvidia-cutlass-dsl"] }
+{
+  "python-depends": ["nvidia-cutlass-dsl"],
+  "version": 1
+}
 ```
 
 The following dependencies are the only ones allowed at this stage: `einops` and `nvidia-cutlass-dsl`
 
 ## 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`.
 
-We recommend using [semver](https://semver.org/) to version kernels.
+The version **must** be bumped in the following cases:
+
+- The kernel API is changed in an incompatible way.
+- Support for one or more PyTorch version is dropped. This is usually done
+  when the API is extended and the builds for older PyTorch versions are not
+  updated anymore to have the new extensions.
 
 ## Native Python module
 

docs/source/layers.md

@@ -214,20 +214,18 @@ kernel_layer_mapping = {
         "cuda": LayerRepository(
             repo_id="kernels-community/activation",
             layer_name="SiluAndMul",
-            version=">=0.0.4,<0.1.0",
+            version=1,
         ),
         "rocm": LayerRepository(
             repo_id="kernels-community/activation",
             layer_name="SiluAndMul",
-            version=">=0.0.4,<0.1.0",
+            version=1,
         )
     }
 }
 ```
 
-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.
+This will get the layer from lates version on the version 1 branch.
 
 ### Registering kernels for specific modes
 
@@ -242,10 +240,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 +273,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 +313,7 @@ kernel_layer_mapping = {
         ): LayerRepository(
             repo_id="kernels-community/activation",
             layer_name="SiluAndMul",
+            version=1,
         ),
         Device(
             type="cuda",
@@ -319,6 +323,7 @@ kernel_layer_mapping = {
         ): LayerRepository(
             repo_id="kernels-community/activation-hopper",
             layer_name="SiluAndMul",
+            version=1,
         ),
     }
 }
@@ -359,6 +364,7 @@ with use_kernel_mapping(
                 repo_path="/home/daniel/kernels/activation",
                 package_name="activation",
                 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

src/kernels/_versions.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Dict, Optional
 
 from huggingface_hub import HfApi
@@ -6,9 +7,27 @@
 from packaging.version import InvalidVersion, Version
 
 
-def _get_available_versions(repo_id: str) -> Dict[Version, GitRefInfo]:
+def _get_available_versions(repo_id: str) -> Dict[int, GitRefInfo]:
     """Get kernel versions that are available in the repository."""
     versions = {}
+    for branch in HfApi().list_repo_refs(repo_id).branches:
+        if not branch.name.startswith("v"):
+            continue
+        try:
+            versions[int(branch.name[1:])] = branch
+        except ValueError:
+            continue
+
+    return versions
+
+
+def _get_available_versions_old(repo_id: str) -> Dict[Version, GitRefInfo]:
+    """
+    Get kernel versions that are available in the repository.
+
+    This is for the old tag-based versioning scheme.
+    """
+    versions = {}
     for tag in HfApi().list_repo_refs(repo_id).tags:
         if not tag.name.startswith("v"):
             continue
@@ -20,33 +39,49 @@ def _get_available_versions(repo_id: str) -> Dict[Version, GitRefInfo]:
     return versions
 
 
-def resolve_version_spec_as_ref(repo_id: str, version_spec: str) -> GitRefInfo:
+def resolve_version_spec_as_ref(repo_id: str, version_spec: int | str) -> GitRefInfo:
     """
     Get the locks for a kernel with the given version spec.
 
     The version specifier can be any valid Python version specifier:
     https://packaging.python.org/en/latest/specifications/version-specifiers/#version-specifiers
     """
-    versions = _get_available_versions(repo_id)
-    requirement = SpecifierSet(version_spec)
-    accepted_versions = sorted(requirement.filter(versions.keys()))
-
-    if len(accepted_versions) == 0:
-        raise ValueError(
-            f"No version of `{repo_id}` satisfies requirement: {version_spec}"
+    if isinstance(version_spec, int):
+        versions = _get_available_versions(repo_id)
+        ref = versions.get(version_spec, None)
+        if ref is None:
+            raise ValueError(
+                f"Version {version_spec} not found, available versions: {', '.join(sorted(str(v) for v in versions.keys()))}"
+            )
+        return ref
+    else:
+        warnings.warn(
+            "Version specifiers are deprecated, support will be removed in a future `kernels` version,\na concrete version instead."
         )
+        versions_old = _get_available_versions_old(repo_id)
+        requirement = SpecifierSet(version_spec)
+        accepted_versions = sorted(requirement.filter(versions_old.keys()))
 
-    return versions[accepted_versions[-1]]
+        if len(accepted_versions) == 0:
+            raise ValueError(
+                f"No version of `{repo_id}` satisfies requirement: {version_spec}"
+            )
+
+        return versions_old[accepted_versions[-1]]
 
 
 def select_revision_or_version(
-    repo_id: str, revision: Optional[str], version: Optional[str]
+    repo_id: str,
+    *,
+    revision: Optional[str],
+    version: Optional[int | str],
 ) -> str:
     if revision is not None and version is not None:
-        raise ValueError("Either a revision or a version must be specified, not both.")
-    elif revision is None and version is None:
-        revision = "main"
+        raise ValueError("Only one of `revision` or `version` must be specified.")
+
+    if revision is not None:
+        return revision
     elif version is not None:
-        revision = resolve_version_spec_as_ref(repo_id, version).target_commit
-    assert revision is not None
-    return revision
+        return resolve_version_spec_as_ref(repo_id, version).target_commit
+
+    return "main"

src/kernels/cli.py

@@ -1,20 +1,17 @@
 import argparse
 import dataclasses
 import json
-import re
 import sys
 from pathlib import Path
 
-from huggingface_hub import create_repo, upload_folder, create_branch
 
 from kernels.compat import tomllib
 from kernels.lockfile import KernelLock, get_kernel_locks
+from kernels.upload import upload_kernels_dir
 from kernels.utils import install_kernel, install_kernel_all_variants
 
 from .doc import generate_readme_for_kernel
 
-BUILD_VARIANT_REGEX = re.compile(r"^(torch\d+\d+|torch-(cpu|cuda|metal|rocm|xpu))")
-
 
 def main():
     parser = argparse.ArgumentParser(
@@ -69,11 +66,13 @@ def main():
     upload_parser.add_argument(
         "--repo-id",
         type=str,
+        required=True,
         help="Repository ID to use to upload to the Hugging Face Hub",
     )
     upload_parser.add_argument(
         "--branch",
-        type=None,
+        type=str,
+        default=None,
         help="If set, the upload will be made to a particular branch of the provided `repo-id`.",
     )
     upload_parser.add_argument(
@@ -169,46 +168,12 @@ def lock_kernels(args):
 
 
 def upload_kernels(args):
-    # Resolve `kernel_dir` to be uploaded.
-    kernel_dir = Path(args.kernel_dir).resolve()
-
-    build_dir = None
-    for candidate in [kernel_dir / "build", kernel_dir]:
-        variants = [
-            variant_path
-            for variant_path in candidate.glob("torch*")
-            if BUILD_VARIANT_REGEX.match(variant_path.name) is not None
-        ]
-        if variants:
-            build_dir = candidate
-            break
-    if build_dir is None:
-        raise ValueError(
-            f"Couldn't find any build variants in: {kernel_dir.absolute()} or {(kernel_dir / 'build').absolute()}"
-        )
-
-    repo_id = create_repo(
-        repo_id=args.repo_id, private=args.private, exist_ok=True
-    ).repo_id
-
-    if args.branch is not None:
-        create_branch(repo_id=repo_id, branch=args.branch, exist_ok=True)
-
-    delete_patterns: set[str] = set()
-    for build_variant in build_dir.iterdir():
-        if build_variant.is_dir():
-            delete_patterns.add(f"{build_variant.name}/**")
-
-    upload_folder(
-        repo_id=repo_id,
-        folder_path=build_dir,
-        revision=args.branch,
-        path_in_repo="build",
-        delete_patterns=list(delete_patterns),
-        commit_message="Build uploaded using `kernels`.",
-        allow_patterns=["torch*"],
-    )
-    print(f"✅ Kernel upload successful. Find the kernel in https://hf.co/{repo_id}.")
+    upload_kernels_dir(
+        Path(args.kernel_dir).resolve(),
+        repo_id=args.repo_id,
+        branch=args.branch,
+        private=args.private,
+    )
 
 
 class _JSONEncoder(json.JSONEncoder):