Hugging Face 的 @huggingface Kernels 项目现已支持 Helion。本文将详细介绍如何通过 Hugging Face Kernels 项目构建、自动调优并发布高性能且可移植的 Helion 内核，使用户能够无缝使用这些内核。了解更多：https://t.co/51KJSXshK5

构建与发布开箱即用的高性能内核 – PyTorch

来源：https://pytorch.org/blog/helion-x-%F0%9F%A4%97-hf-kernels-building-and-shipping-out-of-the-box-performant-kernels/

项目推荐

• Helion 标志 (https://pytorch.org/projects/helion/)
• PyTorch 标志 (https://pytorch.org/projects/pytorch/)

摘要

HuggingFace Kernels 项目现已支持 Helion。本文将介绍如何通过 Hugging Face Kernels 项目构建、自动调优并发布高性能且可移植的 Helion 内核，使用户能够无缝使用这些内核。

引言

Helion (https://github.com/pytorch/helion) 是一种用于编写高性能、可移植的机器学习内核的高级 DSL。🤗 Kernels (https://huggingface.co/docs/kernels/en/index) 项目允许内核开发者以一致且可重现的方式在 Hugging Face Hub 平台上打包和分发他们的内核。同时，它也让内核用户无需管理复杂的依赖关系即可无缝使用这些内核。

在本文中，我们将讨论 Helion 如何在 Kernels 项目中获得支持、用户如何受益于 Helion 的一流自动调优支持，以及如何发布预调优的内核配置以减少冷启动时间。我们还将展示 Helion 内核的示例，以及如何针对特定问题规模进行调优以获得性能提升。

附注：在后文中，我们将用首字母大写的“Kernels”来指代该项目，以区分实际的“内核”（kernels）。

简介：Helion

Helion 是一种分块 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 out

Helion 的吸引力不仅在于其简洁的语法，更在于其刻意不指定的部分。当你编写 hl.tile 时，你只是说明迭代空间应该被分块——而块的大小如何，或如何从内存中获取数据，这些都未指定。Helion 将这些决策转化为一个可供自动调优的搜索空间。

关键是，自动调优器不仅扫描像分块大小这样的数值参数，它还会搜索降级策略——即内核的实际实现方式：使用哪种内存访问模式（指针算术、块指针、TMA），如何排序和扁平化嵌套循环，归约操作应该是持久化的还是循环的，等等。

在 Triton 或 CUDA 中，切换这些选择意味着要完全重写内核；而在 Helion 中，最优选择是通过算法找到的。正是这个自动调优过程，使得 Helion 内核在针对大量形状进行基准测试时，通常能够超越用更底层语言手工编写的内核。

不过，自动调优有时是一个耗时的过程，因此有一个既定的方法来发布捆绑了预调优配置的内核是很有益的。这正是 Kernels 项目发挥作用的地方。

简介：Kernels

当前内核打包和分发的格局是分散的，表现为不一致的源代码结构、不同的工具以及有限的兼容性支持。因此，用户常常面临冗长的构建时间，即使有预构建的 wheel 包也是如此。

Kernels 项目通过为 AOT（提前编译）和 JIT（即时编译）内核建立标准化的、统一的打包和构建流程来解决这些挑战。该项目分为两个主要组件：

• kernel-builder：一个供开发者可靠地跨不同框架版本和系统配置打包和分发内核的工具。它强制执行标准，以确保可预测的源代码结构、构建可重现性、原生 PyTorch 兼容性以及便于社区共享。
• kernels：一个面向用户的 Python 库，允许用户通过简单的命令（如 get_kernel("org/name", version=1)）轻松加载现成的内核，无需管理依赖问题，就像从 Hugging Face Hub 拉取模型或数据集一样。

对于内核用户，我们希望提供无缝的体验，能够立即加载内核并准备好使用。让我们看一个如何加载流行的 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(...)

我们为广泛的兼容性矩阵提供了预构建的 AOT 内核二进制文件，例如 Flash Attention 3。这对最终用户非常有益，尤其是在内核的上游仓库可能没有特定构建可用时。用户可以在 Hugging Face Hub 平台 (hf.co/kernels) 上浏览各种各样的内核。我们将此内核集合称为 Kernels Hub。

在 Kernels 中打包和使用 Helion

Helion 内核是纯 Python 的。它们在第一次调用时自行编译，因此 kernel-builder 无需提前编译任何内容。Helion 提供了用于针对特定工作负载和硬件调优这些内核的实用程序（稍后详述），以便用户可以一次调优，之后重复使用。Helion 内核也是无架构（noarch）内核：你发布源代码，Helion 在用户的机器上完成其余工作。

在本节中，我们将讨论如何为使用 kernel-builder 构建而脚手架化和组织一个 Helion 内核。

从脚手架开始

kernel-builder init 会给你一个可编辑的内核项目：

kernel-builder init --name myorg/vector-add-helion --backends cuda rocm xpu -- vector-add-helion
cd vector-add-helion

注意目录名之前的 --。--backends 可以接受任意数量的值，因此没有分隔符时目录名会被读作另一个后端。

脚手架假设是一个编译好的内核，所以删除你不需要的部分：

rm -rf vector_add_helion_cuda vector_add_helion_xpu torch-ext/torch_binding.{cpp,h}

这剩下三个需要编辑的文件。

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]

有两点值得注意：

• python-depends = ["helion"] 记录了内核在运行时需要 Helion。当有人加载该内核时，kernels 会检查 Helion 是否可导入，如果不可导入则会给出明确的错误信息。
• [torch-noarch] 表示没有提前编译。

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"]

这里使用了一个通过 config= 固定下来的硬编码配置。在后面的章节中，我们将详细介绍预调优以及如何发布一个覆盖大量形状的配置决策树。

flake.nix

脚手架生成的无需更改：

{ inputs.kernel-builder.url = "github:huggingface/kernels";
  outputs = { self, kernel-builder, ... }:
    kernel-builder.lib.genKernelFlakeOutputs { inherit self; path = ./.; };
}

构建与发布

kernel-builder check-config .
kernel-builder build-and-copy .

构建并将构建产物发布到 Hub：

kernel-builder build-and-upload .

构建会为每个后端生成一个目录，每个目录包含你的 __init__.py 以及一个生成的 metadata.json，该文件携带了 Helion 依赖项：

{
  "name": "vector-add-helion",
  "python-depends": ["helion"],
  "backend": {
    "type": "cuda"
  }
}

以下是在 Hub 上发布的内核示例：sayakpaul/vector-add-helion (https://huggingface.co/kernels/sayakpaul/vector-add-helion)。

使用内核

from kernels import get_kernel
kernel = get_kernel("myorg/vector-add-helion", version=1, trust_remote_code=True)
out = kernel.vector_add(x, y)

用户需要安装 Helion (pip install helion) 以及一个能识别它的 kernels 版本。旧版本会直接拒绝该依赖并报错 unsupported kernel dependency: helion。

预调优与发布预调优配置

为了确保发布的 Helion 内核在不同的输入问题形状和不同的 GPU 代际上都能表现出色，通常预调优是很有益的。工作流程包含三个步骤：

1. 在一组具有代表性的形状上对内核进行调优。
2. 让 Helion 构建一个决策树，将每个形状映射到其最佳配置。
3. 将该决策树与你的内核一起发布。

在加载时，Helion 读取该树（存储为与内核源文件相邻的普通 python 文件），并为每次调用选择一个配置——无需在用户的机器上进行调优。

为此，第一个内核源代码变更是添加一个装饰器：

--- @helion.kernel(config=...)
+++ @helion.aot_kernel(static_shapes=True)

编写一个小脚本，对你想要在预调优期间覆盖的每个形状调用该内核：

# bench.py
import torch
from vector_add_helion.vector_add import vector_add

for n in [1024, 1 << 16, 1 << 20, 1 << 24]:
  x = torch.randn(n, device="cuda")
  y = torch.randn(n, device="cuda")
  vector_add(x, y)

然后将该脚本交给 Helion 的 AOT 运行器：

python -m helion.autotuner.aot_runner \
  --phase all --goal max_slowdown --threshold 1.01 --max-configs 8 \
  -- python bench.py

该运行器驱动 bench.py，完成三个阶段：

• 收集：独立地对每个形状进行自动调优。
• 测量：在每个形状上重新对每个发现的配置进行基准测试。
• 构建：选择最小的配置集合，使得每个形状的性能保持在其自身最佳性能的 –threshold 范围内（1.01 即 1% 以内），最多不超过 –max-configs 个配置。

如果一个配置能满足所有形状，那么你就只发布这一个；如果形状差异较大，你就会得到一个包含多个配置的树。

运行器会在你的内核源文件旁边生成一个纯 Python 文件，命名为 _helion_aot___.py，该文件可以与内核源文件以及其他文件一起在 Hub 上的构建中发布。因此，最终的文件结构可能如下所示：

vector-add-helion/
├── build.toml
├── flake.nix
└── torch-ext/vector_add_helion/
    ├── __init__.py
    ├── vector_add.py               # @helion.aot_kernel
    ├── _helion_aot_vector_add_cuda_sm90.py   # H100 配置
    └── _helion_aot_vector_add_cuda_sm100.py  # B200 配置

当使用者通过 get_kernel 访问此内核并在张量上调用它时，Helion 将使用决策树来识别一个适合运行时输入形状的预调优配置。

示例

Attention

在 HelionDSL/attention (https://huggingface.co/kernels/HelionDSL/attention) 中，我们展示了一个通过 Kernels 发布的预调优 Helion attention 内核示例。它包含为 NVIDIA H100 GPU 预调优的配置，使用上述工作流程和结构创建。

我们测量了这些预调优配置的性能，并在 H100 上与 PyTorch 的 scaled_dot_product_attention 实现的 FLASH 后端进行了比较。此发布的内核在 19 个预调优形状中的 19 个上均优于 SDPA，几何平均加速比为 1.20。在调优期间未见过的保持集形状上，该内核在 10 个形状中的 9 个上优于 SDPA，几何平均加速比为 1.17。

Linear Attention

在 HelionDSL/linear-attention (https://huggingface.co/kernels/HelionDSL/linear-attention) 中，我们发布了七个预调优的线性注意力内核：linear attention、simple GLA、retention、GLA、delta rule、gated delta rule 和 KDA。它包含为 NVIDIA B200 GPU 预调优的配置。

我们测量了这些预调优配置与 FLA 的性能。在六个预调优形状上，发布的内核在所有七个变体上都比 flash-linear-attention 更快，在设备时间（CUDA 图重放）上的几何平均加速比为 1.41，端到端（包括 CPU 调度）为 1.33。在调优期间未见过的六个保持集形状上，设备时间几何平均加速比为 1.35，端到端为 1.31。前向和后向传播合并在预调优形状上的几何平均加速比为 1.55。

结论

在本文中，我们讨论了来自 Meta 的高级 DSL Helion 和来自 Hugging Face 的 Kernels 项目。我们还展示了这两个项目如何相辅相成，以简化计算优化内核的开发、分发和使用过程。我们欢迎您尝试 Helion 并在 Hub 上发布您酷炫的 Helion 内核 🤗

致谢： 感谢 Daniël de Kok 和 Lysandre Debut 对本文的审阅。