#
Writing Hub kernels with kernel-builder

##
Introduction

The Kernel Hub allows Python libraries and applications to load compute kernels directly from the Hub. 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.

`kernel-builder` is a set of tools that can build conforming kernels. It takes care of:

- Building kernels for all supported PyTorch configurations (C++98/11 and different CUDA versions).
- Compatibility with old glibc and libstdc++ versions, so that kernels also work on older Linux distributions.
- Registering Torch ops, such that multiple versions the same kernel can be loaded without namespace conflicts.

`kernel-builder` builds are configured through a `build.toml` file. `build.toml` is a simple format that does not require intricate knowledge of CMake or setuptools.

This page describes the directory layout of a kernel-builder project, the format of the `build.toml` file, and some additional Python glue that `kernel-builder` provides. We will use a simple ReLU kernel as the running example. After reading this page, you may also want to have a look at the more realistic ReLU kernel with backprop and `torch.compile` support.

##
Setting up environment

###
Quick install

The fastest way to get started is to run the install script. This installs Determinate Nix and `kernel-builder` in a single command:

    curl -fsSL https://raw.githubusercontent.com/huggingface/kernels/main/install.sh | bash

This will:

1. Install Determinate Nix (if not already installed).
2. Configure the Hugging Face binary cache (to avoid building dependencies from source).
3. Install `kernel-builder` via `nix profile install`.

To update `kernel-builder` later:

    nix profile upgrade --all

For a step-by-step breakdown of what the script does, see Using the kernel builder with Nix.

###
Cloud environment

In the `terraform` directory, we provide an example of programatically spinning up an EC2 instance that is ready with everything needed for you to start developing and building kernels.

If you use a different provider, the Terraform bridges should be similar and straightforward to modify.

##
Starting a new kernel

The easiest way to start a new kernel is by using the `init` subcommand of `kernel-builder`. This creates a minimal, compilable kernel:

    $ kernel-builder init --name myorg/mykernel
    Initialized `myorg/mykernel` at /home/daniel/git/kernels/examples/kernels/mykernel

This creates a kernel named `mykernel` in the directory `mykernel`. The kernel is configured to upload to the `myorg/mykernel` Hub repository when an upload command is used.

By default, the `init` subcommand creates a CUDA kernel. You can specify another backend with the `--backends` option:

    $ kernel-builder init --name myorg/mykernel --backends xpu

You can also make a multi-backend kernel by adding all the backends that you would like to support as arguments to `--backends`:

    $ kernel-builder init --name myorg/mykernel --backends cuda xpu
    Initialized `myorg/mykernel` at /home/daniel/git/kernels/examples/kernels/mykernel

Finally, if you want to create a kernel for all supported backends, you can use `--backends all`.

##
Kernel project layout

Kernel projects follow this general directory layout:

    mykernel
    ├── benchmarks
    │   └── benchmark.py
    ├── build.toml
    ├── CARD.md
    ├── example.py
    ├── flake.nix
    ├── mykernel_cuda
    │   └── mykernel.cu
    ├── tests
    │   ├── __init__.py
    │   └── test_mykernel.py
    └── torch-ext
    ├── mykernel
    │   └── __init__.py
    ├── torch_binding.cpp
    └── torch_binding.h

In this example we can find:

- The build configuration in `build.toml`.
- One or more top-level directories containing kernels (`mykernel_cuda`).
- The `torch-ext` directory, which contains:
  - `torch_binding.h`: contains declarations for kernel entry points (from `kernel_a` and `kernel_b`).
  - `torch_binding.cpp`: registers the entry points as Torch ops.
  - `torch_ext/mykernel`: contains any Python wrapping the kernel needs. At the bare minimum, it should contain an `__init__.py` file.
- Kernel tests in the directory `tests`.
- Benchmarks in the directory `benchmarks`.
- A kernel card template in `CARD.md`. This placeholders in the card are filled during the kernel build.
- The Nix flake configuration in `flake.nix`.
- An example script that uses the kernel in `example.py`.

##
build.toml

`build.toml` tells `kernel-builder` what to build and how. It looks as follows for the `mykernel` kernel:

    [general]
    backends = [
      "cuda",
    ]
    name = "mykernel"
    version = 1

    [general.hub]
    repo-id = "myorg/mykernel"

    [torch]
    src = [
      "torch-ext/torch_binding.cpp",
      "torch-ext/torch_binding.h",
    ]

    [kernel.mykernel]
    backend = "cuda"
    depends = ["torch"]
    src = ["mykernel_cuda/mykernel.cu"]
    # If the kernel is only supported on specific capabilities, set the
    # cuda-capabilities option:
    #
    # cuda-capabilities = [ "9.0", "10.0", "12.0" ]

The following sections enumerate all supported options for `build.toml`.

###
general

- `name` (required): the name of the kernel. The Python code for a Torch extension must be stored in `torch-ext/<name>`.
- `version` (int): the major version of the kernel. The version is written to the kernel’s `metadata.json` and is used by the `kernels upload` command to upload the kernel to a version branch named `v<version>`.
- `backends` (required): a list of supported backends. Must be one or more of `cpu`, `cuda`, `metal`, `rocm`, or `xpu`.
- `python-depends` (experimental): a list of additional Python dependencies that the kernel requires. The only supported dependencies are `einops` and `nvidia-cutlass-dsl`.

###
general.hub

- `repo-id`: the Hub repository to upload the kernel to when the `upload` or `build-and-upload` subcommands of `kernel-builder` are used.

###
general.cuda

- `maxver`: the maximum CUDA toolkit version (inclusive). This option must not be set under normal circumstances, since it can exclude Torch build variants that are required for compliant kernels. This option is provided for kernels that cause compiler errors on newer CUDA toolkit versions.
- `minver`: the minimum required CUDA toolkit version. This option must not be set under normal circumstances, since it can exclude Torch build variants that are required for compliant kernels. This option is provided for kernels that require functionality only provided by newer CUDA toolkits.

###
torch

This section describes the Torch extension. In the future, there may be similar sections for other frameworks. This section has the following options:

- `src` (required): a list of source files and headers.
- `pyext` (optional): the list of extensions for Python files. Default: `["py", "pyi"]`.
- `include` (optional): include directories relative to the project root. Default: `[]`.
- `maxver` (optional): only build for this Torch version and earlier. Use cautiously, since this option produces non-compliant kernels if the version range does not correspond to the required variants.
- `minver` (optional): only build for this Torch version and later. Use cautiously, since this option produces non-compliant kernels if the version range does not correspond to the required variants.

###
kernel.<name>

Specification of a kernel with the name `<name>`. Multiple `kernel.<name>` sections can be defined in the same `build.toml`. See for example `kernels-community/quantization` for an example with multiple kernel sections.

The following options can be set for a kernel:

- `backend` (required): the compute backend of the kernel. The currently supported backends are `cpu`, `cuda`, `metal`, `rocm`, and `xpu`. The `cpu` backend is currently experimental and might still change.
- `depends` (required): a list of dependencies. The supported dependencies are listed in `deps.nix`.
- `src` (required): a list of source files and headers.
- `include` (optional): include directories relative to the project root. Default: `[]`.

Besides these shared options, the following backend-specific options are available:

####
cuda

- `cuda-capabilities` (optional): a list of CUDA capabilities that the kernel should be compiled for. When absent, the kernel will be built using all capabilities that the builder supports. The effective capabilities are the intersection of this list and the capabilities supported by the CUDA compiler. It is recommended to leave this option unspecified unless a kernel requires specific capabilities.
- `cuda_flags` (optional): additional flags to be passed to `nvcc`. Warning: this option should only be used in exceptional circumstances. Custom compile flags can interfere with the build process or break compatibility requirements.

####
rocm

- `rocm-archs`: a list of ROCm architectures that the kernel should be compiled for.

####
xpu

- `sycl_flags`: a list of additional flags to be passed to the SYCL compiler.

##
Torch bindings

###
Defining bindings

Torch bindings are defined in C++, kernels commonly use two files:

- `torch_binding.h` containing function declarations.
- `torch_binding.cpp` registering the functions as Torch ops.

For instance, the `mykernel` kernel discussed above has the following declaration in `torch_binding.h`:

    #pragma once

    #include <torch/torch.h>

    void mykernel(torch::Tensor &out, torch::Tensor const &input);

This function is then registered as a Torch op in `torch_binding.cpp`:

    #include <torch/library.h>

    #include "registration.h"
    #include "torch_binding.h"

    TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      ops.def("mykernel(Tensor! out, Tensor input) -> ()");
    #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
      ops.impl("mykernel", torch::kCUDA, &mykernel);
    #endif
    }

    REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

This snippet uses macros from `registration.h` to register the function. `registration.h` is generated by `kernel-builder` itself. A function is registered through the `def`/`ops` methods. `ops` specifies the function signature following the function schema. `impl` associates the function name with the C/C++ function and the applicable device.

##
Using kernel functions from Python

The bindings are typically wrapped in Python code in `torch_ext/<name>`. The native code is exposed under the `torch.ops` namespace. However, we add some unique material to the name of the extension to ensure that different versions of the same extension can be loaded at the same time. As a result, the extension is registered as `torch.ops.<name>_<unique_material>`.

To deal with this uniqueness, `kernel_builder` generates a Python module named `_ops` that contains an alias for the name. This can be used to refer to the correct `torch.ops` module. For example:

    from typing import Optional

    import torch

    from ._ops import ops

    def mykernel(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
        if out is None:
            out = torch.empty_like(x)
        ops.mykernel(out, x)
        return out

##
Kernel tests

Kernel tests are stored in the `tests` directory. Since running all kernel tests in CI may be prohibitively expensive, the `pyproject.toml` generated by the builder adds support for the special `kernels_ci` PyTest marker that can be used as follows:

    import pytest

    @pytest.mark.kernels_ci
    def test_mykernel():
      ...

We recommend that you to pick tests that together would catch most error cases while running within 60 seconds.

You can run the tests (e.g. in CI) using:

    $ nix run .#ci-test

If the kernel supports multiple backends, it will run the test for the first supported backend that was found, obeying the following order: CUDA, ROCm, XPU, Metal, CPU. If you would like to the tests for a specific build variant, you can use `nix run .#ciTests.<variant>`. For instance:

    $ nix run .#ciTests.torch210-cxx11-cpu-x86_64-linux

When running the tests on a non-NixOS systems, make sure that the CUDA driver library can be found.

##
Kernel docs

We provide a utility to generate a system card for a given kernel, utilizing information from its `build.toml` and metadata. This system card provides a reasonable starting point and is meant to be edited afterward by the kernel developer.

The template card is generated as a part of `kernel-builder init` command and is serialized in the root directory of the kernel.

The card will be filled automatically by the builder when using the `build-and-upload` or `build-and-copy` command. It will be serialized to the `build` sub-directory inside the main kernel directory. It will be uploaded as `README.md` to the Hub.

Update on GitHub