pytorch/helion

GitHub: pytorch/helion

Helion 是一种嵌入式 Python DSL,旨在通过自动调优大幅简化高性能 ML kernel 的编写,并编译为 Triton 后端代码。

Stars: 904 | Forks: 158

Helion Logo
# 活动 - **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 文件。
标签:Android, Apex, DSL, GPU编程, PyTorch, Triton, Vectored Exception Handling, 凭据扫描, 机器学习, 算子开发, 编译器, 逆向工具