Helion 接入 Hugging Face Kernels:构建与分发高性能
Helion x 🤗 HF Kernels: Building and Shipping Out-of-the-box Performant Kernels
给做底层算子优化的同学看,Helion 结合 HF Kernels 解决了算子分发和自动调优的工程痛点,有完整可复现的代码流程。
TL;DR
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.
HuggingFace Kernels 项目现已支持 Helion。本文介绍了如何通过 Hugging Face Kernels 项目构建、自动调优并部署高性能且可移植的 Helion 内核,使用户能够无缝地使用这些内核。
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.
Helion 是一种高级领域特定语言(DSL),用于编写高性能、可移植的机器学习内核。Kernels 项目允许内核开发者以一致且可复现的方式在 Hugging Face Hub 平台上打包和分发他们的内核。它还允许内核用户无缝地消费这些内核,而无需陷入依赖地狱。
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.
在本文中,我们将讨论如何在 Kernels 项目中支持 Helion,用户如何从 Helion 的一等公民自动调优支持中受益,以及如何交付预调优的内核配置以减少冷启动时间。我们还将展示 Helion 内核的示例,并说明针对特定问题规模进行调优如何带来性能提升。
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”.
附注:在本文的其余部分,我们将用大写字母“K”指代 Kernels 项目,以区别于实际的“kernels”(内核)。
Intro: Helion
引言: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:
Helion 是一种基于分块(tiled)的 DSL,用于编写高性能的 ML 内核。其编程模型常被描述为“带有分块的 PyTorch”——内核操作于 PyTorch 张量之上,分块级别的操作通过普通的 PyTorch 张量算子指定。作为一个快速示例,以下函数展示了在 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 outimport 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 outWhat 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.
Helion 令人青睐的原因不仅在于其简洁的语法,更在于它故意留白的部分。当你编写 hl.tile 时,你仅声明迭代空间应当被分块——而不指定分块的大小,也不指定数据如何从内存中获取。Helion 将这些决策转化为一个待自动调优的搜索空间。关键在于,自动调优器不仅仅扫描如分块大小之类的数值参数,它还会搜索降低策略(lowering strategies)——即内核的实际实现:使用哪种内存访问模式(指针算术、块指针、TMA)、如何排序和展平嵌套循环、归约操作是持久化的还是循环式的,等等。在 Triton 或 CUDA 中,在这些选项之间切换意味着完全重写内核;而在 Helion 中,最优选择是通过算法自动找到的。
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.
这就是为什么在对大量形状进行基准测试时,Helion 内核通常能超越用低级语言手写内核的原因。话虽如此,自动调优有时是一个漫长的过程,因此建立一种随附预调优配置打包内核的分发方法是有利的。这正是 Kernels 项目发挥作用的地方。
Intro: Kernels
简介: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.
当前的内核打包和分发格局支离破碎,其特征是源代码结构不一致、工具各异以及兼容性支持有限。因此,即使有预构建的 wheel 包可用,用户也常常面临艰难的构建时间。
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:
Kernels 项目通过为 AOT(提前编译)和 JIT(即时编译)内核建立标准化、统一的打包和构建流程来解决这些挑战。该项目分为两个主要组件:
- 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.
- kernel-builder:一个供开发人员可靠地跨不同框架版本和系统配置打包和分发内核的工具。它强制执行标准,以确保可预测的源代码结构、构建可重复性、原生 PyTorch 兼容性以及易于社区共享。
- kernels:一个面向消费者的 Python 库,允许用户通过类似 get_kernel("org/name", version=1) 的简单命令轻松加载即用型内核,而无需担心依赖管理问题,这就像从 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:
对于内核用户,我们希望提供无缝的内核加载体验,使其能够立即准备就绪。让我们看看如何加载流行的 Flash-Attention 3 内核的一个示例:
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(...)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.
我们为涵盖广泛兼容性的提前编译内核(如 Flash Attention 3)提供了预构建的二进制文件。这对终端用户非常有益,特别是当内核的上游仓库可能没有特定的构建版本可用时。
Users can browse a wide variety of kernels on the Hugging Face Hub platform: hf.co/kernels:
用户可以在 Hugging Face Hub 平台上浏览各种各样的内核:hf.co/kernels:
We refer to this collection of kernels as the Kernels Hub.
我们将这一系列内核称为 Kernels Hub。
Packaging and using Helion in Kernels
在 Kernels 中打包和使用 Helion
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.
Helion 内核是纯 Python 代码。它们在首次调用时自行编译,因此 kernel-builder 无需提前编译任何内容。Helion 提供了针对特定工作负载和硬件调优这些内核的工具(稍后详细介绍),以便用户可以一次性调优并在后续复用。Helion 内核也是 noarch(无架构特定)内核:你分发源代码,其余工作由 Helion 在用户的机器上完成。
In this section, we discuss how to scaffold and structure a Helion kernel for building with kernel-builder.
在本节中,我们将讨论如何使用 kernel-builder 来生成和构建 Helion 内核的结构。
Start from the scaffold
从脚手架开始
kernel-builder init gives you a working kernel project to edit:
kernel-builder init 会为你提供一个可编辑的工作内核项目:
kernel-builder init --name myorg/vector-add-helion --backends cuda rocm xpu -- vector-add-helion
cd vector-add-helionkernel-builder init --name myorg/vector-add-helion --backends cuda rocm xpu -- vector-add-helion
cd vector-add-helionNote the -- before the directory name. --backends takes any number of values, so without the separator the directory name is read as another backend.
注意目录名称前的 --。--backends 可以接受任意数量的值,因此如果没有分隔符,目录名会被读取为另一个后端。
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}rm -rf vector_add_helion_cuda vector_add_helion_xpu torch-ext/torch_binding.{cpp,h}That leaves three files to edit.
这样就剩下三个文件需要编辑。
build.toml
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][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.
- python-depends = ["helion"] 记录了内核在运行时依赖 Helion。当有人加载该内核时,kernels 会检查 Helion 是否可导入,如果不可导入则给出明确的错误提示。
- [torch-noarch] 表示没有提前编译(ahead-of-time compilation)。
torch-ext/vector_add_helion/__init__.py
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"]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.
此示例使用硬编码的配置,并通过 config= 固定。在后面的章节中,我们将详细介绍预调优以及如何分发覆盖大量形状的配置决策树。
flake.nix
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 = ./.; };
}{
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 .kernel-builder check-config .
kernel-builder build-and-copy .Build and publish the builds to the Hub:
将构建结果发布到 Hub:
kernel-builder build-and-upload .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:
构建会为每个后端生成一个目录,其中包含你的 __init__.py 以及一个生成的 metadata.json,后者携带了 Helion 的依赖信息:
{
"name": "vector-add-helion",
"python-depends": ["helion"],
"backend": { "type": "cuda" }
}{
"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.
以下是 Hub 上已发布的内核示例:sayakpaul/vector-add-helion。
更进一步:量化金融体系
看懂新闻只是起点——沿量化金融路径,把它变成能交付的工程能力