RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators

本文提出了 RedFuser 框架,通过形式化理论方法自动识别并融合级联归约操作,在 AI 加速器上生成优化内核,从而在多种工作负载中实现比现有编译器快 2 至 5 倍的加速,并达到手写内核的性能水平。

Xinsheng Tang, Yangcheng Li, Nan Wang, Zhiyi Shu, Xingyu Ling, Junna Xing, Peng Zhou, Qiang Liu

发布于 Thu, 12 Ma
📖 1 分钟阅读☕ 轻松阅读

Each language version is independently generated for its own context, not a direct translation.

这篇文章介绍了一个名为 RedFuser 的新工具,它就像是一个**“超级大厨”**,专门负责优化人工智能(AI)模型在芯片上的烹饪过程,让 AI 跑得更快、更省电。

为了让你更容易理解,我们可以把 AI 模型的运行过程想象成在一家繁忙的餐厅里做一道复杂的菜

1. 遇到的难题:繁琐的“流水线”

在 AI 模型(比如现在的聊天机器人)中,有一种非常常见的计算任务叫**“级联归约”(Cascaded Reductions)**。

  • 比喻:想象你要做一道菜,步骤非常严格:
    1. 先切菜(计算最大值)。
    2. 把切好的菜拿去炒(减去最大值)。
    3. 把炒好的菜拿去煮(计算总和)。
    4. 最后把煮好的菜装盘(归一化)。
  • 现状:目前的 AI 编译器(相当于餐厅的“自动点餐系统”)比较笨。它会让厨师(GPU 芯片)做完第一步,把菜端回厨房放好,然后再去拿回来做第二步,做完第二步再放回去,以此类推。
  • 后果
    • 来回跑断腿:厨师在厨房(内存)和灶台(计算单元)之间来回搬运食材,浪费了大量时间。
    • 无法并行:因为第二步必须等第一步的结果,所以厨师不能同时做两件事,效率很低。

2. 解决方案:RedFuser 的“融合魔法”

RedFuser 就是为了解决这个问题而生的。它提出了一种**“融合”的魔法,把上述四个步骤合并成一个连续的、流畅的动作**。

  • 核心思想
    • 不再来回跑:RedFuser 告诉厨师:“别把菜端回去了!就在灶台上,切完直接炒,炒完直接煮,煮完直接装盘。”
    • 增量计算(Incremental Computation):这是 RedFuser 最聪明的地方。
      • 传统做法:厨师必须等所有菜都切完、炒完,才能开始下一步,这需要很大的案板(显存)来堆放所有半成品。如果菜太多,案板不够大,厨师就卡住了。
      • RedFuser 做法:它采用**“边做边更新”**的策略。就像你在切菜时,切一片就立刻处理一片,不需要等所有菜都切完。这样,无论你要做多少菜(数据量多大),厨师只需要一个小盘子(极小的显存)就能完成,而且速度极快。

3. 它是怎么工作的?(三步走)

  1. 数学侦探(理论分析)
    RedFuser 首先像数学家一样,分析这些复杂的烹饪步骤。它发现,虽然步骤看起来不同,但在数学上,它们其实可以拆解和重组。只要满足特定的数学条件(比如加法、乘法、最大值等),就可以把它们“融合”在一起。

  2. 自动翻译官(自动融合算法)
    它会自动检查你的 AI 模型代码,识别出哪些步骤是这种“级联”关系。一旦识别出来,它就自动把原本分开的代码“缝合”成一个超级高效的代码块。这就像把原本分开的四张食谱,合并成一张**“终极融合食谱”**。

  3. 金牌大厨(生成高性能代码)
    最后,RedFuser 会根据具体的硬件(比如 NVIDIA 的显卡),生成专门优化的代码。它知道怎么利用显卡的每一个小核心,怎么安排数据搬运,让硬件发挥 100% 甚至 120% 的性能。

4. 效果有多好?(实测数据)

作者们在真实的 AI 任务上测试了 RedFuser,效果非常惊人:

  • 速度提升:相比目前最先进的 AI 编译器(如 TVM、PyTorch 自带的编译器),RedFuser 让任务运行速度提升了 2 倍到 5 倍
  • 媲美手工:它生成的代码,性能几乎和由世界顶级专家手工编写的优化代码一样快。
  • 通用性强:以前,专家为了优化某个特定功能(比如注意力机制),需要专门写一段代码。现在,RedFuser 可以自动搞定各种类似的任务,包括:
    • 注意力机制(AI 理解上下文的核心)。
    • 专家混合模型(MoE)路由(决定 AI 调用哪个“专家”大脑)。
    • FP8 量化(一种让 AI 模型变小、变快的技术)。

5. 总结:为什么这很重要?

想象一下,如果 RedFuser 普及了:

  • 对于开发者:不再需要花费几个月时间去手动优化代码,编译器自动帮你把菜做得更快。
  • 对于用户:AI 聊天机器人反应更快,更流畅,而且因为计算效率高了,电费也省了
  • 对于未来:它让那些原本因为太慢、太占内存而无法运行的复杂 AI 模型,现在可以在普通的显卡上跑起来了。

一句话总结
RedFuser 就像是一个AI 领域的“流水线优化大师”,它通过把原本割裂、低效的“切菜 - 炒菜 - 煮菜”过程,变成一气呵成的“融合烹饪”,让 AI 模型跑得飞快,既省资源又高效。