Skip to content

huggingface/kernel-builder

Repository files navigation

kernel-builder

kernel-builder logo

Build and Push Docker Image GitHub tag GitHub package


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.

πŸš€ Quick Start

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 huggingface

Then quick start a build with:

cd examples/relu
nix run .#build-and-copy \
  --max-jobs 2 \
  --cores 8 \
  -L

Where --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 build

See dockerfiles/README.md for more options, including a user-level container for CI/CD environments.

🎯 Hardware Support

Hardware Kernels Support Kernel-Builder Support Kernels Validated in CI Tier
CUDA βœ“ βœ“ βœ“ 1
ROCm βœ“ βœ“ βœ— 2
XPU βœ“ βœ“ βœ— 2
Metal βœ“ βœ“ βœ— 2
Huawei NPU βœ“ βœ— βœ— 3

πŸ“š Documentation

Credits

The generated CMake build files are based on the vLLM build infrastructure.

About

πŸ‘· Build compute kernels

Resources

License

Security policy

Stars

Watchers

Forks

Packages