Skip to content

MonakaResearch/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

Releases

No releases published

Packages

No packages published

Languages

  • Nix 38.9%
  • Rust 22.0%
  • Python 15.2%
  • CMake 11.7%
  • Shell 6.7%
  • PowerShell 3.9%
  • Other 1.6%