PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

在现代计算机内部，图形处理器（GPU）扮演着一个强大的引擎角色，旨在同时处理海量数据。为了充分发挥这个引擎的效能，软件开发人员编写专门的指令，告诉硬件如何精确地移动数据并执行计算。这些指令是用一种称为 PTX 的低级语言编写的，它赋予了程序员对机器内部运作机制的直接控制权。虽然存在更高层级的工具使这一过程变得更加容易，但这些工具有时会遗漏最新硬件世代的独特特性。随着这些机器的快速演进，使软件与它们的特定能力保持完美同步正成为一项艰巨的挑战。问题随之而来：人工智能，特别是大语言模型，能否学会自主编写这些精确的低级指令，还是说它仍然依赖于那些会导致性能损失的陈旧、通用的方法？

一个研究小组致力于通过创建一个名为 PTXBench 的新测试场来回答这个问题。他们想看看这些 AI 模型是否不仅能编写出可以运行的代码，还能编写出能积极利用最新 GPU 特定先进特性的代码。研究人员将重点放在了两款最强大的显卡——H100 和 B200 上，并在涉及人工智能所用复杂数学运算的任务上对模型进行了测试，例如矩阵乘法和注意力机制。其目标是衡量三个方面：代码是否在运行中无错误、代码是否确实触发了研究人员要求的特定硬件指令，以及代码的运行速度是否快于现有的最佳软件库。

结果揭示了模型的实际能力与顶尖性能要求之间存在显著差距。虽然 AI 模型在编写适用于简单任务且功能正确的代码方面通常很成功，但在处理更复杂的运算时却显得相当吃力，尤其是涉及注意力机制反向传播（backward pass）的任务，而这对于训练 AI 系统至关重要。即使模型成功编写出了能够执行所请求硬件指令的代码，这些代码也很少能达到领先专业库的速度。在许多情况下，AI 生成的代码运行速度较慢，这表明仅仅让指令运行起来并不足以实现真正的优化。研究发现，模型经常退回到使用通用的编码模式，而不是利用新型芯片独特的、针对特定架构的特性，即便这些特性已被明确描述给它们。

为了了解如何改进这些模型，研究人员进行了一项受控实验，通过一种称为“监督微调”的方法来教导特定的 AI 模型。他们向模型展示了其自身的错误示例，以及经过修正的代码和关于为何这些修正有效的解释。这种被称为 Fixit 的方法有助于模型在多项任务上有所提升，特别是在生成功能正确代码方面。然而，这种提升并非均衡分布。模型仍然难以将其学到的新技能泛化到不同的问题规模或其在训练期间未曾见过的任务变体中。研究人员发现，修正的质量和训练数据的平衡性与数据的纯粹数量同样重要。在一个多样化问题集上训练的模型，其表现优于在一个规模更大但缺乏变化的数据集上训练的模型，这表明学习体验的类型比容量更为重要。

该研究还强调了为 AI 提供正确上下文的重要性。当研究人员向模型提供有关硬件特定能力的详细信息时，模型更有可能使用所请求的指令。如果没有这些特定知识，模型很少会使用请求的 PTX。这表明，虽然 AI 可以学习进行这种低层级的编程，但它需要精确的引导和高质量的示例才能有效地做到这一点。研究人员得出结论，虽然目前的 AI 模型展现出了潜力，但它们尚未准备好取代专家级人类工程师来为最新的硬件优化代码。未来的路径在于更好的训练数据、更平衡的学习体验，以及更深入地理解如何教导这些模型去推理不断演进的计算机架构中的特定约束。