AI芯片为何偏爱脉动阵列:从矩阵乘法到TPU的架构解析
发布时间:2026/10/9 0:30:22 锦皓数字建站

1. 从矩阵乘法说起为什么AI芯片偏偏看上了脉动阵列很多人第一次听到“脉动阵列”这个词脑子里浮现的是一排排整齐的运算单元像心跳一样同步跳动的画面。这个直觉其实相当准确。要理解它为什么在AI芯片领域被反复提起得先回到一个最朴素的问题神经网络到底在算什么答案出奇地简单——绝大部分计算量都压在矩阵乘法和卷积上。一个典型的Transformer层里面全是Q、K、V的矩阵乘法一个卷积层本质上是把卷积核展开成矩阵做乘法。换句话说谁能把矩阵乘法做得又快又省电谁就抓住了AI芯片的命脉。传统的CPU做矩阵乘法是什么路子它从内存里取一个数送到运算单元算完写回内存再取下一个数。这个过程里运算单元大部分时间在等数据。就像一个厨师每炒一道菜都要亲自跑去菜市场买一次菜灶台大部分时间是冷的。GPU稍微好一点它有很多个厨师同时炒菜但买菜这件事还是各自跑各自的数据搬运的能耗远远超过计算本身的能耗。脉动阵列的思路完全不同。它把运算单元排成一个二维网格数据像血液一样从边缘流入在网格中一格一格地“脉动”前进每经过一个运算单元就完成一次乘加操作。数据被复用了不需要反复从内存里搬运。这个设计的精妙之处在于数据流动本身就是计算过程搬运和计算合二为一。我第一次在论文里看到这个结构时觉得它像一条工厂流水线——原料从一端进去每经过一个工位就被加工一次最后从另一端出来就是成品。而不是每个工位都派人去仓库领料。这个类比帮助我理解了为什么脉动阵列在能效比上有天然优势它把数据复用做到了极致。那么为什么偏偏是AI芯片把脉动阵列推到了聚光灯下因为AI计算有一个极其鲜明的特征计算量大但数据模式高度规整。矩阵乘法中每个输出元素都需要一行和一列的内积这种规整的依赖关系天然适合用规则的数据流来驱动。而脉动阵列恰好就是为这种规整计算量身定做的。相比之下通用计算中大量存在分支跳转、不规则访存脉动阵列反而施展不开。这里需要点明一个常见的误解脉动阵列并不是什么新发明。它的概念可以追溯到上世纪七八十年代H.T. Kung等人在CMU做的一系列工作。但那个年代它没有成为主流因为通用计算的需求不匹配。直到深度学习兴起矩阵乘法成为绝对主角脉动阵列才真正找到了自己的主场。Google的TPU v1是把这个结构推向大规模商用的标志性产品之后几乎所有做AI加速的公司都在自己的架构里或多或少借鉴了脉动阵列的思想。所以当你看到“AI芯片脉动阵列”这个标题时它背后其实是一条清晰的逻辑链AI计算以矩阵乘法为核心矩阵乘法需要高数据复用率脉动阵列提供了极致的数据复用因此成为AI芯片的主流计算架构之一。理解这条链比记住“脉动阵列是什么”重要得多。2. 拆开一个脉动阵列数据怎么流计算怎么叠2.1 一个最简单的4x4脉动阵列长什么样要真正搞懂脉动阵列光看文字描述不够得在脑子里把它“跑”起来。我们拿一个最经典的4x4权重固定型脉动阵列做例子。假设我们要计算矩阵A4x4乘以矩阵B4x4得到C4x4。在权重固定型设计里矩阵B的每一列被预先加载到阵列的每一列运算单元中作为固定权重。矩阵A的元素从左侧流入每次流入一行逐拍向右脉动。每个运算单元做一件事把从左边进来的数据和本地存储的权重相乘加到从上方进来的部分和上然后把新的部分和往下传把左边进来的数据往右传。用更具体的方式说第0拍A[0][0]进入左上角单元PE(0,0)和B[0][0]相乘得到部分和。第1拍A[0][1]进入PE(0,1)同时A[0][0]从PE(0,0)脉动到PE(0,1)PE(0,1)做A[0][0]*B[1][0]并累加。以此类推数据像波浪一样在阵列中推进。这个过程的关键在于每个运算单元在每个时钟周期都在做有效的乘加没有空闲周期。数据到达每个单元的时间经过精心设计使得所有部分和恰好在正确的时间汇聚到正确的位置。这就是“脉动”的含义——数据有节奏地流动计算有节奏地发生。2.2 权重固定、输出固定、行固定三种数据流的取舍脉动阵列不是只有一种形态。根据哪个矩阵被固定在运算单元中可以分为权重固定型、输出固定型和行固定型。这三种形态各有各的适用场景选错了会让效率大打折扣。权重固定型是最常见的。权重提前加载到每个PE中输入数据从左侧和上方流入。它的优势是权重加载一次可以复用很多次适合卷积神经网络中卷积核固定的场景。但缺点是如果权重矩阵很大需要分块加载块与块之间的切换有开销。输出固定型则是把输出矩阵的部分和固定在PE中输入和权重都流动。这种设计适合输出矩阵较小但输入矩阵很大的场景比如全连接层中batch size很大的情况。部分和不需要在PE之间传递减少了数据搬运。行固定型介于两者之间把输入矩阵的一行固定在PE中权重和输出流动。这种设计在某些特定的矩阵维度下能获得更好的利用率。我在实际评估芯片架构时发现很多论文只讲权重固定型但真正做产品时选择哪种数据流取决于你的目标网络结构。如果你的芯片主要跑卷积权重固定型通常是最优解如果主要跑全连接或注意力机制可能需要考虑输出固定型或者混合方案。没有一种数据流是万能的。2.3 为什么脉动阵列的能效比能做到这么高脉动阵列最被人称道的就是能效比。同样做一次矩阵乘法脉动阵列的能耗可能只有传统架构的几分之一甚至十分之一。这个优势从何而来核心在于数据复用。在一个NxN的脉动阵列中每个输入数据元素被复用了N次每个权重被复用了N次。而在传统的乘加器中每个数据元素从内存读出后只用一次就丢弃了。内存访问的能耗通常是计算能耗的几十倍甚至上百倍所以减少内存访问次数就是减少能耗的最直接手段。具体算一笔账假设一个4x4的脉动阵列计算4x4矩阵乘法。传统方式需要从内存读取4x44x432个数据元素做64次乘加。脉动阵列方式输入数据从边缘流入每个元素只从内存读取一次总共读取8个元素4个输入行4个权重列做64次乘加。内存访问次数从32次降到8次减少了75%。如果阵列扩大到16x16这个比例会更夸张。另一个能耗优势来自局部互联。脉动阵列中数据只在相邻PE之间传递连线短电容小驱动能耗低。相比之下如果用一个大的交叉开关连接所有PE长距离连线的能耗会迅速上升。这也是为什么脉动阵列的扩展性相对较好——增加阵列规模时互联的复杂度是线性增长的而不是平方级增长。不过这里有一个容易被忽略的代价脉动阵列的延迟。数据从阵列边缘流到中心需要多个时钟周期对于小规模矩阵乘法这个启动延迟可能比计算本身还长。所以脉动阵列适合大批量、大矩阵的场景对于小矩阵或者延迟敏感的应用反而不一定划算。3. 从TPU到边缘端脉动阵列在真实芯片里的样子3.1 Google TPU v1脉动阵列的成名之战2016年Google公开TPU v1的架构时整个行业都注意到了那个256x256的脉动阵列。这个阵列在700MHz频率下能提供92 TOPS的峰值算力而功耗只有40W左右。这个能效比在当时是碾压性的。TPU v1的脉动阵列是典型的权重固定型。权重从片外DRAM加载到阵列的每一列输入数据从左侧流入。整个数据流经过精心编排使得大部分时间阵列的利用率都在80%以上。Google在论文里特别强调了一点他们不需要做复杂的调度因为矩阵乘法的数据依赖是确定的编译器可以静态地安排好所有数据的流动。但TPU v1也有它的局限。它主要面向推理不支持训练。训练需要反向传播数据流方向会变化权重也需要频繁更新这对脉动阵列的静态数据流设计提出了挑战。后来的TPU v2/v3增加了对训练的支持架构上做了不少调整比如增加了更多的片上存储和更灵活的数据路由。从TPU v1的经验里我总结出一个关键点脉动阵列的效率高度依赖编译器的数据编排能力。硬件只是提供了规则的计算网格怎么把神经网络的计算图映射到这个网格上怎么安排数据的流入流出顺序这些都需要编译器在编译时解决。TPU的编译器XLA在这方面做了大量工作这也是为什么TPU的实际性能往往能接近峰值而一些架构类似但编译器跟不上的芯片实际利用率可能只有30%到40%。3.2 边缘AI芯片为什么也爱用脉动阵列很多人以为脉动阵列只适合数据中心的大芯片其实边缘端同样有它的用武之地。边缘AI芯片的功耗预算通常在几瓦甚至几百毫瓦对能效比的要求比数据中心更苛刻。脉动阵列的高数据复用率正好切中这个需求。不过边缘端的脉动阵列规模通常小得多可能是8x8、16x16或者32x32。规模小了之后启动延迟和填充/排空开销的占比会上升。所以边缘端的脉动阵列设计往往需要做一些优化比如支持可变阵列规模小矩阵用小阵列跑大矩阵用大阵列跑或者采用深度可分离卷积专用的数据流减少不必要的计算。我见过一些边缘芯片的设计把脉动阵列和SIMD向量单元结合在一起。卷积层用脉动阵列跑全连接层或者一些非规则计算用向量单元跑。这种异构设计在边缘端很常见因为边缘端的负载更加多样化单一架构很难覆盖所有场景。另一个边缘端特有的问题是权重存储。边缘芯片的片上SRAM通常只有几百KB到几MB而一个稍微大点的模型权重就有几十MB。所以边缘端的脉动阵列往往需要配合权重压缩技术比如量化到INT8甚至INT4或者用稀疏化技术减少有效权重数量。这些技术在脉动阵列上的实现方式会直接影响最终的能效比。3.3 脉动阵列不是万能药它搞不定的那些计算说了这么多脉动阵列的好处也得说说它的边界。脉动阵列最擅长的是密集的、规则的矩阵乘法。一旦计算模式偏离这个前提它的效率就会急剧下降。比如稀疏矩阵乘法。如果权重矩阵中有大量零脉动阵列仍然会为这些零分配计算周期因为它的数据流是静态编排的不会因为某个权重是零就跳过。虽然有一些研究在做稀疏脉动阵列比如通过门控时钟关掉不活跃的PE或者动态跳过零权重但这些方案都会增加控制逻辑的复杂度抵消一部分能效优势。再比如非矩阵乘法的操作。激活函数、归一化、池化这些操作在脉动阵列上跑效率很低通常需要专门的单元来处理。所以实际的AI芯片往往是脉动阵列加上其他专用单元的异构架构脉动阵列只负责它最擅长的部分。还有一个容易被忽视的问题是灵活性。脉动阵列的数据流是相对固定的一旦芯片流片能支持的计算模式就确定了。如果几年后出现了新的网络结构数据流模式发生了根本变化脉动阵列可能就无法高效支持。这也是为什么一些公司在设计AI芯片时会保留一定的可重构性让脉动阵列的互联方式可以根据不同的计算模式重新配置。4. 自己动手用Python模拟一个脉动阵列的完整过程4.1 为什么值得亲手模拟一遍看再多论文和架构图不如自己写一遍模拟代码。脉动阵列的很多细节——数据到达时间、部分和的累加顺序、填充和排空的过程——只有亲手实现过才能真正理解。而且模拟代码可以帮助你验证自己对数据流的理解是否正确比如你可以故意改变数据流入的顺序看看输出会变成什么样。下面我用Python写一个权重固定型脉动阵列的模拟。代码不追求性能只追求逻辑清晰方便你对照着理解每一拍发生了什么。4.2 核心模拟代码与逐拍解析import numpy as np def systolic_array_weight_stationary(A, B): 模拟权重固定型脉动阵列计算 C A B A: 输入矩阵 (M, K) B: 权重矩阵 (K, N) 返回: 输出矩阵 C (M, N) M, K A.shape K2, N B.shape assert K K2, 矩阵维度不匹配 # 初始化输出矩阵和部分和累加器 C np.zeros((M, N)) # 每个PE的部分和累加器形状 (M, N) psum np.zeros((M, N)) # 权重固定在每个PE中PE[i][j] 存储 B[i][j] # 输入数据从左侧流入第i行数据从第i行左侧进入 # 数据在阵列中向右脉动部分和向下脉动 # 总拍数 M N K - 2考虑填充和排空 total_cycles M N K - 2 for cycle in range(total_cycles): # 计算每个PE在当前拍的操作 # PE(i, j) 在 cycle 时刻接收来自左侧的 A 元素 # A 元素进入 PE(i, j) 的时间cycle - j # 该元素对应的 k 索引cycle - j - i for i in range(M): for j in range(N): k cycle - j - i if 0 k K: # 当前PE做乘加 psum[i][j] A[i][k] * B[k][j] # 检查是否有输出完成 # PE(i, j) 的输出在 cycle i j K - 1 时完成 for i in range(M): for j in range(N): if cycle i j K - 1: C[i][j] psum[i][j] return C # 测试 A np.array([[1, 2, 3], [4, 5, 6]]) B np.array([[7, 8], [9, 10], [11, 12]]) C_systolic systolic_array_weight_stationary(A, B) C_numpy A B print(脉动阵列结果:) print(C_systolic) print(NumPy结果:) print(C_numpy) print(是否一致:, np.allclose(C_systolic, C_numpy))这段代码的核心逻辑在于那个k cycle - j - i的公式。它决定了在哪个时钟周期哪个数据元素到达哪个PE。这个公式不是随便写的它来自脉动阵列的数据流设计数据从左侧进入后每经过一个PE就延迟一拍部分和从上方进入后每经过一个PE也延迟一拍。所以PE(i, j)在cycle时刻处理的元素其k索引正好是cycle减去它在行方向和列方向上的延迟。运行这段代码你会看到输出和NumPy的矩阵乘法完全一致。但更重要的是你可以在循环里加打印语句观察每一拍每个PE在做什么。比如在cycle0时只有PE(0,0)在做A[0][0]*B[0][0]cycle1时PE(0,0)在做A[0][1]*B[1][0]PE(0,1)在做A[0][0]*B[0][1]PE(1,0)在做A[1][0]*B[0][0]。数据像波浪一样在阵列中推进每个PE在每个时刻都在做有效的乘加。4.3 模拟中暴露出的三个关键细节亲手跑过这段代码后有三个细节会变得非常清晰而这三个细节在纯看论文时很容易被忽略。第一个细节是填充和排空的开销。从代码里可以看到总拍数是MNK-2而实际做有效计算的拍数只有MNK次乘加。当M、N、K都很小的时候填充和排空的周期占比很高。比如MNK4时总拍数是10但有效计算分布在不同的PE上平均利用率可能只有60%左右。这就是为什么脉动阵列需要大矩阵才能发挥优势。第二个细节是部分和的累加顺序。在代码中psum[i][j]是按照k从0到K-1的顺序累加的。这个顺序是由数据到达的时间决定的不能随意改变。如果改变数据流入的顺序部分和的累加顺序就会变化虽然最终结果在数学上一样但在浮点运算中可能会因为舍入误差导致细微差别。在实际芯片设计中这个顺序会影响数值稳定性。第三个细节是权重的加载时机。代码中假设权重已经预先加载到每个PE中。但在实际芯片中权重加载需要时间而且如果权重矩阵很大需要分块加载块与块之间的切换会打断计算流水线。这个开销在模拟中看不出来但在实际芯片的性能分析中必须考虑。5. 设计脉动阵列时最容易踩的五个坑5.1 阵列规模不是越大越好很多刚接触脉动阵列的人会有一个直觉阵列越大算力越强能效比越高。这个直觉在某种程度上是对的但有一个重要的前提被忽略了——你的矩阵维度是否匹配阵列规模。假设你设计了一个256x256的脉动阵列但你的目标网络中最常见的矩阵乘法是64x64的。那么每次计算时阵列只有1/16的PE在工作其余15/16的PE要么空闲要么在做无效计算。利用率只有6.25%能效比的优势荡然无存。更糟糕的是大阵列的填充和排空延迟更长。256x256的阵列数据从边缘流到中心需要128拍对于小矩阵来说这128拍的延迟可能比计算本身还长。所以阵列规模的选择必须基于目标工作负载的矩阵维度分布。如果目标网络的矩阵维度变化很大可以考虑支持可变阵列规模或者用多个小阵列组合成大阵列。TPU的做法是固定256x256因为Google的目标负载是自家的大模型矩阵维度足够大。但如果你做的是边缘芯片目标负载可能是各种不同维度的卷积层那就需要更灵活的设计。5.2 数据位宽的选择直接影响面积和功耗脉动阵列中每个PE都需要乘法器和加法器。乘法器的面积和功耗与数据位宽呈平方关系。FP32的乘法器面积大约是INT8乘法器的16倍功耗差距更大。在AI推理场景中INT8通常已经足够精度损失在可接受范围内。所以大多数AI芯片的脉动阵列都采用INT8或混合精度设计。但这里有一个坑累加器的位宽必须足够大。即使输入是INT8累加结果可能超出INT8的范围。通常累加器需要INT32否则会溢出。另一个坑是权重的动态范围。如果权重分布不均匀量化到INT8后可能会有较大的精度损失。这时候需要考虑per-channel量化即每个通道用不同的缩放因子。这个技术在脉动阵列上的实现需要在权重加载时做额外的处理增加了控制逻辑的复杂度。5.3 片上存储的带宽必须匹配阵列的吞吐脉动阵列的算力很容易算PE数量乘以频率乘以每PE每拍的操作数。但算力上去了数据供应能不能跟上这是很多设计失败的原因。一个NxN的脉动阵列每个时钟周期需要从左侧流入N个数据元素从上方流入N个权重元素如果是权重固定型权重只在加载阶段需要。如果频率是1GHzN256那么输入带宽需要256GB/s假设每个元素1字节。这个带宽对于片上SRAM来说是可以做到的但如果数据需要从DRAM读取带宽就可能成为瓶颈。所以脉动阵列的设计必须和存储层次一起考虑。通常的做法是在阵列旁边放一块大的SRAM作为数据缓冲区DRAM的数据先搬到SRAM再从SRAM流入阵列。SRAM的带宽要足够高延迟要足够低。TPU v1的SRAM带宽是34GB/s配合256x256的阵列在700MHz下运行刚好匹配。5.4 控制逻辑的复杂度容易被低估脉动阵列看起来规则但控制逻辑并不简单。数据什么时候流入、什么时候切换权重、什么时候读出结果这些都需要精确的时序控制。而且不同的矩阵维度需要不同的控制序列如果支持动态维度控制逻辑会更复杂。一个常见的坑是地址生成单元的设计。数据在SRAM中的存储方式会影响地址生成的复杂度。如果数据按行存储但脉动阵列需要按列流入就需要转置操作这会增加延迟和功耗。所以在设计存储布局时就要考虑脉动阵列的数据流方向尽量让数据的自然存储顺序和流入顺序一致。另一个坑是边界处理。当矩阵维度不是阵列规模的整数倍时需要填充零或者做边界裁剪。填充零会浪费计算周期边界裁剪会增加控制复杂度。通常的做法是填充零因为实现简单但代价是利用率下降。5.5 验证和测试的难度远超预期脉动阵列的验证是一个大工程。因为数据在阵列中流动任何一个PE的故障都可能导致输出错误而且错误可能只在特定的数据模式下才出现。所以测试向量需要覆盖各种数据模式包括全零、全一、随机数、边界值等。更麻烦的是时序验证。脉动阵列的时序非常紧凑建立时间和保持时间的余量很小。在先进工艺节点下线延迟和时钟偏斜的影响很大需要仔细做时序收敛。我见过一些设计在RTL仿真时功能完全正确但流片后因为时序问题导致数据在PE之间传递时出错。所以脉动阵列的设计需要从架构阶段就考虑可测试性比如在PE中插入扫描链在关键路径上插入观测点。这些在初期看起来是额外的开销但流片后调试时会感谢自己当初做了这些。6. 脉动阵列之后它会被什么取代或者如何进化6.1 脉动阵列面临的三个根本性挑战脉动阵列在AI芯片领域的地位目前还很稳固但它面临的挑战也在增加。第一个挑战是稀疏化。越来越多的模型采用稀疏化技术来减少计算量比如剪枝后的模型可能有70%以上的权重是零。脉动阵列对稀疏计算的支持天然不足因为它的数据流是静态的不会因为权重是零就跳过计算。第二个挑战是动态性。新的网络结构比如动态路由、条件计算、MoE混合专家计算模式不再是静态的矩阵乘法而是根据输入动态变化的。脉动阵列的静态数据流很难高效支持这种动态性。第三个挑战是存内计算。一些新兴的存内计算技术比如基于RRAM或SRAM的存内计算把计算直接放在存储阵列里做彻底消除了数据搬运的开销。虽然目前存内计算在精度和可靠性上还有不少问题但长期来看可能对脉动阵列构成威胁。6.2 可重构脉动阵列在灵活性和效率之间找平衡面对这些挑战一个明显的进化方向是可重构脉动阵列。基本思路是在保留脉动阵列高数据复用率的同时增加互联的可配置性让阵列可以适应不同的数据流模式。比如可以在PE之间增加可配置的旁路通路让数据可以选择跳过某些PE从而支持稀疏计算。或者把大阵列划分成多个小阵列每个小阵列可以独立配置数据流支持不同的计算模式。还有一些设计把脉动阵列和NoC片上网络结合让数据可以通过NoC在阵列之间灵活路由。可重构的代价是面积和功耗的增加。每个可配置的开关都会增加面积动态配置也会增加功耗。所以关键是在灵活性和效率之间找到平衡点。我的经验是如果目标负载相对固定不需要过度追求可重构性如果目标负载变化很大那可重构带来的收益可能超过它的开销。6.3 从架构层面看脉动阵列的未来位置脉动阵列不会消失但它的角色可能会变化。在未来的AI芯片中脉动阵列可能不再是唯一的计算核心而是异构计算系统中的一个组成部分。它负责处理密集的、规则的矩阵乘法而其他类型的计算由其他专用单元处理。比如一个典型的未来AI芯片可能包含一个脉动阵列用于密集矩阵乘法一个向量单元用于元素级操作一个稀疏计算单元用于稀疏矩阵一个存内计算单元用于低精度的大规模并行计算。这些单元通过片上网络连接由编译器根据计算图自动分配任务。这种异构架构的挑战在于编程模型和编译器。程序员不需要知道每个计算在哪个单元上跑编译器需要自动完成映射和调度。这比单一的脉动阵列架构复杂得多但也是发挥各种架构优势的唯一途径。从更宏观的视角看脉动阵列代表了一种设计哲学用规则的计算结构匹配规则的计算模式用数据流动代替数据搬运。这个哲学不会过时即使具体的架构形态会变化。理解脉动阵列不仅是理解一个具体的硬件结构更是理解一种在能效约束下做计算架构设计的方法论。这个方法论在AI计算需求持续爆发的背景下只会越来越重要。我在实际项目中最大的体会是不要为了用脉动阵列而用脉动阵列。先分析你的目标负载看看矩阵乘法的占比有多大矩阵维度分布是什么样的稀疏性如何然后再决定脉动阵列的规模、数据流形态和精度配置。架构设计没有银弹只有针对具体场景的最优解。脉动阵列是一个强大的工具但工具的价值取决于使用它的人对问题的理解深度。
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。