Helion x šŸ¤— HF Kernels: Building and Shipping Out-of-the-box Performant Kernels

TL;DR The HuggingFace Kernels project now has Helion support. This blog walks through how to build, autotune, and ship performant and portable Helion kernels via the Hugging Face Kernels project, allowing users to consume these kernels seamlessly. Introduction Helion is a high-level DSL for writing…

TL;DR The HuggingFace Kernels project now has Helion support. This blog walks through how to build, autotune, and ship performant and portable Helion kernels via the Hugging Face Kernels project, allowing users to consume these kernels seamlessly. Introduction Helion is a high-level DSL for writing high-performance, portable kernels for machine learning. The Kernels project lets kernel developers package and distribute their kernels on the Hugging Face Hub platform in a consistent and reproducible manner. It also lets kernel users consume these kernels seamlessly without managing dependency hell. In this post, we will discuss how Helion is supported within the Kernels project, how users can benefit from first-class autotuning support in Helion, and how to ship pre-tuned kernel configs to reduce cold-start times. We will also show examples of Helion kernels and how tuning them for specific problem sizes can yield performance benefits. P.S.: Throughout the rest of the post, we will refer to the Kernels project with ā€œkā€ in capital letters to distinguish it from actual ā€œkernelsā€. Intro: Helion Helion is a tiled DSL for writing performant ML kernels. The programming model is often described as ā€œPyTorch with tilesā€ – the kernel operates on PyTorch tensors, and tile-level operations are specified via ordinary PyTorch tensor operators. As a quick example, the following function shows a tiled matmul implemented in Helion: import torch, helion, helion.language as hl @helion.kernel() def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: m, k = x.size() k, n = y.size() out = torch.empty([m, n], dtype=x.dtype, device=x.device) for tile_m, tile_n in hl.tile([m, n]): acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) for tile_k in hl.tile(k): acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n]) out[tile_m, tile_n] = acc return out What makes Helion desirable is not just its concise syntax but what it leaves deliberately unspecified. When you write hl.tile, you say only that the iteration space should be tiled – not how large the tiles are, or how their data is fetched from memory. Helion turns these decisions into a search space to be autotuned over. Crucially, the autotuner does not merely sweep numerical parameters like tile sizes, it also searches over lowering strategies – the actual implementation of the kernel: which memory-access pattern to use (pointer arithmetic, block pointers, TMA), how to order and flatten nested loops, whether a reduction should be persistent or looped, and more. In Triton or CUDA, switching between these choices means rewriting the kernel entirely; in Helion the optimal choice is found algorithmically. This autotuning process is why a Helion kernel can often outperform a hand-written kernel in a lower-level language, when benchmarked on a large set of shapes. With that said, autotuning is sometimes a lengthy process, so it is beneficial to have an established approach for shipping a kernel bundled with pre-tuned configs. This is where the Kernels project comes in. Intro: Kernels The current landscape of kernel packaging and distribution is fragmented, characterized by inconsistent source structures, disparate tooling, and limited compatibility support. Consequently, users often face arduous build times, even when pre-built wheels are available. The Kernels project addresses these challenges by establishing a standardized, unified packaging and build process for both AoT and JIT kernels. The project is divided into two primary components: kernel-builder: A tool for developers to reliably package and distribute kernels across different framework versions and system configurations. It enforces standards to ensure predictable source structures, build reproducibility, native PyTorch compatibility, and easy community sharing. kernels: A consumer-facing Python library that allows users to effortlessly load ready-to-use kernels without dependency management issues via a simple command like get_kernel("org/name", version=1), much like pulling a model or dataset from the Hugging Face Hub. For kernel users, we want to provide a seamless experience of loading kernels and getting them ready to use right away. Let’s take a look at an example of how one could load the popular Flash-Attention 3 kernel: from kernels import get_kernel kernel_module = get_kernel("kernels-community/flash-attn3", version=1) flash_attn_func = kernel_module.flash_attn_func flash_attn_func(...) We provide prebuilt binaries for a comprehensive compatibility matrix of ahead-of-time kernels, such as Flash Attention 3. This is quite beneficial to end users, particularly when the kernel’s upstream repository may not have a specific build available. Users can browse a wide variety of kernels on the Hugging Face Hub platform: hf.co/kernels: We refer to this collection of kernels as the Kernels Hub. Packaging and using Helion in Kernels Helion kernels are plain Python. They compile themselves the first time you call them, so there is nothing for kernel-builder to compile ahead of time. Helion provides utilities to tune these kernels for specific workloads and hardware (more on that in a bit) so that users can tune once and reuse later. Helion kernels are also noarchĀ kernels: you ship the source, and Helion does the rest on the user’s machine. In this section, we discuss how to scaffold and structure a Helion kernel for building with kernel-builder.Ā  Start from the scaffold kernel-builder init gives you a working kernel project to edit: kernel-builder init --name myorg/vector-add-helion --backends cuda rocm xpu -- vector-add-helion cd vector-add-helion Note the --Ā before the directory name. --backendsĀ takes any number of values, so without the separator the directory name is read as another backend. The scaffold assumes a compiled kernel, so delete the parts you don’t need: rm -rf vector_add_helion_cuda vector_add_helion_xpu torch-ext/torch_binding.{cpp,h} That leaves three files to edit. build.toml [general] name = "vector-add-helion" license = "Apache-2.0" backends = ["cuda", "rocm", "xpu"] version = 1 edition = 5 python-depends = ["helion"] [general.hub] repo-id = "myorg/vector-add-helion" [torch-noarch] Two things worth paying heed to: python-depends = ["helion"] records that the kernel needs Helion at runtime. When someone loads the kernel, kernels checks that Helion is importable and gives a clear error if it isn’t. [torch-noarch] says there is no ahead-of-time compilation.Ā  torch-ext/vector_add_helion/__init__.py import helion import helion.language as hl import torch @helion.kernel(config=helion.Config(block_sizes=[1024], num_warps=4)) def vector_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: out = torch.empty_like(x) for tile in hl.tile(x.size(0)): out[tile] = x[tile] + y[tile] return out __all__ = ["vector_add"] This example here uses a hardcoded config, pinned via config=. In a later section we’ll go into details on pre-tuning and shipping a decision tree of configs covering a large set of shapes. flake.nix The scaffolded one needs no changes: { inputs.kernel-builder.url = "github:huggingface/kernels"; outputs = { self, kernel-builder, ... }: kernel-builder.lib.genKernelFlakeOutputs { inherit self; path = ./.; }; } Build and publish kernel-builder check-config . kernel-builder build-and-copy . Build and publish the builds to the Hub: kernel-builder build-and-upload . The build produces one directory per backend, each holding your __init__.py alongside a generatedĀ metadata.jsonĀ that carries the Helion dependency forward: { "name": "vector-add-helion", "python-depends": ["helion"], "backend": { "type": "cuda" } } Below is an example of the published kernel on the Hub: sayakpaul/vector-add-helion. Using the kernel from kernels import get_kernel kernel = get_kernel("myorg/vector-add-helion", version=1, trust_remote_code=True) out = kernel.vector_add(x, y) Users would need Helion installed (pip install helion) and a version of kernels that knows about it. Older releases reject the dependency outright with unsupported kernel dependency: helion. Pre-tuning and Shipping Pre-tuned configs To ensure the shipped Helion kernel performs well on different input problem shapes and different GPU generations, it’s often beneficial to pre-tune them. The workflow has three steps: tune the kernel across a representative set of shapes, let Helion build a decision tree that maps each shape to its best config ship that tree beside your kernel. At load time, Helion reads the tree (stored as a plain python file next to the kernel source file) and picks a config per call – no tuning on the user’s machine. To do this, the first kernel source change is a decorator: --- @helion.kernel(config=...) +++ @helion.aot_kernel(static_shapes=True) Write a small script that calls the kernel on every shape you want covered during pre-tuning: # bench.py import torch from vector_add_helion.vector_add import vector_add for n in [1024, 1

Source: PyTorch — Published — Category: Open Source

šŸ”— Read full article on PyTorch →