Join us on Discord for questions and discussions!
This repo contains a Nix package that can be used to build custom machine learning kernels for PyTorch. The kernels are built using the PyTorch C++ Frontend and can be loaded from the Hub with the kernels Python package.
This builder is a core component of the larger kernel build/distribution system.
We recommend using Nix to build kernels. To speed up builds, first enable the Hugging Face binary cache:
# Install cachix and configure the cache
cachix use huggingface
# Or run once without installing cachix
nix run nixpkgs#cachix -- use huggingfaceThen quick start a build with:
cd examples/relu
nix run .#build-and-copy \
--max-jobs 2 \
--cores 8 \
-LWhere --max-jobs specifies the number of build variant that should be built concurrently and --cores the number of CPU cores that should be used per build variant.
The compiled kernel will then be available in the local build/ directory.
We also provide Docker containers for CI builds. For a quick build:
# Using the prebuilt container
cd examples/relu
docker run --rm \
--mount type=bind,source=$(pwd),target=/kernelcode \
-w /kernelcode ghcr.io/huggingface/kernel-builder:main buildSee dockerfiles/README.md for more options, including a user-level container for CI/CD environments.
| Hardware | Kernels Support | Kernel-Builder Support | Kernels Validated in CI | Tier |
|---|---|---|---|---|
| CUDA | β | β | β | 1 |
| ROCm | β | β | β | 2 |
| XPU | β | β | β | 2 |
| Metal | β | β | β | 2 |
| Huawei NPU | β | β | β | 3 |
- Writing Hub kernels
- Building kernels with Nix
- Framework-specific notes:
- Building kernels with Docker (for systems without Nix)
- Local kernel development (IDE integration)
- Kernel security
- Why Nix?
The generated CMake build files are based on the vLLM build infrastructure.
