Triton for MTIA: Bridging the Programming Model Gaps for Custom AI Accelerators
本文介绍了 Triton 编程语言在 Meta 自研 MTIA-2i 加速器上的首次生产级部署，证明了新的编译器后端和极简的语言扩展能够实现高性能、高效的算子开发，从而弥合机器学习框架与定制硬件之间的差距。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下，你正试图建造一个庞大的、高速的火车系统，用以在大陆间运送乘客（数据）。多年来，市面上只有标准的、高科技的电力机车（GPU）可用，大家都知道如何驾驶它们。这些火车有自动门、内置 GPS，且行驶平稳。但最近，一种新型的火车引擎被发明了（定制 AI 加速器）。这些新引擎极其强大，专为重型货物设计，但它们很古怪。它们没有自动门，你必须手动开关；它们没有 GPS，你必须自己绘制地图；而且它们的运行轨道也与旧火车不同。
问题在于，建造火车车厢的工程师们（软件开发者）已经习惯了那种易于驾驶的旧款电力机车。如果他们想使用这些超快速的新引擎，就必须学习一种完全不同且困难的新驾驶方式，或者必须为每一条新路线从头开始建造一辆全新的火车车厢。这拖慢了所有进度。核心问题是：我们能否教会这种旧的、易于使用的驾驶风格，使其能在这些奇特的、新的引擎上运行，同时又不损失速度优势？本文就在探索这一点，试图在熟悉的标准 AI 软件世界与像 Meta 这样公司为运行其 AI 模型而构建的奇异定制硬件之间架起一座桥梁。
故事：教一个新引擎说“旧语言”
Meta，即 Facebook 和 Instagram 背后的公司，一直在构建自己的特殊计算机芯片，称为 MTIA-2i，以更快速、更廉价的方式运行其 AI 模型。你可以把 MTIA-2i 想象成一辆定制打造的赛车引擎。它在执行特定任务时表现惊人，但它的构造方式与大多数 AI 程序员所熟悉的标准引擎（GPU）大不相同。
在 AI 世界中，有一种流行的编程语言叫做 Triton。你可以把 Triton 想象成一个“通用遥控器”。开发者不需要编写复杂、混乱的代码来告诉标准 GPU 如何移动数据，而是编写简单、整洁的 Triton 代码，然后由一个智能编译器将其翻译成 GPU 能理解的机器语言。这就像是在对一位懂得如何与机器交流的翻译员说英语。
然而，Triton 最初仅为标准的 GPU 引擎设计。当 Meta 尝试将其应用于其定制的 MTIA-2i 赛车引擎时，它失效了。“遥控器”无法适配新引擎的按钮。MTIA-2i 引擎以异步方式工作（它接收指令并稍后执行），以一种奇特的“循环缓冲区”方式管理自身的内存，并且需要非常具体的指令来保持一切流畅运行。标准的 Triton 代码无法处理这些特性。
本文的研究内容：构建一个新的适配器
Meta 团队决定解决这个问题。他们并没有强行让旧的遥控器去适配，而是专门为 MTIA-2i 构建了一个新的编译器后端。这个新系统就像一个超级智能的适配器。它接收开发者编写的简单、易读的 Triton 代码，并将其完美地翻译成 MTIA-2i 引擎所需的复杂的底层指令。
以下是他们如何实现这一点的，我们使用了一些有趣的类比：
- “FIFO”自助餐线： 在标准 GPU 上，计算机像一个智能自助餐一样自动管理内存，盘子会被即时取走并更换。而在 MTIA-2i 上，它更像是一条手动的流水线，带有“先进先出”（FIFO）缓冲区。新编译器学会了完美地管理这条流水线，确保数据在机器需要时准确到达，而无需程序员手动推动每一个盘子。
- 双厨厨房： MTIA-2i 芯片拥有两个“核心”（就像厨房里的两位厨师）可以同时工作。编译器找到了如何在两者之间分配烹饪任务，使它们既不会互相碰撞，也不会闲置等待。它甚至允许专家级程序员精确地告诉厨师谁负责切洋葱，谁负责搅动汤汁，如果他们想玩得更高级的话。
- “之字形”配送路线： 想象你要向 64 户人家派送包裹。如果你只是沿着街道直线前进（线性路径），前几家会完成得很快，但最后几家则会等待很久。团队发现，“之字形”配送路线（来回穿梭）能更好地平衡工作量。他们为 Triton 添加了一个功能，让程序员可以轻松选择这种更聪明的路线。
结果：快速、灵活且面向现实世界
团队不仅是在实验室里完成了这项工作；他们已将其投入到现实世界的生产中。他们成功地将这些新的 Triton 内核部署到了 60 种不同类型的 AI 模型中。
以下是他们的发现：
- 速度： 使用简单的 Triton 语言编写的代码，运行速度与使用复杂的低级 C++（该芯片的“原生”语言）编写的代码一样快。事实上，对于通用矩阵乘法（GEMM）——一项繁重的数学任务——Triton 代码达到了该芯片理论最大速度的 80% 以上。
- 采用率： 在短短一个季度内，使用 Triton 的模型类型数量增加了三倍。
- 影响： Triton 现在处理了这些模型中 50% 的层以及 47% 的非 GEMM 执行时间。这几乎有一半的工作是由这个新的、易于使用的系统完成的！
为什么这很重要
本文认为，你不需要在“易于编写”和“超高速”之间做选择。在此之前，如果你想使用像 MTIA 这样的定制芯片，你必须成为一名低级编码巫师并从头开始编写一切。如果你想使用像 Triton 这样的高级语言，你就会被困在标准 GPU 上。
通过弥合这一差距，Meta 展示了高级语言可以被适配并运行在几乎任何定制硬件上。这意味着在未来，随着公司构建出更加奇特且专门化的 AI 芯片，开发者将不必为每一个芯片都学习一种新的、困难的语言。他们可以继续使用熟悉的、强大的工具，而编译器会处理新硬件中那些杂乱的细节。
论文还指出，虽然他们早前尝试过一种不同的、非常底层的语言叫做 KNYFE，但由于其学习难度过高，他们放弃了它。他们证明了坚持使用像 Triton 这样灵活的高级语言（即使进行一些小的定制化调整）才是保持 AI 开发快速高效的制胜策略。
简而言之，本文证明了通过正确的编译器魔法，我们可以让“通用遥控器”在几乎任何“电视机”上都能正常工作，即使是那些奇特的、定制制造的电视机，而且不会牺牲画质。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。