AI算力反碎片化：世界上最快的统一矩阵乘法

LLVM之父Chris Lattner创办的Modular团队一直在布道AI系统和工具的碎片化带来的负面影响，他们认为，这种现状不光抬高了AI开发人员的开发难度以及开发成本，还抑制了技术创新速度。

此前，他们撰写了《模块化设计决定AI前途，不服来辩》一文讨论了AI软件模块化的重要性，而《为什么我们要重建AI基础设施软件》一文则分享了他们对重建AI基础设施软件解决方案的思考。近期，他们发布了一些实际工作成果，包括一个统一推理引擎以及面向AI开发者的Mojo编程语言。本文则介绍了他们开发的世界上最快的跨多个CPU的矩阵乘算法。

（以下内容由OneFlow编译发布，转载请联系OneFlow获得授权。原文：https://www.modular.com/blog/the-worlds-fastest-unified-matrix-multiplication）

作者 | Modular团队（Abdul Dakkak、Chad Jarvis、Eric Johnson、Hengjie Wang、AI Performance Engineer）

OneFlow编译

翻译 | 徐佳渝、贾川、杨婷

1 

Matmul，AI性能的缩影

在之前的文章《AI算力碎片化：矩阵乘法的启示》中，Modular团队说明了为什么AI需要解决计算碎片化问题以发挥全部潜力，以及通过矩阵乘法（matmul）举例说明了为什么这仍是一个未解决的问题。本篇文章将探讨Modular解决这个问题的方法及其具有的颠覆性好处，包括与现有解决方案相比，在CPU上实现SOTA性能的新标准。

在深入介绍这种算法之前，先回顾一下已有实现存在哪些问题，以及为什么构建一个通用解决方案非常困难。如今的AI行业正受到硬件性能和内存容量的限制，其结果是，硬件供应商开发了大量多样化的并行硬件架构，并支持高度优化的kernel库。

AI开发人员面临的问题是：这些kernel库是单体性质（monolithic）的“单点解决方案”，每个解决方案仅支持一小部分的行业硬件和用例。它们通常用汇编语言编写以最大限度地提高性能，结果却牺牲了可组合性、可编程性（hackability）和跨多种硬件架构的可移植性。由于需要针对特定形状和数据类型进行定制，它们的代码量都很庞大。

2

新颖方法

Modular利用多年来构建AI基础设施、服务数十亿用户的经验，发明了一种新颖的方法，以解决整个行业面临的问题。为此，我们对技术栈进行了完整的第一性原理再思考，构建了一些在行业中真正具有差异化和独特性的东西。

我们没有采用传统的方法，即编写硬编码kernels或matmul编译器，而是构建了一种更通用、可扩展的技术，将两种方法的最佳特性结合起来。这种技术使kernel开发人员能够快速开发跨形状、布局、数据类型和硬件架构的高性能kernel。本文重点讨论我们所提出方法的优点和贡献。（关于这项技术的细节可参考：https://www.modular.com/blog/a-unified-extensible-platform-to-superpower-your-ai）

统一始于单一真实信息源

如果深入研究OneDNN等库的源代码，你会发现matmul的许多实现——每个实现都是硬编码的，专门用于不同的用例，也会发现不同的数据类型（FP16、F32、F64、Int8）、各种内存布局(转置或非转置)、特殊宽高比(正方形或高瘦)、各种指令集特征等等。

由于代码太多，这些碎片化的单点解决方案让工程师很难改进库，以涵盖所有可能的用例。这会给用户带来麻烦，因为这些库会占用大量磁盘空间，增大容器和发行版（distributions）的大小，比如OneDNN占用51MB，MKL占用14MB，cuBLAS占用150MB。

Modular将通常有很多特定硬件实现的部分整合到一个“单一真实信息源（Single Source of Truth）”中。因此，专业的kernel开发者可以在不同的架构和用例中构建一个可组合、可扩展、可移植的单一代码库。这种方法能快速地重复利用代码和模式，并且适用于问题的优化子变量，并在特殊情况下轻松采用异构硬件功能。矩阵乘法的模块化实现通常不超过100kb的机器代码，这使得它在包括移动端、Web和IoT在内的许多不同用例中都可以得到实际应用。

性能可移植性

正如我们在上一篇文章中讨论的那样，为任何单个芯片实现高性能matmul都具有挑战性。然而，由于人工智能行业采用了异构硬件，包括各种类型的CPU、GPU、TPU等等，这一挑战变得更加复杂。然而，现在的kernel库仅本地支持非常有限的目标体系结构。

例如，虽然OneDNN支持ARM核心，但它是作为对ARM自己的硬件专用软件库ARM Compute Library（ACL）的封装器实现。与此同时，OneDNN并未针对AMD进行优化，AMD有自己的OneDNN分支ZenDNN，利用的是AMD的高级计算库（AOCL）。移动端则是另一个难题，常常使用像Google的Ruy这样的库。

所有这些专门定制的库给AI框架和最终用户带来了很大麻烦。框架被不同的变体和分支碎片化，用户必须混用不同版本，在各种问题中权衡。这可能会导致理论性能和实际性能之间存在巨大差距，因为用户通常不了解（或不愿意了解）这个软件层次的问题。

Modular热爱所有硬件，并且我们的方法具有广泛适用性，适用于许多种不同的架构。这使得我们能够提供一个统一的解决方案，从而解决kernel之上的框架软件碎片化的问题。这也使得我们能够更快地为新的硬件类型实现高性能支持，而且只需要一个相对较小的工程团队和更少的成本。

动态特性

一些系统采用由高级编译器（包括Google的XLA、Apache TVM和OneDNN Graph）即时编译（JIT，just in time）或预先编译（AOT，ahead of time）的kernel，它们生成的kernel专用于特定的矩阵大小，这减小了分发的代码大小，但仍需要在执行时存有许多kernel。其他库如MKL在其kernel库针对流行模型使用的特定矩阵大小进行特殊处理。

然而，随着动态语言模型如BERT（以及无数的大型语言模型、分割模型、目标检测等）的兴起，这些挑战变得更加严峻。这些模型需要处理任意大小的输入（文本或图像），但挑战在于，系统只能在推理时了解输入大小，而无法在模型训练或编译时知晓。基于静态形状的系统需要对输入进行填充（padding）或替换许多专用于不同大小的模型版本，通常这会导致性能下降、代码大小增加和模型管理困难。一些系统尝试通过JIT代码生成来解决这个问题，但这会引入不可预测的长尾延迟问题。

Modular的方法完全支持动态特性，从而彻底消除了上述问题。Modular矩阵乘法和其他核心功能可以完全支持动态形状，避免了需要JIT或AOT专门化的问题，并支持许多现有系统难以处理的其他形式的动态特性，如非常规控制流、不寻常的数据类型等。这使得整个系统更简单、更可预测。

可组合性

在神经网络中，矩阵乘法通常不会单独执行，而在其之前和之后会有其他操作，例如激活函数或逐元素元算（elementwise operations）。一般而言，通过将这些运算的代码与矩阵乘法进行“融合（fusing）“，可以显著提高性能，改善内存局部性（memory locality）并减少调度的额外开销。为了解决这一问题，通常采用两种路径：提供有限数量的预融合特殊案例，或提供一个特定领域的编译器进行融合。

第一种是最广为人知的方法，被TensorFlow、PyTorch和其他许多专业框架（如 ONNXRuntime、TFLite、TensorRT 等）所广泛采用。这种方法非常强大且灵活，非研究人员和领域专家也可以扩展系统。但是，这种方法的挑战在于运算的组合数量非常多（这也是TensorFlow和PyTorch拥有成千上万个kernel的原因之一），手动融合（hand-fusing）这些运算会进一步加剧上面提到的代码大小和可维护性问题。

第二种方法则在权衡空间（tradeoffs space）方面提供了不同的视角。像OneDNN graph、XLA和NVFuser这样的AI编译器提供了各种内核融合（kernel fusions）的选择，而无需为每种情况进行特殊处理。但是，这种方法的缺点在于，它们让你从一小部分固定的运算符集合中选择，缺乏扩展性。此外，尽管新的融合可以提供巨大的好处，但这些编译器通常无法达到传统人工编写的融合kernel库的性能水平。

Modular的方法提供了两大优势：一是支持与各种运算符的广义融合，而无需手动编写和维护（kernel）变体；更重要的是，Modular的方法具有通用性和可扩展性，无需重新编译系统，也无需用户成为编译器工程师。我们认为，这将为可能不了解编译器内部原理的专家提供重要的新研究方向和应用。

3

无与伦比的性能

虽然灵活性、通用性和易用性听起来很不错，但如果以牺牲性能为代价，那就毫无意义。性能成本直接影响运营成本，所有企业都希望自身能变得更加高效。接下来将分享我们在CPU上的早期测试结果，这只是Modular系统的开端，我们还有很多工作要做。

我们在亚马逊网络服务平台（AWS）上对多个可比较系统进行分析：Intel Skylake (c5.4xlarge)、AMD Zen-2 (c5a.4xlarge) 和 Amazon Graviton 2 (c6g.4xlarge)。它们使用了两种完全不同的指令集：Intel和AMD是X86-64，Graviton是ARM Aarch64；三种不同的向量设计：AVX-512、AVX2和NEON；还有三种不同的向量长度：512位、256位和128位。

我们使用了在各个系统上最知名的SOTA库以比较Modular方法的性能表现，如Intel上的MKL和OneDNN，AMD上的AOCL（也是ZenDNN的底层库），ARM上的ACL和Ruy，还有Eigen库，因为Eigen库是一个被广泛使用的kernel库，已经被移植到了许多体系结构。在写这篇文章时，我们使用的都是这些库的最新版本，具体如下：MKL v2023.1.0、OneDNN v2023.1.0、Eigen v3.4、AOCL v4.0、ACL v23.02.1和Ruy（从主分支#363f252拉取）。

方法论

就评估而言，我们遵循了与Google相同的基准测试方法，即先预热（warmed up）每个基准并重复运行完2秒。对需要额外设置的库，我们会在主基准测试循环之外执行设置。为了避免干扰、提高稳定性，我们确保每个基准测试调用都得到一个冷缓存，并禁用超线程以提高基准测试稳定性。虽然Modular实现确实支持预打包，但并非所有我们正在评估的库都支持它。因此，为保持公平性，我们不会在启用预打包的情况下对我们的实现进行基准测试。我们注意到，针对支持预打包的库，我们对预打包的实现进行了基准测试，并且正在和这些库竞争。

正如我们在之前文章中所讨论的，矩阵乘法的用例非常多。在这项研究中，我们决定按照AI行业最重要的矩阵形状来衡量自己，这些形状最有可能被现有库优化。因此，我们选择了从流行AI模型中挖掘的矩阵形状，例如BERT（序列长度为128和256），GPT和DLRM。列出的形状采用MxNxK形式，其中matmul的左侧操作数大小为MxK，右侧操作数大小为KxN。形状按端到端模型执行的重要性排序。

最后，虽然有许多有趣的数据类型，如Int4、FP8和bfloat16，但我们希望保持数据的简单和可比性。当然，Modular系统可以支持任何类型的系统，但对于此分析，我们专注于FP32，在我们引用的每个实现中，FP32都得到了广泛的应用和精调。

性能结果

考虑到这一点，我们首先来看英特尔Skylake（c5.4xlarge）系统上的单线程性能。单线程性能无法利用整个芯片的资源，但有助于归一化结果（即消除多处理（multi-processing）、错误共享、NUMA问题等高阶因素）。除此之外，它还是多线程性能的基础，对移动和游戏引擎中的某些用例至关重要。

下图展示了英特尔系统（MKL、OneDNN和Eigen）的Modular方法和其他SOTA实现的性能（每秒十亿次浮点运算，GFLOP/s）性能。Modular的矩阵乘实现取得了与现有SOTA解决方案不相上下或更好的性能。事实上，事实上，对于英特尔系统，Modular的速度比OneDNN大约快1.5倍。

虽然单线程性能是有用的数据点（datapoint），但对服务器用例而言，重要的是完整的多线程性能。这给机器带来了更大的压力，因为全峰值实现可能遇到峰值FLOP/s、DRAM带宽、缓存利用率等限制。下面我们展示了Modular针对这些系统的性能——Modular比OneDNN快1.46倍，考虑到之前讨论的通用性和其他好处，这是一项了不起的成就。

虽然在一种硬件平台上的强势结果很重要，但Modular实现的基本价值主张是：单一的真实信息来源（a single source of truth）可以在多硬件上提供高性能。下面来看看与AOCL库（AMD上的SOTA，是ZenDNN的基础）和上面提到的OneDNN库相比，Modular在AMD硬件上的表现。

观察下表可以看出：OneDNN的性能在AMD硬件上表现不佳，而AOCL库虽然提供了显著的性能提升，但是Modular方法在AMD系统上比现有SOTA方案快约1.6倍。

我们还在Amazon Graviton 2系统上进行了相同实验，涵盖了Ruy和ACL库。Ruy是边缘框架（如TensorFlow Lite）使用的库，而ACL是支持ARM平台的OneDNN库的基础。

尽管ACL和Ruy在ARM上具有竞争力，但Modular实现的平均性能明显更好（比ACL好1.8倍，比Ruy好1.2倍）。

除了在给定系统上比较矩阵乘法的不同实现之外，交叉比较不同系统的绝对性能也很有趣。这些是非常复杂的机器，有很多活动部件，但我们可以用粗粒度（coarse grain）来看待事物。英特尔系统的优势在于512位长矢量，而不是较短的256或128位矢量。尽管128位矢量较短，但Graviton 2系统性能表现良好，因为它有16个物理kernel，相比之下，Intel和AMD系统只有8个物理kernel。

Kernel融合感知

最后，为证明实现的可组合性，我们来看看matmul如何与其他运算融合。为了与其它经高度精调的实现的常见运算进行比较，我们需要使用“全连接（FC）”块，由等式“activation（matmul（A，B）+bias）”定义，其中activation是一个激活函数（下面使用Relu）。

OneDNN等库已经为FC block提供了融合路径，为保证结果公平，我们不与无法支持可融合方式定义FC模块的库进行比较。因此，在此分析中，我们仅与OneDNN和Eigen进行比较，它们提供了表达融合模式的方法。为了呈现结果的一致性，我们使用与上文相同的形状和三种硬件配置。

下面是英特尔Skylake架构上的FC模块的性能。尽管Modular方法具有通用性和灵活性，但它的平均性能高出OneDNN 1.45倍，高出Eigen 1.8倍，代表着新的SOTA。

我们在AMD系统上看到了类似的强势结果，其中Modular方法的性能比OneDNN高出了2.1倍，比Eigen高出了2.3倍。

Amazon Graviton 2系统是明显不同的架构，Modular堆栈不像在X86-64上那样经过精调。尽管如此，Modular还是比Eigen高了1.3倍，比Ruy高了1.1倍。OneDNN/ACL没有为ARM系统提供融合的FC层。

从上面的数据我们可以看出：如果执行得当，kernel融合可以带来显著的性能提升。随着融合区域增长，性能提升会更加明显。Modular方法从一开始构建时就积极促进融合，这让它能够支持非常通用的融合（即远远超出逐元素运算，不受matmul的限制）。我们认为，性能可移植、动态和可组合的单一事实信息来源的实现是关键贡献，它支持新的研究和生产用例。

4

AI的未来

虽然对性能结果感到兴奋，但最重要的是我们可以在不损害最初目标的情况下实现性能结果，而这仅仅只是开始！

Modular matmul具有单一真实信息来源，支持多种不同架构、动态形状和可扩展融合。除了为用户提供“眼前”的好处之外，Modular架构的通用性还能从根本上简化上述堆栈，产生更可预测的用户体验，并以人们从未体验过的方式快速启动新硬件。

值得强调的是，现有的SOTA实现是大量优秀工程师数十年研究和开发的产物。虽然Modular团队人数相对较少，但都非常优秀，由于堆栈关键技术的进步，团队能够快速交付优质结果。更重要的是，这也是通过第一性原理重新思考和愿意真正投资Modular建立的“自下而上重建”方法的结果。

从更大的层面上说，我们希望实现AI人人可用，处处可及，我们想要发挥AI的真正影响力，利用它去完成更有意义、更有用的事。

通过构建和创建AI基础设施新方法，我们描绘了一个世界，可以帮助推动整个行业向前发展，使其更快、更高效、更安全地开发和部署人工智能系统，最终实现全世界共享AI成果。我们希望帮助硬件行业构建创新型计算架构，以软件推动硬件创新。

其他人都在看

• AI算力碎片化：矩阵乘法的启示

• 关于大型语言模型的争论和局限

• John Schulman：通往TruthGPT之路

• 为什么ChatGPT用强化学习而非监督学习

• Geoffrey Hinton：智能进化的下一个阶段

• OneEmbedding:单卡训练TB级推荐模型不是梦

• GLM训练加速：性能最高提升3倍，显存节省1/3

试用OneFlow: github.com/Oneflow-Inc/oneflow/http://github.com/Oneflow-Inc/oneflow/