pytorch/helion

Helion – 一个用于编写 GPU 内核的 Python 嵌入式 DSL

它是什么 – Helion 允许你直接在 Python 中编写 内核 代码（在 GPU 上运行的小程序），其语法感觉就像普通的 PyTorch。在底层，代码会被转换为 Triton 内核，这是一种低级语言，可编译为快速的 GPU 机器代码。

为什么重要 – 编写高效的 GPU 内核很困难：你必须决定分块大小、网格维度、内存访问模式以及许多其他低级参数。Helion 自动化了大部分工作。一个 Helion 函数会被编译为 一个 GPU 内核，但在首次启动之前，它会运行一个 自动调优器，尝试数百种 Triton 变体并选择最快的。结果是内核既易于编写，又接近手工调优的性能。

主要特性

• Python 优先的 API – 使用 @helion.kernel() 定义内核，并在内部使用普通的 PyTorch 张量。外部 Python 代码在 CPU 上运行（分配、形状检查），内部循环变成 GPU 代码。
• 自动分块和网格选择 – hl.tile 分割迭代空间；Helion 决定分块维度、循环顺序和网格大小。
• 丰富的自动调优 – 搜索大型配置空间（块大小、循环顺序、索引策略、warp 特化等），并可以缓存最佳配置以供后续运行。你也可以提供自己的配置以跳过昂贵的搜索。
• 内存访问策略 – 自动在基于指针、块指针或 TensorDescriptor（TMA）索引之间选择，并支持覆盖。
• 支持归约、掩码和持久内核 – 无需额外代码即可处理常见模式，如 softmax、求和或大型归约。
• 后端 – 主要是 Triton，但也有实验性的 TileIR（用于 NVIDIA Blackwell GPU）和 CuTe 后端（CUDA 13+）。
• 调试和开发工具 – 打印生成的 Triton 代码，生成可复现的脚本，并在解释器模式下运行内核以进行逐步调试。

典型工作流程

1. 使用 @helion.kernel() 和 hl.tile 编写内核函数。
2. 使用 CUDA 张量调用该函数。第一次调用会触发自动调优（矩阵乘法示例约需 10 分钟）。
3. 可选地提取最佳的 helion.Config 并硬编码以避免未来的调优。
4. 使用保存的配置部署内核，以获得可预测的启动时间。

谁应该使用 – 需要自定义 GPU 内核但不想手写 Triton 或 CUDA 的研究人员和工程师。它非常适合 PyTorch 代码库，并且可以加速现有库缺乏专门实现的负载。

安装 – pip install helion（需要 Python 3.10-3.14、PyTorch ≥ 2.9、Triton ≥ 3.5）。开发设置使用 uv 工具，并为 CuTe 后端提供可选扩展。

社区 – 讨论在 GPU MODE Discord（#helion）上进行。该项目遵循治理模型，并采用 BSD 许可证。