Helion 编译时启发式是什么:PyTorch 称不用自动调优也能拿到接近调优的内核速度,此前一个算子调一次要 10 小时
PyTorch 团队称,此前有团队为单个 paged-attention 内核在 72 种形状上做自动调优,仅 quick 模式就花了 10 小时,full 模式则需 25 小时。为解决这一瓶颈,他们为 Helion(PyTorch 的高性能可移植内核 DSL)开发了一套编译时启发式:在编译阶段读取内核的 IR 与数据流,直接给出配置参数,跳过自动调优。在线性注意力的 96 个内核×形状组合上,该启发式前向性能超过手写 FLA Triton 内核 1.18 倍(H100)。
它解决什么问题
据 PyTorch 官方博客,Helion 的一个内核定义可以展开为数千种不同的 Triton 内核,自动调优器需要逐一编译和基准测试大量候选配置来选出最优解,而这个成本是按内核×形状×硬件后端叠加的。此前 Helion 的选择是二元的:要么接受未调优的默认配置,在线性注意力上只有手写 FLA Triton 的 0.52 倍性能;要么花费大量 GPU 时间做自动调优,达到 1.26 倍。编译时启发式的目标是在性能与等待时间之间增加一个新的折中点——不编译、不基准测试,直接给出一个足够好的配置。
它怎么做
据 PyTorch 博客,启发式的工作方式可以拆为以下几步。
第一步:分类。 扫描内核的设备 IR(中间表示),根据其中包含的操作将内核归入三个族:matmul(包含任何矩阵乘操作如 hl.dot 或 torch.baddbmm)、reduction(包含归约但无矩阵乘)、pointwise(其余)。如果内核的维度在编译时不完全已知,则不触发启发式,退回原始默认配置。
第二步:资源估算。 对编译器 IR 做活跃性分析(liveness analysis),在内核的每个步骤上估算给定候选 block size 下一个 CTA(协作线程阵列)需要多少寄存器、共享内存,以及在 Blackwell 架构上需要多少张量内存。由此推算占用率(occupancy,即一个 SM 上能同时容纳多少 CTA)和寄存器溢出风险。PyTorch 团队承认这是近似估算,因为最终由 Triton 编译器决定值的存放位置。
第三步:按族选参数。 据 PyTorch 博客,以 matmul 启发式为例:如果当前配置产生的 CTA 数量少于 GPU 的 SM 数量,就缩小 block size 以增加并行度;num_stages(流水线级数)尽量多以隐藏加载延迟,但要避免共享内存占满导致占用率下降;num_warps 则需要平衡指令选择(四个 warp 以上且 M block 不小于 64 时才会使用 Hopper 的 wgmma 或 Blackwell 的 tcgen05 高吞吐指令)和寄存器压力。当一个内核包含多个矩阵乘时,按估算重要性排序——优先考虑带循环累加器的 dot,因为其累加器在整个循环中常驻——由排名最高的 dot 决定共享参数。
第四步:可行性检查。 最后做一次悲观估算,确认张量内存、共享内存和寄存器加总后配置确实能启动。
据 PyTorch 博客,启发式分析本身的开销可以忽略不计,与关闭启发式相比仅多出几十毫秒,淹没在 Triton 编译的噪声中。
和以前的做法有何不同
据 PyTorch 博客,在线性注意力的 9 种变体、6 组生产形状(批大小 1–8、8–96 头、序列长度 1024–16384、头维度 64–256)共 96 个内核×形状组合上,启发式配置前向超过手写 FLA Triton 1.18 倍,含反向为 1.12 倍;比 Helion 此前的未调优默认配置快 2.13 倍(H100)。完整自动调优的结果为 1.26 倍,启发式拿到了其中的大部分收益。
PyTorch 博客还给出了 vLLM 归约内核的数据:6 个生产内核在 H100 上的 153 个内核×形状组合中,启发式配置比 vLLM 自己的手写 CUDA/C++ 算子快 1.27 倍,比 Helion 旧默认快 2.2 倍;完整自动调优为 1.31 倍,意味着调优仅比启发式多出约 3%。
据 PyTorch 博客,在通用内核方面,10 个归约内核的 80 个组合上,启发式配置达到 torch.compile max-autotune-no-cudagraphs 的 1.09 倍;13 个 matmul 内核族的 70 个形状上达到 1.88 倍,但团队注明后者并非完全对等比较——部分 PyTorch 实现会物化中间结果而 Helion 做了融合,且 jagged HSTU attention 上启发式(3.27 倍)反而不如旧默认(3.72 倍)。
此外,启发式配置还可作为自动调优的种子。据 PyTorch 博客,团队在 38 个 matmul 及多 matmul 内核上测试发现,使用启发式种子后达到”尚可”水平所需的候选配置数量减少到无种子时的约 1/35,达到最优 25% 以内减少到约 1/41,但达到最优 5% 以内仅快 1.5 倍,最终结果与无种子搜索几乎相同。团队总结:种子加速了搜索前期,但不改善最终配置质量。
目前的边界
据 PyTorch 博客,启发式目前已合并进 Helion 并在 H100 和 B200 的 Triton 后端上默认启用。团队明确列出了几项限制:启发式本质上是在预测 Triton 编译器的行为,资源估算是近似的;当循环所遍历的范围仅在运行时才已知时,启发式不会触发,退回旧默认;目前只设置 block size、num_warps、num_stages 及部分 L2 策略,张量描述符与持久化启动等参数尚未覆盖;所有常数和阶梯值仅针对 H100 和 B200 调校,虽然方法本身不绑定 NVIDIA,但尚未在其他后端验证。PyTorch 团队称,完整自动调优在需要为固定形状集合发布内核且能承担一次性成本时仍是更优选择;Helion 开箱行为已从 0.52 倍手写 FLA Triton 提升到 1.18 倍,无需任何编译或基准测试。
主题分类:GPU 编译器/内核
参考资料
- Fast by Default: Compile-Time Heuristics for Helion Kernels — PyTorch, 2026-11-13