脉动阵列承担了现代AI芯片超过95%的计算任务，它通过在处理单元之间本地移动矩阵数据来减少内存流量。更大的阵列能提升峰值吞吐量，但更难达到饱和，因此编译器调度常常决定了芯片能否实现有效的利用率。

AI芯片的动力核心

脉动阵列是每一款现代AI芯片的动力核心，承担了其总计算量的95%以上。没有它们，AI工作负载的速度会慢20到100倍。在这篇文章中，我们将讨论其设计原理以及为何如此高效。

GIF

一个AI模型由密集的线性代数运算构成：逐元素的激活函数、转置、散列和收集，以及最重要的矩阵乘法（matmul）。矩阵乘法占据了AI模型中绝大部分的计算量，因此它是硬件中需要高效实现的最关键操作。没有快速的矩阵乘法，即使模型中其他所有操作都瞬间完成，推理速度也会变得极其缓慢。

令人惊讶的是，尽管不同芯片的整体架构、内存系统、供电和网络连接千差万别，但驱动矩阵乘法的硬件却在市场上所有厂商中趋同于几乎完全相同的设计：脉动阵列。

这是因为矩阵乘法要求输入矩阵以几何模式相互作用。看这张Nvidia为其Tensor Core提供的示意图，它展示了这种交互几何。

两个输入矩阵显示在这个三阶张量的外部。三阶张量中的每个元素都是每个输入矩阵中一个元素的乘积。具体地说，三阶张量中的（M,N,K）元素等于（M,K）乘以（N,K）。最终的输出矩阵是对三阶张量沿K维度所有垂直元素进行求和归约的结果。这是一个熟悉的定义，但请关注可视化及其传达的关于矩阵乘法所需通信模式的几何信息。两个输入矩阵具有一对正交方向和一个共享的归约维度。

脉动阵列的类型

脉动阵列在技术上比专门用于加速矩阵乘法的特定子类型更通用。“systolic“一词源于心脏的泵血运动，脉动阵列是一个将数据泵送到所需位置的系统。加速矩阵乘法的脉动阵列是矩形网格形状，通常为正方形，反映了正交的通信模式。每个节点中都有一些存储和计算单元。节点可以与其基本方向上的最近邻居通信。脉动阵列有两种工作模式：

累加器固定

网格中的每个节点与输出矩阵中的一个元素相关联。每个节点初始化为一个累加器。输入矩阵从顶部和侧面流式输入，每个周期引入一个新的k索引。网格中的（m,n）元素每个周期接收一个新的（m,k）和（n,k）元素。它在本地执行融合乘加运算，并用累加器跟踪归约结果。在输入流结束后，最终的输出矩阵从阵列中卸载。

权重固定

网格中的每个节点与其中一个输入矩阵元素相关联。在这种情况下，另一个输入矩阵从一侧流式输入，每个周期计算部分归约结果并向下传递。这种模式需要的通信更少，但为了实用，权重需要在一个大的输入矩阵上共享。

脉动阵列是一种专用硬件单元，用于极其高效地执行矩阵乘法。名称中的“systolic“来源于医学领域，指的是数据以脉冲方式在阵列中传输，类似于心脏泵血。输入矩阵按行和列流式输入，输出矩阵迭代式流出。

元素直接从顶部和左侧流向底部和右侧，每一步都直接流向相邻的处理单元，无需进行远距离的全局数据传输。每个处理单元从左侧接收一个输入，将其与内存中固定的值相乘，然后与来自顶部的输入相加，再将结果发送给正下方的处理单元。这是一种典型的“权重固定“设置，其中权重存放在处理单元的内存中，激活值从左流入，累加值从上而下流动。另一种形式则是权重从顶部流入，累加值固定存放在处理单元内存中。

脉动阵列具有出色效率的核心原则在于最小化数据移动。由于计算严格在本地进行，阵列可以完全并行地执行大量融合乘加运算。阵列还可以通过简单地以相同模式放置更多处理单元来直接扩展，从而一次计算更大的数据块。

脉动阵列效率如此之高，以至于通常加速器的整体效率直接取决于：A) 你能将脉动阵列做到多大，以及 B) 是否能持续向这些阵列输送数据。这里存在一个权衡：更大的阵列允许更大的数据块和更多的并行计算，从而带来更高的整体效率，但灵活性较低，也更难完全饱和。较小的阵列更容易饱和，对实际中动态变化的工作负载更灵活，但数据局部性受限于较小的数据块大小，因此可达到的峰值效率较低。

实际中的权衡

我们可以在实际产品中看到这种权衡的体现。Nvidia GPU最初采用了大量较小的脉动阵列，后来转向了较少但更大的阵列。这反映了其芯片的演进：从最初的8x8，到32x32，现在达到了128x128（Nvidia将他们的脉动阵列称为“TensorCore“）。另一方面，Google的TPU选择了更少但更大的脉动阵列，范围从TPU v2的128x128到TPU v6和v7的256x256。

这种差异显示了更深层次的理念分歧：TPU更依赖于提前编译调度来布局数据并安排所需的数据移动和指令，以使脉动阵列完全饱和，如果饱和，可以实现非常高的吞吐量。而GPU则较少依赖提前调度，允许内核在运行时调度数据移动，以使相对较小的阵列饱和。

在GPU方面，我们看到随着时间的推移，其理念逐渐向TPU靠拢。较老的GPU包含较小的（8x8）脉动阵列，数据始终来自完全动态可编程的CUDA核心。然而，在较新的几代产品中，我们看到阵列（TensorCore）逐渐变大（Blackwell上的128x128）。还引入了其他特殊功能单元，如TMA，以减轻通用CUDA核心的压力。我们很可能会在Rubin和Feynman上继续看到这一趋势：更少、更大的核心以及更多的特殊功能单元，如果通过复杂的软件正确驱动，可以实现更高的吞吐量。

系统软件为何重要

协调所有这些元素协同工作的软件，意味着10%与80%利用率之间的差距。调度不当的数据移动或不必要的同步可能会导致一切停滞。

这就是为什么Luminal编译器如此专注于最小化数据移动并在指令调度中最大化重叠。通过以流水线方式操作每个硬件单元，我们可以使芯片的每个部分尽可能独立运行，并最大限度地减少空闲时间。

很多时候，人工编写或代理编写的内核会出现“气泡“现象，导致脉动阵列闲置等待数据。大规模搜索使Luminal能够尝试许多不同的指令顺序和重叠方式，以找到能持续向这些“饥饿“的阵列输送数据的调度方案。

深入理解运行这些模型的硬件，使我们能够精准攻克正确的瓶颈，实现光速般的性能！