Helion 简介
在现代机器学习领域，对高性能计算的需求催生了大量定制化内核。尽管这些内核能够提供卓越的性能，但它们往往是用低级且特定于硬件的语言编写的。这带来了沉重的长期维护负担：一个为某种硬件架构精心优化的内核很快就会变成技术债务，且难以以低成本移植到其他架构。这一挑战阻碍了开发与创新，迫使开发者必须在生产力与性能之间做出权衡。
Helion 通过将一种嵌入 Python 的高级领域专用语言 (DSL) 编译为自动调优的 Triton 代码，解决了这一矛盾。它建立了一个新的抽象层，架起了 PyTorch 用户友好型简单性与底层语言高性能之间的桥梁。通过自动化处理张量索引、内存管理和特定硬件调优等枯燥且易错的任务，Helion 让开发者能够专注于算法逻辑，而非硬件相关的实现细节。Helion 通过将熟悉的以 PyTorch 为中心的语法与强大的自动调优引擎相结合，实现了这种平衡，该引擎可自动完成复杂内核配置的搜索过程。最终，该系统在保证跨硬件架构性能可移植性的同时，大幅降低了开发工作量。
开发新 DSL 的动机
为内核开发选择合适的抽象级别是一项战略决策，它直接影响性能、可维护性和开发速度。当前的编程语言和抽象方式迫使开发者在底层控制和高级生产力之间做出二选一的权衡，而这两个极端各有优劣。
- CUDA/Gluon/TLX： 使用 CUDA 等语言直接编写内核虽然提供了最大限度的控制权，但要达到高性能通常需要投入巨大的精力。这些内核高度专业化，往往针对特定硬件，难以适配新架构。
- Triton： 虽然 Triton 代表了重大的进步，但仍需要大量的人工介入。开发者需要负责显式管理张量索引、定义自动调优的搜索空间、管理内核参数，且更改优化策略往往需要大幅重写代码。
- PyTorch： 虽然 PyTorch 和
torch.compile等框架提供了极佳的易用性，但它们提供的细粒度控制有限。需要指定精确融合策略的用户往往会发现这种高级抽象过于受限。
Helion 编程模型：“带有平铺 (Tiles) 的 PyTorch”
Helion 编程模型的目标是最大限度减少样板代码，并利用开发者现有的 PyTorch 知识。这种设计理念通过提供一种被称为“带有平铺 (Tiles) 的 PyTorch”的熟悉且直观的语法，加速了正确、高效内核的创建。
一个典型的 Helion 内核（例如下方的矩阵乘法示例）由两个协同工作的不同部分组成：
import torch, helion, helion.language as hl @helion.kernel() def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # --- Host Code (runs on CPU) --- m, k = x.size() k, n = y.size() out = torch.empty([m, n], dtype=x.dtype, device=x.device) # --- Device Code (compiles to a Triton kernel) --- 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
- 宿主代码 (Host Code)： 最外层
hl.tilefor 循环之外的代码为标准 PyTorch 代码。它主要用于分配输出张量和计算形状等设置任务。Helion 会自动处理这些值向设备代码的传递，从而无需人工管理参数。 - 设备代码 (Device Code)： 最外层
hl.tilefor 循环之内的代码是内核的核心。该部分会被编译成单个高性能 Triton 内核，并在 GPU 上并行执行。
核心语言结构 hl.tile 将内核的迭代空间细分为多个平铺块 (Tiles)。程序员只需指定如何平铺迭代空间，而具体的实现细节（如平铺大小、迭代顺序和内存布局优化）均由 Helion 的自动调优器处理，它会系统地探索目标硬件的最佳配置。
在内核体内，开发者可以使用标准的 PyTorch 算子，如 torch.addmm 及其他逐点 (pointwise) 或归约 (reduction) 操作。Helion 利用 PyTorch 2 的核心组件 TorchInductor，自动将这些 PyTorch 调用映射到对应的底层 Triton 实现。这带来了一个强大的可用性优势：熟悉 PyTorch 意味着你已经掌握了大部分 Helion。
Helion 还包含模板功能，允许将可能在闭包中捕获额外参数的 lambda 函数作为参数传递给内核。如该示例所示，这对于实现带有可自定义尾部的通用内核特别有用。例如，lambda 函数可以捕获周围作用域中定义的张量。Helion 的编译器会自动检测该变量并将其作为生成的 Triton 内核中的一个参数。这消除了大量用于在多层函数调用中传递新输入的样板代码，实现了高可重用性和通用内核的创建。
Helion 使内核实现变得极其简单：Attention 内核在 Helion 中仅需 30 行代码，而在 Triton 中则需要 120 行，在 CUDA 中更是需要数千行。这种高级声明式编程模型之所以能保持高性能，得益于支撑整个系统的核心机制：自动调优引擎。
Helion 的自动调优器：通过隐式搜索空间生成最优内核
Helion 的核心差异化优势在于其自动化的提前 (AOT) 自动调优引擎。在 Triton 中，开发者必须手动定义优化搜索空间，这需要显式枚举每一个待测试的配置，是一个既枯燥又限制探索范围的过程。
Helion 通过隐式搜索空间改变了这种动态。高级语言会自动构建一个关于实现选择的庞大、多维搜索空间。例如，一个 hl.tile 调用会隐式指示自动调优器去探索不同的块大小、循环顺序以及是否将迭代空间展平为单个维度。因此，一个 Helion 内核定义可以映射到数千个 Triton 配置，使自动调优器能够创建更大、更丰富的搜索空间，从而发现更优的配置。
自动调优工作流
当内核在没有指定配置的情况下首次运行时，自动调优器会启动自动搜索。该过程通常耗时约 10 分钟，利用差分进化 (Differential Evolution) 或模式搜索 (Pattern Search) 等搜索策略，针对给定的输入形状和硬件评估数千个候选 Triton 内核配置，以识别最优参数集。完成后，自动调优器会输出发现的最佳配置。
[586s] Autotuning complete in 586.6s after searching 1520 configs. One can hardcode the best config and skip autotuning with: @helion.kernel(config=helion.Config(block_sizes=[64, 64, 64], loop_orders=[[0, 1]], l2_groupings=[4], range_unroll_factors=[0, 1], range_warp_specializes=[None, False], range_num_stages=[0, 3], range_multi_buffers=[None, False], range_flattens=[None, None], num_warps=8, num_stages=6, indexing='block_ptr', pid_type='flat'))
开发者可以将此配置复制到源代码中的 @helion.kernel() 装饰器中。这指示 Helion 在后续运行中完全跳过搜索过程。在生产环境中，这能实现快速、确定性的编译，生成单个预优化的 Triton 内核，以极小的工作量获得等同于手工精心调优的内核性能。
@helion.kernel(config=helion.Config( block_sizes=[64, 64, 64], loop_orders=[[0, 1]], l2_groupings=[4], range_unroll_factors=[0, 1], range_warp_specializes=[None, False], range_num_stages=[0, 3], range_multi_buffers=[None, False], range_flattens=[None, None], num_warps=8, num_stages=6, indexing='block_ptr', pid_type='flat' )) def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: ...
开发者也可以在 @helion.kernel() 中指定一个配置列表，此时 Helion 将仅探索这些配置以选择最快的实现。
配置空间
配置空间代表了 Helion 自动化的实现选择集合。该空间是 Helion 性能可移植性的主要来源，因为它允许单个内核定义适应不同硬件架构和输入张量大小的独特性质。探索这一空间正是 Helion 相较于手动编写内核（通常针对特定条件调优）的优势所在。
自动调优器会探索从数据移动到线程映射等广泛的参数。下表详细说明了配置选项。
| 参数 | 描述 | |
|
在 Triton 中，开发者可以在三种不同的内存访问方法中做出选择：指针算术、块指针 (block pointers) 和张量描述符 (tensor descriptors，利用 NVIDIA Hopper/Blackwell GPU 上的 Tensor Memory Accelerators, TMA)。最佳选择取决于硬件架构和内存访问模式，但切换这些方法通常需要大幅重写代码。
Helion 通过 |
|
|
block_sizes 参数是一个列表，对应 hl.tile 循环中每个维度的平铺大小，它决定了每个线程块处理的数据量。这会影响寄存器使用率、共享内存需求和并行度。 |
|
|
flatten_loops 选项控制是否将 hl.tile 循环的多维平铺空间展平为单个维度，从而在无需开发者重写内核代码的情况下扩大了自动调优器的搜索空间。 |
|
|
为了优化数据局部性，Helion 提供了两个配置旋钮。loop_orders 参数允许自动调优器对嵌套平铺的迭代顺序进行置换，这会根据张量布局影响缓存命中率。l2_grouping 配置启用 PID swizzling（一种重新排序线程块分配以改善 L2 缓存中数据重用的技术）。在 Triton 中，这些转换将需要重写复杂的循环结构和索引计算。 |
|
|
在执行归约（例如对张量维度进行 sum()）时，持久化归约 (persistent reduction) 会处理单个平铺的整个归约维度，这对小维度来说很快。然而，如果归约规模很大，这种方法会造成高寄存器压力，导致寄存器溢出和低性能。或者，开发者可以编写一个循环，以更小的分块来迭代归约维度。
虽然在这两种策略之间切换通常需要重写代码，但 Helion 通过 |
|
|
Helion 自动计算网格大小并将程序 ID (PID) 映射到数据平铺块。pid_type 配置允许自动调优器在不进行任何手动代码更改的情况下探索各种映射策略。
|
|
|
Helion 会对 Triton 的 tl.load 的 eviction_policy 参数进行自动调优，从而影响 GPU L1 缓存驻留情况。自动调优可以挑选出最适合内核内存访问模式的驱逐提示组合。 |
|
Triton 配置
|
Helion 会自动探索标准的 Triton 可调参数，减轻了开发者进行手动调优的负担。 | |
性能分析与基准测试
我们对比了 Helion 与 torch.compile（开启 max-autotune）以及手写 Triton 的性能，衡量了它们在各种内核和形状下相对于 eager 模式执行的加速比，测试平台为 NVIDIA B200 和 AMD MI350X GPU。大多数手写 Triton 内核来自 Liger-Kernel 基准测试套件。
NVIDIA B200 上的性能
下表总结了在 NVIDIA B200 GPU 上的性能，绿色突出显示的单元格表示每个内核相对于 eager 模式执行的最高加速比。在所有基准测试中，Helion 实现了 3.27 倍的最高几何平均加速比，其次是 torch.compile（开启 max-autotune）的 2.7 倍，以及手写 Triton 内核的 1.76 倍。平均而言，Helion 比 torch.compile 快 1.21 倍，比 Triton 内核快 1.85 倍。性能提升在 softmax kernel 中尤为显著（Helion 比 torch.compile 快 2.28 倍），在 jsd 内核中，Helion 的表现比手写 Triton 高出 6.22 倍。
AMD MI350X 上的性能
AMD MI350X 上的性能表现出类似的趋势，Helion 实现了相对于 eager 模式执行 2.37 倍的最高几何平均加速比，而 torch.compile 为 2.26 倍，Triton 内核为 1.65 倍。平均而言，Helion 比 torch.compile 快 1.05 倍，比 Triton 内核快 1.44 倍。在 int4_gemm 和 jsd 内核中，性能提升尤为明显，Helion 分别比手写 Triton 内核高出 4.5 倍和 4.4 倍。
案例研究 1：超越高度优化的 CuTe DSL 内核
Helion 实现的 RMSNorm 反向内核 在不到一天的时间内编写完成，其性能证明与用 CuTe DSL 编写的高度手动优化 Quack 内核 相当或更优。在 H100 GPU 上的各种归约维度测试中，Helion 一致匹配或优于手动调优的内核，证明了其在更高抽象级别开发所带来的生产力提升下，仍能匹配专家级性能的能力。我们也与 torch.compile（开启 max-autotune）及手写 Triton 进行了对比，Helion 在许多情况下均表现更佳。
这些结果得益于 Helion 底层的编译器架构，该架构旨在高效支持自动调优的大规模搜索。
案例研究 2：Helion 与 TileLang 的基准测试
我们还将 Helion 与 TileLang（一个用于开发高性能 GPU 内核的开源 DSL）进行了性能对比。我们在 Helion 中实现了 Mamba-2-chunk-scan 内核（Mamba-2 架构的核心选择性扫描操作），并与 TileLang 和 Triton 的实现进行了对比。在 H100 上，Helion 展现了最高性能，在不同配置下分别比 TileLang 快 2.12–2.63 倍，比 Triton 快 1.2–1.85 倍。
高级编译器架构
Helion 的编译器架构旨在通过 TorchInductor 将 Python 函数逐步降级为高度优化的 Triton 代码。
编译流水线按以下关键阶段进行：
- Python AST 解析： 该过程首先将内核的 Python 源代码解析为抽象语法树 (AST)。
- 类型传播与元数据： 自定义通道遍历 AST，用类型信息和其他必要元数据标注每个节点，以创建一个扩展的 AST。
- 降级至设备 IR (Device IR)： 此标注后的树被降级为 Helion 的主中间表示 (IR)。设备 IR 是一组静态单赋值 (SSA) 形式的 FX 图，每个图代表程序的一个基本块。图中的每个节点都包含一个指向 Inductor IR 节点的指针，用于代码生成。
- 编译器通道 (Compiler Passes)： 对设备 IR 应用一系列转换通道。这些通道实现了关键语义更改，例如将持久化归约转换为循环归约的“归约回滚”优化。
- 配置代码生成： 在最后阶段，代码生成器接收两个输入：转换后的设备 IR 和已自动调优的配置。它利用这些输入生成最终的输出 Triton 代码。
一个关键的架构决策是：性能关键型配置仅在流水线的最后阶段即代码生成时才应用。这使得从解析到 IR 转换的大部分编译过程可以在自动调优搜索之前只运行一次，从而使数千个配置的探索变得计算高效。
结论
Helion 填补了当今机器学习内核编写领域的一个关键空白。通过结合熟悉的、类似 PyTorch 的高级语法与强大的提前自动调优引擎，它提供了开发人员生产力、细粒度控制和性能可移植性之间的独特平衡。它赋能开发者编写可移植、面向未来的内核，在无需深厚硬件专业知识的情况下实现最前沿的性能，为高性能机器学习内核建立了一种更新、更具生产力的范式。
Helion 于 2025 年 10 月 22 日发布测试版，欢迎社区提供反馈、错误报告和贡献。更多信息请查看：
致谢
Helion 是众人协作的成果，贡献者包括：Jason Ansel, Oguz Ulgen, Will Feng, Jongsok Choi, Markus Hoehnerbach, Manman Ren, Jie Liu, Paul Zhang, Driss Guessous, Joy Dong, Xuan Zhang, Karthick Panner Selvam, Peng Wu, Hongtao Yu, Neil Dhar, Nick Riasanovsky, Shane Nay, Alexey Loginov，以及 Meta、NVIDIA、AMD 和 Intel 的团队。
注：B200 性能数据已于 10/23 更新。