引言
正如在之前的博客文章中所介绍的，Helion 是一种高级 DSL（领域特定语言），它使开发者能够使用熟悉的类似 PyTorch 的语法编写高性能 ML 内核，并将复杂的优化任务委托给其自动调优引擎。该自动调优器会探索一个巨大的、高维的实现选择空间（包括块大小、循环顺序、内存访问模式等），以发现能在目标硬件上实现性能最大化的配置。因此，Helion 可以比 torch.compile 甚至 Triton 或 CuTe DSL 中高度优化的手写内核实现显著的加速。
然而，自动调优带来的性能提升是有代价的：漫长的挂钟时间（wall-clock time）。典型的自动调优过程可能需要 10 分钟以上，评估数千个候选配置，对于复杂的内核，甚至可能需要数小时。自发布以来，漫长的自动调优时间一直是用户反映最多的问题，也是内核开发周期中最大的痛点之一。虽然 Helion 为开发者提供了缩短自动调优过程的选项（例如减少搜索步数），但这通常会导致内核性能下降，从而产生一种不理想的权衡。
在这篇博客中，我们讨论了为改善自动调优体验所做的持续努力。特别是，我们介绍了一种为解决这些问题而开发的新型搜索算法——LFBO 模式搜索 (LFBO Pattern Search)，它采用了机器学习 (ML) 技术来提高自动调优引擎的效率。该搜索算法通过训练一个 ML 模型来智能地过滤候选配置，从而大幅减少了需要评估的候选对象数量。重要的是，该模型仅使用搜索过程中收集的数据，无需用户提供任何额外数据。
利用机器学习，我们可以在不牺牲性能的情况下大幅减少自动调优时间
- 在我们的一组基准 NVIDIA B200 内核上，我们将自动调优时间缩短了 36.5%，同时平均内核延迟降低了 2.6%。
- 在 AMD MI350 内核上，我们将自动调优时间缩短了 25.9%，同时内核延迟降低了 1.7%。
对于某些内核，改进效果尤为显著：对于 B200 LayerNorm 内核，挂钟时间减少了多达 50%，而对于 B200 Helion FlashAttention 内核，内核延迟甚至改善了 >15%。由于其优越的性能，它已成为撰写本文时的默认搜索算法。
内核自动调优的挑战
自动调优引擎会搜索内核配置，对它们的延迟进行基准测试，并利用结果确定下一组要测试的配置。虽然编译和测量单个配置的延迟仅需几秒钟，但自动调优引擎通常需要搜索数千个配置才能达到最佳性能。由于设计空间内在的几个因素，寻找最佳内核配置是一个极具挑战性的优化问题：
- 高维组合空间： 块大小、展开因子等所有可能组合的空间是高维且巨大的。即使是像 LayerNorm 这样简单的内核，也有超过 8 千万亿（10^16）种可能的配置。然而，虽然搜索空间很大，但只有极少数配置具有良好的性能。
- 编译时间长： 某些内核配置的编译可能需要大量时间，从而不必要地延长了自动调优过程的挂钟时间。
- 配置错误和超时： 搜索空间可能还包含会导致编译错误、产生不准确结果或编译时间过长的配置。
之前的默认搜索策略（模式搜索）从多个有前景的配置（“搜索副本”）开始，通过穷举评估所有单参数扰动来探索相邻配置。虽然该方法很彻底，但效率低下：绝大多数相邻配置并不能带来性能提升，但每一个都要经过编译和基准测试。此外，将移动限制为单参数更改，限制了算法快速穿越高维搜索空间的能力。
无需似然函数的贝叶斯优化模式搜索 (LFBO Pattern Search)
为了解决这些低效问题，我们借鉴了贝叶斯优化 (Bayesian Optimization) 的灵感。这是机器学习的一个子领域，利用概率代理模型（例如高斯过程）来智能地选择下一个要评估的点（可在 botorch 和 Ax 等库中使用）。为了尽量减少额外的挂钟时间，我们采用了无需似然函数的贝叶斯优化 (LFBO)，它使用轻量级的分类模型作为代理。我们将模式搜索的局部搜索启发式方法与 LFBO 分类模型相结合，仅过滤出最有希望的候选者进行基准测试，而不是进行穷举搜索。
LFBOPatternSearch 算法步骤如下：
- 与 PatternSearch 类似，我们首先对一组随机生成的配置进行基准测试，并识别出一小部分最有前景的配置（“搜索副本”）。
- 我们通过在多个参数上进行随机扰动从搜索副本中生成候选对象，探索范围比 PatternSearch 更广。
- 我们在迄今收集到的延迟数据上训练一个分类模型（随机森林）。我们不直接预测延迟，而是预测一个二元标签，指示该配置在延迟方面是否属于前 10%。
- 我们根据 ML 模型的预测对候选对象进行排名。与典型的 LFBO 不同，我们还为与先前排名候选对象的相似性添加了惩罚项，以鼓励进一步的探索。
- 我们选择前 10% 的候选对象进行编译和基准测试。我们根据性能最好的配置更新搜索副本，并将延迟添加到数据集中。
我们讨论了一些关键的设计决策，这些决策对于实现我们所观察到的挂钟时间和延迟的改进至关重要：
分类与回归： 基于回归的方法（即训练模型直接预测延迟）是系统/编译器研究中成本建模的事实标准。然而，我们发现基于分类的方法能更好地将模型容量集中在性能最好的配置上，而不是试图学习所有配置（无论好坏）的延迟。其次，分类损失使模型能够学会避免那些报错或超时编译的配置（因为它们被标记为负标签）。然而，对于回归方法来说，这些点没有任何有效的延迟数据可供学习。
鼓励多样性： 通常配置是批量编译的，以利用并行预编译的优势。随机森林分类器可能会重复选择聚集在一起的相似配置，这可能会在冗余样本上浪费批处理预算，从而提供较少的新信息。为了缓解这种情况，我们根据随机森林模型中的叶节点共现情况计算相似度分数，并对与先前排名配置的相似性进行惩罚。
当我们调查 LFBO 模式搜索的行为时，我们发现各内核、硬件类型和形状的性能改进确实源于它能够通过更少的评估找到具有更好运行时的配置。下图是 B200 LayerNorm 内核的自动调优轨迹示例图，显示了自动调优器随时间获得的最佳配置的延迟。我们不仅看到 LFBO 更早地完成了自动调优（约 5 分钟而不是约 9 分钟），而且与模式搜索相比，它能以更大幅度的性能提升更快地找到更好的配置。
我们看到 LFBO 通过比模式搜索更广泛的探索实现了这一点。下图是 LFBO 模式搜索和模式搜索为同一个 B200 LayerNorm 内核采样的配置图，我们应用主成分分析 (PCA) 进行可视化（因为配置是高维的）。我们看到，虽然 LFBO 模式搜索评估的配置数量不到前者的一半，但其采样的配置比模式搜索的分布更广，后者由于仅进行单参数扰动而高度聚集在一起。在分类器的指导下，LFBO 模式搜索能够进行更大规模但更有针对性的跨越。
最后，我们使用自动调优器日志（从 PatternSearch 收集）数据集，与其他代理模型（特别是涉及随机森林、梯度提升树和多层感知机 (MLP) 的回归方法）进行了消融实验。我们计算了一个与自动调优器性能最直接相关的指标：使用代理过滤下一批候选对象时，内核延迟的预期改进。下图中，我们绘制了预期改进（以延迟的相对百分比改进计）与代理被允许选择的候选对象百分比的关系。我们发现，基于 LFBO 的方法提供了最大的预期改进，而多样化选择也带来了有意义的改进。值得注意的是，当我们只能选择 10% 的候选对象时，回归方法的效果等同于甚至劣于简单的随机选择，因为回归并不总是与排名性能保持一致。
结论
在这篇博客中，我们展示了机器学习 (ML) 如何加速自动调优引擎并改善 Helion 中的内核开发体验。通过利用在搜索过程中收集的延迟数据，我们可以让自动调优器专注于更有前景的配置，从而节省时间并发现更快的内核配置。我们非常乐意应用更多的机器学习技术来增强自动调优器，包括来自强化学习 (RL) 和大语言模型 (LLM) 的方法，并欢迎任何贡献。