# 活动
- **2026年6月15日**:Helion 教程,[使用 Helion 简化性能可移植 kernel 的编写](https://pldi26.sigplan.org/details/pldi-2026-tutorials/1/Writing-Performance-Portable-Kernels-Simplified-with-Helion) @ PLDI 2026,科罗拉多州博尔德
# 关于
📚 **[查看文档](https://helionlang.com)** 📚 |
🎥 **[观看演讲](https://youtu.be/BW-Ht-5IxgM)** 🎥 |
🚀 **[在 Colab 中尝试](https://colab.research.google.com/github/pytorch/helion/blob/main/notebooks/softmax.ipynb)** 🚀 |
**[在 AMD DevCloud 中尝试](https://amd-ai-academy.com/github/pytorch/helion/blob/main/notebooks/softmax.ipynb)**
**Helion** 是一种嵌入在 Python 中的领域特定语言 (DSL),用于编写机器学习 kernel,旨在将其编译为 [Triton],后者是用于 GPU 和其他设备编程的高性能后端。与 Triton 相比,Helion 旨在提高抽象级别,使得编写正确且高效的 kernel 变得更加容易,同时在自动调优过程中实现更多自动化。
*Helion* 这个名字指的是氦-3 原子的原子核,而 *Triton* 指的是氢-3。
Helion 既可以被视为 *带有 tile 的 PyTorch*,也可以被视为 *更高层级的 Triton*。与
Triton 相比,Helion 通过自动调优减少了手动编码的工作量。Helion 会花费更多时间(约 10 分钟)进行自动调优,因为它会评估从单个 Helion kernel 生成的数百种潜在的 Triton 实现。这种更大的搜索空间也使得 kernel 在不同硬件之间具有更好的性能可移植性。Helion 自动化并自动调优以下内容:
1. **Tensor 索引:**
* 自动计算 stride 和索引。
* 自动调优各种索引方法(pointer、block pointer、TensorDescriptor)之间的选择。
* 支持针对每个操作的索引策略,以实现对加载和存储的细粒度内存访问控制。
2. **掩码:**
* Helion 中的大多数掩码都是隐式的,并且在不需要时会被优化掉。
3. **Grid 大小和 PID 计算:**
* 自动确定 grid 大小。
* 自动调优从 Program ID (PID) 到数据 tile 的多种映射。
4. **隐式搜索空间定义:**
* 无需手动定义搜索配置。
* 自动生成配置标志和探索空间。
5. **Kernel 参数管理:**
* 自动处理 kernel 参数,包括 tensor 大小和 stride。
* 将全局变量和(嵌套的)闭包提升为 kernel 参数,从而实现更好的模板化。
6. **循环归约:**
* 可以自动将大型归约转换为循环实现。
7. **自动化优化:**
* PID swizzling,以提高 L2 缓存复用率。
* 循环重排。
* 持久化 kernel 策略。
* Warp specialization 选择、展开等。
## 示例
Helion 中一个最小化的矩阵乘法 kernel 如下所示:
```
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` 循环之外的代码是在 CPU 上执行的标准 PyTorch 代码。它通常用于分配输出 tensor 和执行形状计算等任务。
`for` 循环之内的代码会被编译为 Triton kernel,从而生成一个单一的 GPU kernel。一个单一的 Helion kernel 总是会被精确地编译为一个 GPU kernel。
`hl.tile` 函数将迭代空间(在本例中为 `m` 乘以 `n`)细分为多个 tile。这些 tile 在 GPU 上并行执行。Tiling 的细节,如维度(1D 或 2D)、tile 大小和循环顺序,由 Helion 的自动调优器自动决定。或者,也可以使用 `helion.kernel` 中的 `config=` 参数显式指定这些细节。
* 外层的 `for` 循环被映射到所生成 kernel 的 grid 上。grid 大小是根据所选的 tile 大小自动确定的。
* 内层的 `for` 循环转换为生成 kernel 内部的一个循环,其 tile 大小也是自动确定的。
在 Helion kernel 中,标准的 PyTorch 操作(如 `torch.addmm`)会使用 [TorchInductor](https://github.com/pytorch/pytorch/tree/main/torch/_inductor) 自动映射到 Triton 操作。
因此,熟悉 PyTorch 就意味着你已经掌握了 Helion 的大部分内容。Helion 支持多种操作,包括逐点操作(`add`、`sigmoid` 等)、归约操作(`sum`、`softmax` 等)、视图操作以及矩阵乘法操作。Helion kernel 支持任意函数调用,但必须可以使用 [make_fx](https://pytorch.org/docs/stable/generated/torch.fx.experimental.proxy_tensor.make_fx.html) 进行追踪。
## 自动调优
上述示例可以通过以下方式执行:
```
out = matmul(torch.randn([2048, 2048], device="cuda"),
torch.randn([2048, 2048], device="cuda"))
```
当 kernel 第一次运行时,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.
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'))
```
由于自动调优可能非常耗时(在上述示例中大约需要 10 分钟),你可能希望手动指定通过自动调优找到的最佳配置,以避免重复调优:
```
@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 执行更轻量级的自动调优过程:
```
@helion.kernel(configs=[
helion.Config(...),
helion.Config(...),
])
def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
...
```
在这种情况下,Helion 会评估所提供的配置并选择最快的一个。
此外,Helion 还提供了编程式 API,可以直接从你的代码中管理自动调优和配置。
**对于生产环境部署**,我们建议使用提前调优好的配置,而不是依赖运行时自动调优。自动调优过程可能非常耗时且消耗大量资源,这使其不适用于对性能和启动时间有严格预测要求的生产环境。
### 静态形状与自动调优键
默认情况下,Helion 使用静态形状 (`static_shapes=True`)。这意味着每个唯一的输入形状/stride 签名都会被视为其自身的特化,并进行单独调优。这通常能带来最佳性能,但当遇到多种形状时,可能会增加调优时间。
如果你想通过在不同形状之间共享配置来减少调优时间,请设置 `static_shapes=False`。在此模式下,自动调优键会忽略确切的大小,从而允许在多种形状之间重用单一的调优配置。与完全特化的静态形状相比,这可能会带来性能上的损耗。
```
@helion.kernel(static_shapes=False)
def my_kernel(x: torch.Tensor) -> torch.Tensor:
...
```
## 配置
Helion 配置包含以下选项:
* **block\_sizes** (`list[int]`):
控制对应于 kernel 中传递给 `hl.tile` 或调用 `hl.register_block_size` 的每个维度的 tile 大小。
* **loop\_orders** (`list[list[int]]`):
包含对应于每个具有两个或更多维度的 `hl.tile` 调用的一项,允许你排列 tile 的迭代顺序。
* **flatten\_loops** (`list[bool]`):
包含对应于每个具有两个或更多维度的 `hl.tile` 调用的一项,允许你将迭代空间展平为单个维度。
* **range\_unroll\_factors** (`list[int]`):
包含对应于每个循环维度的一项,指定 `tl.range()` 调用的展开因子。小于 1 的值会忽略 `loop_unroll_factor` 参数。
* **range\_num\_stages** (`list[int]`):
包含对应于每个循环维度的一项,指定 `tl.range()` 调用的阶段数。小于 1 的值会忽略 `num_stages` 参数。
* **range\_multi\_buffers** (`list[bool | None]`):
包含对应于每个循环维度的一项,控制 `tl.range()` 调用的 `disallow_acc_multi_buffer` 参数。`True` 允许多缓冲(设置 `disallow_acc_multi_buffer=False`),`False` 禁止多缓冲(设置 `disallow_acc_multi_buffer=True`),而 `None` 则忽略该参数。
* **range\_flattens** (`list[bool | None]`):
包含对应于每个循环维度的一项,控制 `tl.range()` 调用的 `flatten` 参数。`True` 设置为 `flatten=True`,`False` 设置为 `flatten=False`,而 `None` 则忽略该参数。
* **range\_warp\_specializes** (`list[bool | None]`):
包含对应于每个循环维度的一项,控制 `tl.range()` 调用的 `warp_specialize` 参数。`True` 设置为 `warp_specialize=True`,`False` 设置为 `warp_specialize=False`,而 `None` 则忽略该参数。
仅在启用了 `allow_warp_specialize` 设置且具有 Blackwell 或更新架构的 CUDA 设备上可用。
* **static\_ranges** (`list[bool]`):
包含对应于每个具有静态边界的循环维度的一项,控制是否使用 `tl.static_range()` 调用。`True` 生成 `tl.static_range()` 并忽略该循环的 range_* 配置。`False` 生成 `tl.range()`。
* **reduction\_loops** (`list[int | None]`):
包含对应于每个归约维度的一项(参见 `examples/softmax.py`)。使用 `None` 会触发持久化归约,即整个归约都在单个 tile 中处理。指定整数 block 大小会将归约转换为循环,这对于超过可用寄存器容量的较大归约非常有利。
* **l2\_groupings** (`list[int]`):
重排生成 kernel 的 Program ID (PID),以改善 L2 缓存行为。值为 `1` 会禁用此优化,而更高的值指定分组大小。
* **indexing** (`"pointer"`, `"tensor_descriptor"`, `"block_ptr"` 或这些项的列表):
指定加载和存储操作的内存索引策略。可以是:
- 单一策略(适用于所有加载和存储):`indexing="block_ptr"`
- 策略列表(按执行顺序,每个加载/存储一个):`indexing=["pointer", "pointer", "block_ptr"]`
- 空/忽略(默认所有操作使用 `"pointer"`)
- 使用列表时,请按顺序提供策略:`[load1, load2, ..., store1, store2, ...]`
`"tensor_descriptor"` 选项使用 Tensor Memory Accelerator (TMA),但需要 Hopper 或更新的 GPU 以及最新开发版本的 Triton。
* **pid\_type** (`"flat"`, `"xyz"`, `"persistent_blocked"` 或 `"persistent_interleaved"`):
指定 Program ID 映射策略。`"flat"` 仅使用 x 维度,`"xyz"` 利用多个 grid 维度,而持久化策略则启用持久化 kernel 以提高 SM 利用率。
* **num\_warps** (`int`):
设置 kernel 将使用的 warp 数量。
* **num\_stages** (`int`):
定义要传递给 Triton 的流水线阶段数。
* **load_eviction_policies** (`list[str]`):
控制在设备循环中发现的加载所使用的逐出策略。每一项对应一个加载点;允许的值为 `""`(无策略)、`"first"`(映射到 Triton `evict_first`)和 `"last"`(映射到 Triton `evict_last`)。
在 `hl.load` 上显式指定 `eviction_policy=...` 会覆盖此配置。
更改这些选项通常会导致生成的 Triton 代码大不相同,从而使自动调优器能够从单个 Helion kernel 中探索广泛的实现。
## TileIR 后端 (Blackwell GPU)
Helion 支持 NVIDIA Blackwell GPU(计算能力 10.x/12.x)的 [Triton-TileIR 后端](https://github.com/triton-lang/Triton-to-tile-IR)。该后端提供了针对 [TileIR](https://docs.nvidia.com/cuda/tile-ir/latest/index.html) 的优化代码生成,并带有额外的调优旋钮。
要启用 TileIR 后端:
1. 安装 [Triton-TileIR 后端](https://github.com/triton-lang/Triton-to-tile-IR)
2. 设置环境变量:
```
export ENABLE_TILE=1
```
有关详细文档,请参阅 [TileIR 后端指南](docs/tileir_backend.md)。
## CuTe 后端 (CUDA 13+)
Helion 拥有一个实验性的 [CUTLASS CuTe DSL](https://github.com/NVIDIA/cutlass) 后端,该后端通过 CuTe 而不是 Triton 来降低 kernel 的层级。
**要求:** PyTorch 必须基于 **CUDA 13 或更高版本** 构建 (`torch.version.cuda >= "13"`)。该后端将拒绝在较旧的 CUDA 构建版本上启动。
要启用 CuTe 后端:
1. 安装指定的 CuTe DSL 版本:
pip install -e '.[cute]'
# 或者运行辅助脚本,它也会处理 libs-cu13 的重新安装顺序
./scripts/install_cute.sh
2. 在运行时选择后端:
export HELION_BACKEND=cute
## 开发和调试设置
在使用 Helion 开发 kernel 时,你可能希望跳过自动调优以实现更快的迭代。为此,请设置环境变量 `HELION_AUTOTUNE_EFFORT=none` 或使用装饰器参数 `@helion.kernel(autotune_effort="none")`。**警告:** 默认配置速度很慢,不适用于生产环境或性能测试。
要查看生成的 Triton 代码,请设置环境变量 `HELION_PRINT_OUTPUT_CODE=1` 或在 `@helion.kernel` 装饰器中包含 `print_output_code=True`。这会将 Triton 代码打印到 `stderr`,这对于调试和理解 Helion 的编译过程非常有帮助。还可以使用 `foo_kernel.bind(args).to_triton_code(config)` 将 Triton 代码作为字符串获取。
要生成一个复现脚本,其中包含 Helion kernel 定义、配置器以及一个在调用 Helion kernel 之前重新创建运行时输入的 `helion_repro_caller()` 辅助函数,请设置 `HELION_PRINT_REPRO=1` 或在 `@helion.kernel` 装饰器中包含 `print_repro=True`。这会将复现脚本打印到 `stderr`,这对于调试以及在 GitHub issue 跟踪器上分享最小复现用例非常有帮助。
在 `hl.tile`/`hl.grid` 设备循环中,如果你想使用 `print("x", ...)` 语法打印中间结果,或者使用 Python 内置的 `breakpoint()` 暂停执行,请设置 `TRITON_INTERPRET=1`(运行 Triton 的 CPU 解释器)或 `HELION_INTERPRET=1`(在 eager 模式下运行 Helion kernel)。
要强制进行自动调优并绕过提供的配置,请设置 `HELION_FORCE_AUTOTUNE=1` 或调用 `foo_kernel.autotune(args, force=True)`。
[settings.py](https://github.com/pytorch/helion/blob/main/helion/runtime/settings.py) 中提供了其他设置。如果同时设置了环境变量和 kernel 装饰器参数,则以 kernel 装饰器参数为准,并且将忽略环境变量。
通过设置环境变量 `HELION_LOGS=all` 来启用 INFO 级别的日志记录,或设置 `HELION_LOGS=+all` 来启用 DEBUG 级别的日志记录。或者,你可以使用逗号分隔的列表为特定模块指定日志记录(例如,`HELION_LOGS=+helion.runtime.kernel`)。
## 环境要求
Helion 目前针对 Linux 系统,并且需要较新的 Python 和 PyTorch 环境:
- 基于 Linux 的操作系统
- Python 3.10–3.14
- [PyTorch] 2.9 或更高版本
- [Triton] 3.5 或更高版本
*(旧版本可能可以使用,但将缺少对 Hopper/Blackwell GPU 上 TMA 等功能的支持,并且可能表现出较低的性能。)*
- [Triton-to-tile-IR](https://github.com/triton-lang/Triton-to-tile-IR) *(可选)* 3.5 或更高版本
## 安装说明
我们建议使用 [uv] 来管理隔离的虚拟环境。首先,安装兼容版本的 [PyTorch] 和 [Triton]。
环境设置完成后,你可以安装 Helion:
```
pip install helion
```
或者,出于开发目的,你可以从源码安装。如果使用 `uv`,请先创建并激活虚拟环境:
```
git clone https://github.com/pytorch/helion.git
cd helion
# 使用 uv 创建并激活 virtual environment(一次性)
uv venv .venv
source .venv/bin/activate
# 以 editable 模式安装并包含必需的 dev packages
pip install -e .'[dev]'
```
这将以“可编辑”模式安装 Helion,以便对源代码的更改无需重新安装即可生效。
## Linting
我们使用 `pre-commit` 来自动运行 ruff、pyrefly 和其他检查。
– 一次性设置(安装 git hook):
```
pip install pre-commit
pre-commit install
```
– 在整个存储库中运行所有检查:
```
pre-commit run --all-files
```
注意:你仍然可以通过 `./lint.sh [fix|check|unsafe]` 直接运行底层工具。
## 社区
有疑问或反馈?请在 [GPU MODE Discord](https://discord.gg/gpumode) 的 `#helion` 频道加入我们。
## 许可证
Helion 采用 BSD 风格的许可证,详情请见 LICENSE 文件。