PyTorch Helion 引入编译期启发式，免自动调优即可生成内核配置

TL;DR

Helion 现在使用编译期启发式方法，分析内核的 IR 和数据流，从而在不进行缓慢自动调优的情况下选出快速设置。它将内核分类为 matmul、reduction 或 pointwise，并选择块大小、warp 数和流水线阶段数，免费获得接近自动调优的速度。

引言

Helion 是 PyTorch 用于性能可移植机器学习内核的领域特定语言（DSL），其性能高度依赖自动调优。自动调优时间可能成为主要瓶颈：例如，此前一个团队在 72 种形状上调优单个 paged-attention 内核，在 Helion 的快速模式下花费了 10 小时，全力调优则花费 25 小时。为解决这些问题，我们开发了一组启发式方法，在编译期读取内核的 IR 和数据流，直接生成配置，无需任何自动调优。

我们根据内核包含的操作将其划分为大类——matmul、reduction 或 pointwise——并对每一类应用特定的启发式方法。这些类别有意保持宽泛：例如，matmul 启发式会查找任何包含 matmul 操作（如 hl.dot 或 torch.baddbmm）的内核，而不仅限于标准的 m-n-k 矩阵乘法。这些方法被设计为通用且广泛适用于用户可能编写的自定义内核。

亮点：

• 自动调优器启发式方法。我们的启发式方法覆盖大量内核，包括 GEMM、reduction，甚至线性注意力。
• 线性注意力。在覆盖九种线性注意力变体的超过 96 个内核与形状组合上，计算出的配置在 H100 上前向传播比 Flash Linear Attention（FLA）中手写的 Triton 快 1.18 倍，包含反向传播时快 1.12 倍。总体而言，在 H100 上，该启发式方法比 Helion 此前无自动调优的默认配置快 2.13 倍。
• vLLM reduction 内核。在 H100 上覆盖六个生产内核的 153 个内核与形状组合上，计算出的配置比 vLLM 自己手写的 CUDA/C++ 算子快 1.27 倍，比 Helion 此前的默认配置快 2.2 倍。
• 分析本身开销很低。启用与禁用启发式方法编译内核的差异微乎其微（数十毫秒），淹没在 Triton 编译的噪声中。

我们的方法：在编译期决定配置

单个 Helion 内核可以降级为数千个不同的 Triton 内核。一个定义涵盖块大小、循环顺序、索引模式和流水线策略，但自动调优器必须通过编译和基准测试许多候选方案，来决定生成这数千个中的哪一个。而且这种自动调优成本是按内核和按形状计算的，你希望支持的内核、形状和硬件后端越多，问题就越成倍放大。

我们在编译期计算配置，而不是在调优期搜索配置。我们读取设备 IR 并记录其结构事实：哪些轴被分块、哪些块大小可调、哪些块在内核的哪个部分活跃，等等。启发式方法读取这些事实并生成配置。

这里的重点不是提出比自动调优器更好的配置——后者拥有在选定一个之前对可能数千个配置进行基准测试的优势。而是要在性能与等待时间之间的权衡空间中增加一个点。

我们的启发式规则有两个用途。主启发式规则会生成在禁用自动调优时使用的默认配置。当启用自动调优时，辅助启发式规则会生成大约 5–20 个相关或备选配置，作为自动调优器的种子。两者的目标不同：默认配置应当自身表现良好，而种子则应为自动调优器提供若干有希望的起点，以加速自动调优。在本篇博客中，我们会涉及两者，但重点放在禁用自动调优时生成的主启发式规则上。

启发式规则如何工作

启发式规则为各种配置参数设定取值，这些参数必须仔细调优才能获得高性能，通常由自动调优器来选择。这些参数包括每个分块轴的块大小、num_warps、num_stages、L2 分组、加载驱逐策略等。

我们通过扫描设备 IR 中内核所包含的操作来选择启发式规则。任何 matmul 都会路由到 matmul 启发式规则，无论内核中还有什么其他内容；没有 matmul 的归约会路由到归约启发式规则；其他所有情况则使用逐点启发式规则。如果内核的各个维度范围并非都在编译时已知，则不会触发任何启发式规则，Helion 会回退到原始默认值。

我们的启发式规则都依赖同一套底层分析，该分析估算内核的资源使用情况。在内核的任意步骤，该分析会确定：给定一组候选块大小，一个 CTA 在该时刻需要多少寄存器、共享内存，以及在 Blackwell 上需要多少张量内存。

为此，我们对编译器 IR 进行活跃性分析，并遍历每个图以找出每个值的最后使用位置，同时猜测它将被放置在内存中的何处。这是一个难题：我们试图预测 Triton 会把什么保留在寄存器中、会把什么通过共享内存暂存、会把什么实际物化而不是融合掉。有时这比较容易预测（例如，我们假设 hl.dot 的输出在 B200 上会被放入 TMEM），但估算只是近似的。

利用这一点，我们可以估算对性能有重要影响的事项，例如：

• 占用率：一次能在 SM 上容纳多少个 CTA。哪种空间先耗尽就由它决定，例如，一个 CTA 需要 SM 十分之一的寄存器和十分之一的共享内存，但需要十分之六的张量内存，则每个 SM 只能容纳一个 CTA，因为两个 CTA 会超出张量内存预算 20%（0.6 x 2 = 1.2）。
• 寄存器压力：我们计划启动的线程上峰值活跃寄存器的分布，告诉我们汇编器是否会被迫溢出。

这些考量会影响不同配置参数的选择。我们以 matmul 启发式规则为例来说明：

• block_size：如果启动产生的 CTA 数量少于 GPU 的 SM 数量，我们会在某个网格轴上缩小块大小以产生更多 CTA。
• num_stages：更多阶段能隐藏更多加载延迟，因此我们希望在共享内存允许的范围内尽可能多。但阶段数也可能意外损失占用率：六个阶段可能放得下，但会吃掉 CTA 的全部共享内存预算，导致每个 SM 只有一个 CTA，而三个阶段、每个 SM 两个 CTA 才是更好的权衡。
• num_warps. Two things pull on this:
  • Triton 发出哪条指令。在四个 warp 以下，或 M 块小于 64 时，Triton 使用较小的同步 MMA；否则使用高吞吐的大指令（Hopper 上是 wgmma，Blackwell 上是 tcgen05）。通常我们想要大的那个。但当 matmul 工作量很小时，有时更好的答案是每个 SM 上运行多个 CTA，执行彼此重叠加载的小型 MMA。因此我们统计动态 matmul 工作量（本质上是统计 hl.dot 被动态执行的次数），并据此选择。
  • 寄存器压力。warp 越多，tile 被分摊到更多线程上，因此每个 warp 的压力越小。如果当前 warp 数量下的峰值活跃状态会导致溢出，那就倾向于增加 warp。
• 可行性检查。对整个 kernel 做最后一次悲观检查：累加 tensor memory、共享内存和寄存器，确认该配置确实能够启动。上述调整已经遵守了这些预算，所以这一步主要是防御性的。

当一个 kernel 中有多个 matmul 时（例如 attention 和 linear attention），我们对每个 hl.dot 分别运行 matmul 启发式，并以这些结果作为起点。它们可能互相冲突：两个 dot 可能共用一个 block-size 旋钮——一个 dot 的 N 常常是下一个 dot 的 K——而整个 kernel 只有一个 num_warps。

为解决这些权衡，我们按估计的重要性对每个 dot 排序：优先级最高的是带循环携带累加器的 dot，因为该累加器在整个循环期间常驻，其 tile 决定了 kernel 的占用；其次是它的动态 matmul 工作量；再次是它自身的输出区域。当两个 dot 对某个旋钮意见不一致时，排名更高的获胜（例如 num_warps 仅由排名最高的 dot 决定）。

reduction 启发式通过一组类似的旋钮工作：

• num_warps. Three pressures here, pulling in different directions:
  • 有足够的 warp 来做 reduction。宽 reduction 需要多个 warp 参与，因此从最大的被选 reduction 出发的阶梯确定起点：128 个元素需要一两个 warp，16,384 个则需要更多。
  • 常驻性。每个 CTA 的 warp 越多，每个 SM 的 CTA 就越少（存在每 SM 的 CTA 上限），所以我们也会测试更少的 warp。
  • 寄存器压力。调整 warp 以避免寄存器溢出。
• block_size。每个 SM 容纳更多 CTA 意味着可能使用更小的 tile，因此我们缩小 block 以换取常驻性。约束是：stride 为 1 的轴是承载合并加载的轴，把它缩得太小会失去加载合并。
• 持久化。reduction 可以保持常驻并一次完成，也可以分块处理。真正重要的情况是 kernel 读取两次的行——RMSNorm 先读一次计算 scale，再读一次应用它——当我们能证明这一点时，允许持久化付出一些溢出和常驻性的代价，因为它省去了第二次内存访问。

由于 warp 数量取决于 block，而 block 又取决于 warp 数量，我们迭代直到两者在一整轮中都不再变化（实践中这不会花很长时间）。

结果 1：配置质量，且无需调优时间

Linear attention

我们在 H100 上测量了九种线性注意力变体在仅前向和前向加反向两种情况下的表现，覆盖六种生产形状。这六种形状涵盖 batch 1 到 8、8 到 96 个头、序列长度 1024 到 16384、头维度 64 到 256。基线是 Flash Linear Attention（FLA），一个开源且被广泛使用的手写 Triton 内核库。图表中的自动调优分支（“Helion AOT config”）是对每种形状运行一次完整自动调优器并保存胜出配置的结果——也就是如今针对已知形状集合发布 Helion 内核的方式。

在这项工作之前，选择是二元的：采用 Helion 的无自动调优默认配置，运行速度只有手写 FLA Triton 的 0.52x；或者花费 GPU 小时进行自动调优，达到 1.26x。该启发式方法增加了此前不存在的第三个选项：1.18x，且无需付出自动调优的代价（自动调优所发现的大部分收益可立即获得），并且它使 Helion 的开箱即用行为从输给手写 Triton 变为胜过它，至少对这组内核而言如此。

vLLM 的归约内核

vLLM 是一个被广泛部署的开源 LLM 推理引擎，它为其服务热路径上的操作自带手写 CUDA/C++ 内核。共有六个归约内核——两个 FP8 量化器、分别与量化融合的 RMSNorm 和 SiLU-and-multiply，以及融合的 QK 归一化加 RoPE——vLLM 已经有这些内核的 Helion 实现，由本系列早前一篇文章背后的团队贡献。这使它们成为很好的测试对象：由他人编写的生产内核，对比它们本要取代的手写 CUDA。

我们在另一个代码库中得到了类似的结果。Helion 的无自动调优默认配置运行速度是 vLLM CUDA 的 0.58x；计算出的配置达到 1.27x；完整自动调优达到 1.31x。有趣的数字是最后这个差距：在这些内核上，自动调优相比我们免费计算出的配置只带来约 3% 的提升。

日常内核，对比 torch.compile

上述两项比较都是专业库——线性注意力变体和 FP8 量化器。这一项则刻意相反：人人都写的内核。RMSNorm、LayerNorm、softmax、交叉熵等，类似于介绍文章首次测量 Helion 时基准测试的套件。本节中基准测试的大多数内核可以在这里找到。

这里的参考基线是 torch.compile(mode=”max-autotune-no-cudagraphs”)：我们强制 torch.compile 使用其 Triton 后端，这样我们对比的不是手写库（例如 cuBLAS），而是来自广泛使用编译器的生成 Triton。本节没有调优配置分支，原因很平常：我们没有这些内核的完整自动调优运行结果（那将需要许多 GPU 小时来生成）。

在一个包含十个归约内核和 80 个（内核、形状）组合的套件上，计算出的配置达到 max-autotune torch.compile 的 1.09x，且自身无需调优。Helion 的未播种默认配置为 0.27x。

我们还在 70 种形状上扫描了十三个 matmul 内核家族。计算出的配置总体达到 max-autotune torch.compile 的 1.88x，而 Helion 的默认配置为 0.51x。关于这个 1.88x 有两点说明。这些比较并不总是完全对等的——对其中一些内核，PyTorch 实现会物化中间结果，而 Helion 内核会将其融合。并且 jagged HSTU attention 确实出现了回退：旧的默认配置恰好在那里落到了一个好的配置上（3.72x），而启发式方法表现更差（3.27x）。

结果 2：种子能让搜索更快吗？

以上所有内容都是关于 Helion 在自动调优关闭时输出的配置。但同样的配置可以作为自动调优器的初始种群交给它，原则上，一个好的起点应该能让搜索更早完成。除了这个主配置之外，辅助启发式方法还会添加大约 5 到 20 个新的变体/方法。

为了测试这一点，我们选取了 38 个 matmul 和 multi-matmul 内核——线性注意力、注意力、Mamba，以及 squeeze-and-excitation 等其他类型——每个内核选取一种形状。然后我们对每个单元格运行了两次完整的自动调优：一次从随机配置开始，一次将编译器的种子放入初始种群。

这次运行的简短答案是：种子在早期带来很大收益，但在后期收益小得多。下面是一个单元格，被选为该模式的清晰示例。x 轴表示尝试的配置数量，y 轴表示相对于我们找到的最佳配置的性能。

使用种子时，第一个配置远优于未使用种子的第一个配置。但未使用种子的搜索很快赶上，并且当搜索继续处理数百个更多配置时，它们最终落在相同的答案上。早期优势巨大，而最终优势为零。

这就是全部 38 个单元格的模式。中位数如下；计数包括编译失败或未能产生计时的配置，因为它们也消耗预算。

优势完全集中在开始阶段：达到仅仅尚可的水平所需配置少 35x，达到 25% 以内少 41x，但达到 5% 以内仅少 1.5x，完成搜索仅少 1.2x。

这个结果并不意味着：种子不会让最终配置变得更好。在同一 GPU 上重放两组的获胜者，它们的性能相差不到百分之一。这关乎搜索需要多长时间，而不是它最终落在哪里。

局限性与未来工作

一个核心局限是这些启发式方法预测的是另一个编译器会做什么。每个预算都依赖于猜测一个值最终落在哪里——寄存器、共享内存还是张量内存——而这是 Triton 的决定，不是我们的。其中一些可以安全地押注：在 B200 上，matmul 结果会位于张量内存中。很多地方我们会猜错。未来工作有几个可能的方向。

测量而非预测。 对上述局限最干净的修复是停止猜测：编译一次，读取汇编器报告的寄存器和共享内存使用情况，并据此修正估计。

覆盖更多内核。 目前，当我们缺乏信息来估算任何规模时——例如，一个循环在一个仅在运行时才可知的范围内进行分块——我们会选择放弃。做一个信息较少的猜测会很容易：一个默认范围、一个均值，诸如此类。没有这些信息，配置会更差，但如今的替代方案是基础默认值，而上述结果表明这非常糟糕——所以这个猜测可能仍然值得一试。

触及更多旋钮。 我们设置了块大小、num_warps 和 num_stages，以及一些 L2 分组和 L2 驱逐策略。这留下了很长一串未触及的选项——张量描述符与指针算术、持久化与扁平启动等。单独来看，这些可能不如我们设置的选项有价值，但肯定还有性能留在桌面上。

特化启发式。 它们被有意设计得通用：一个多重矩阵乘法启发式必须同时在注意力和线性注意力上表现良好，一个归约启发式必须同时在 RMSNorm 和 vLLM 的量化内核上表现良好。另一种选择是父启发式与子启发式——一个 vLLMQuantizationHeuristic 从 ReductionHeuristic 的答案出发，然后根据它对该家族的了解进行特化。

其他后端。 这里的一切都在 NVIDIA 上评估，启发式在 Triton 后端上为 H100 和 B200 默认启用。该方法中没有任何 NVIDIA 特有的内容——分析由应适用于各后端的硬件限制参数化——但常量和阶梯是针对这两款芯片调优的。

结论

这些启发式已合并到 Helion 中，并在 Triton 后端上为 H100 和 B200 默认启用。如果您想了解更多，请查看 autotuner heuristics。

重点不是静态分析胜过搜索，而是差距现在小到对某些内核来说跳过搜索是合理的。Helion 的开箱即用行为从手写 FLA Triton 的 0.52x 提升到 1.18x，从 vLLM 手写 CUDA 的 0.58x 提升到 1.27x，且没有编译任何东西，也没有进行任何基准测试。在一般的归约内核套件上，它达到 torch.compile max-autotune 的 1.09x。当你发布一组固定的形状并且可以一次性支付成本时，完整自动调优仍然是正确的选择；改变的是替代方案不再是一个巨大的让步。

文章 Fast by Default: Compile-Time Heuristics for Helion Kernels 首次出现在 PyTorch。