GitHub - pytorch/helion: A Python-embedded DSL that makes it easy to write fast, scalable ML kernels with minimal boilerplate.

2026-06-21T04:19:24.444879

标题：GitHub - pytorch/helion: 一个嵌入 Python 的领域特定语言（DSL），以极少的样板代码轻松编写快速、可扩展的机器学习内核。

原文：
- 2026 年 6 月 15 日：Helion 教程，使用 Helion 简化编写性能可移植内核 @ PLDI 2026, Boulder, CO

📚 查看文档 📚 | 🎥 观看演讲 🎥 | 🚀 在 Colab 中尝试 🚀 | 在 AMD DevCloud 中尝试

Helion 是一个嵌入 Python 的领域特定语言（DSL），用于编写机器学习内核，设计为编译到 Triton（一个用于 GPU 和其他设备编程的高性能后端）。与 Triton 相比，Helion 旨在提高抽象层次，使编写正确高效的内核更容易，同时为自动调优过程提供更多自动化。

名称 Helion 指的是氦-3 原子的原子核，而 Triton 指的是氢-3。

Helion 既可以视为带分块的 PyTorch，也可以视为更高级的 Triton。与 Triton 相比，Helion 通过自动调优减少了手动编码工作量。Helion 花费更多时间（约 10 分钟）进行自动调优，因为它会评估从单个 Helion 内核生成的数百种潜在 Triton 实现。更大的搜索空间也使内核在不同硬件之间具有更好的性能可移植性。Helion 自动化和自动调优的范围包括：

张量索引：
- 自动计算步长和索引。
- 自动调优不同索引方法（指针、块指针、TensorDescriptors）之间的选择。
- 支持逐操作索引策略，以实现对加载和存储的细粒度内存访问控制。
掩码：
- Helion 中大多数掩码是隐式的，不需要时会被优化掉。
网格大小和 PID 计算：
- 自动确定网格大小。
- 自动调优从程序 ID（PID）到数据块的多种映射。
隐式搜索空间定义：
- 消除手动定义搜索配置的需求。
- 自动生成配置标志和探索空间。
内核参数管理：
- 自动处理内核参数，包括张量大小和步长。
- 将全局变量和（嵌套）闭包提升为内核参数，从而实现更好的模板化。
循环归约：
- 可以自动将大型归约转换为循环实现。
自动优化：
- PID 交错以提高 L2 缓存重用。
- 循环重排序。
- 持久内核策略。
- Warp 专业化选择、循环展开等。

一个在 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

for 循环外部的代码是标准的 PyTorch 代码，在 CPU 上执行。通常用于分配输出张量和执行形状计算等任务。

for 循环内部的代码被编译成 Triton 内核，最终生成单个 GPU 内核。单个 Helion 内核始终被编译为恰好一个 GPU 内核。

hl.tile 函数将迭代空间（此处为 m 乘以 n）细分为数据块。这些数据块在 GPU 上并行执行。分块细节（如维度（1D vs 2D）、块大小和循环顺序）由 Helion 的自动调优器自动确定。或者，也可以通过 helion.kernel 中的 config= 参数显式指定这些细节。

外层 for 循环映射到生成内核的网格。网格大小根据所选块大小自动确定。
内层 for 循环转换为生成内核内的循环，其块大小也自动确定。

在 Helion 内核中，标准 PyTorch 运算符（如 torch.addmm）使用 TorchInductor 自动映射为 Triton 操作。因此，熟悉 PyTorch 意味着您已经掌握了 Helion 的大部分内容。Helion 支持多种操作，包括逐元素操作（add、sigmoid 等）、归约操作（sum、softmax 等）、视图操作和矩阵乘法操作。Helion 内核中支持任意函数调用，但必须能够使用 make_fx 进行追踪。

上面的示例可以通过以下方式执行：

out = matmul(torch.randn([2048, 2048], device="cuda"),
             torch.randn([2048, 2048], device="cuda"))

当内核首次运行时，Helion 启动自动调优。典型的自动调优会话输出类似于：

[0s] Starting DifferentialEvolutionSearch with population=40, generations=20, crossover_rate=0.8
[20s] Initial population: failed=4 min=0.0266 mid=0.1577 max=1.2390 best=Config(block_sizes=[64, 32, 64], loop_orders=[[1, 0]], l2_groupings=[8], range_unroll_factors=[3, 1], range_warp_specializes=[True, False], range_num_stages=[1, 0], range_multi_buffers=[True, True], range_flattens=[None, False], num_warps=4, num_stages=7, indexing='block_ptr', pid_type='persistent_blocked')
[51s] Generation 2: replaced=17 min=0.0266 mid=0.0573 max=0.1331 best=Config(block_sizes=[64, 32, 64], loop_orders=[[1, 0]], l2_groupings=[8], range_unroll_factors=[3, 1], range_warp_specializes=[True, False], range_num_stages=[1, 0], range_multi_buffers=[True, True], range_flattens=[None, False], num_warps=4, num_stages=7, indexing='block_ptr', pid_type='persistent_blocked')
[88s] Generation 3: replaced=18 min=0.0225 mid=0.0389 max=0.1085 best=Config(block_sizes=[64, 64, 16], loop_orders=[[0, 1]], l2_groupings=[4], range_unroll_factors=[0, 1], range_warp_specializes=[None, None], range_num_stages=[0, 0], range_multi_buffers=[None, False], range_flattens=[None, None], num_warps=4, num_stages=6, indexing='pointer', pid_type='flat')
...
[586s] Generation 19: replaced=3 min=0.0184 mid=0.0225 max=0.0287 best=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')
[586s] Autotuning complete in 586.6s after searching 1520 configs.

可以通过以下方式硬编码最佳配置并跳过自动调优：

@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:
    ...

由于自动调优可能耗时较长（上述示例中约 10 分钟），您可能希望手动指定从自动调优中找到的最佳配置，以避免重复调优：

```python
@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:
...

查看原文