脉动阵列原理与AI芯片设计:从面试翻车到架构实战 1. 从一次面试翻车聊起为什么脉动阵列值得单独拎出来讲几年前我面一家芯片公司的架构岗面试官问我“给你一个 256x256 的矩阵乘算力峰值 100 TOPS你怎么把数据喂进去”我当时脑子里全是 GPU 那套 SIMT 的思路张口就是“多级缓存加高带宽总线”。面试官笑了笑说“那你有没有想过数据根本不用来回搬”那次之后我才认真去啃了脉动阵列Systolic Array这个东西越啃越觉得它才是理解 AI 芯片算力本质的一把钥匙。脉动阵列这个词听起来玄乎其实核心思想特别朴素让数据像心跳一样有节奏地流过一排排计算单元每个单元只做一件极小的事算完就把结果往下一个单元推。它解决的是 AI 芯片里最要命的一个矛盾——算力上去了但数据搬运跟不上。传统处理器里乘加运算本身很快真正拖后腿的是把数据从内存搬到计算单元、再把结果搬回去这个过程功耗和延迟大头全在这儿。脉动阵列通过让数据在计算单元之间“接力传递”把重复的数据复用做到极致从而把访存次数压到最低。这篇内容适合谁看如果你是做 AI 芯片架构、数字 IC 设计、或者正在准备相关岗位面试的人那脉动阵列是你绕不过去的基本功。如果你只是对“AI 芯片为什么快”这件事好奇想搞明白 TPU 里那块矩阵乘法单元到底怎么工作的这篇也能让你从零建立起直观认识。我会从设计思路、核心原理、参数计算、实操建模、常见坑几个角度把脉动阵列掰开揉碎讲清楚尽量不堆公式多用类比和具体数字。2. 脉动阵列到底在解决什么问题设计思路与方案选型2.1 先搞清楚 AI 计算的真实瓶颈在哪要理解脉动阵列得先理解 AI 计算的特点。深度学习里最核心的运算是矩阵乘法卷积可以展开成矩阵乘全连接本身就是矩阵乘注意力机制里一堆 QKV 也是矩阵乘。矩阵乘的特点是计算量大但数据复用率极高。比如一个 MxK 的矩阵乘一个 KxN 的矩阵理论上需要 MNK 次乘加但输入数据只有 MK KN 个元素。也就是说每个输入元素平均要被用 N 次或 M 次。问题就来了如果按照最朴素的“取两个数、乘一下、加进去、存回去”的方式每个乘加都要访问一次内存那访存次数就是 MNK 量级而计算本身也是 MNK 量级。访存和计算 1:1这在芯片里是灾难性的因为一次内存访问的能耗可能是一次乘加运算的几十倍甚至上百倍。所以 AI 芯片架构的核心命题就一句话怎么让数据在计算单元里多待一会儿少跑几趟内存。常见的解法有几种。一种是靠缓存层次GPU 就是这么干的用共享内存、寄存器堆把数据尽量留在片上但缓存的管理本身有开销而且数据在缓存和计算单元之间还是要来回搬。另一种是数据流架构让数据在计算单元之间直接流动脉动阵列就是这一类里最经典的代表。2.2 脉动阵列的核心思路让数据“流”起来而不是“搬”来“搬”去脉动阵列的基本结构是一堆排列成二维网格的处理单元PEProcessing Element每个 PE 通常只包含一个乘法器、一个加法器和几个寄存器。数据从阵列的边界流入每个时钟周期向前推进一格就像心脏搏动把血液推向前进一样这也是“systolic”这个词的来源。关键在于每个 PE 算完的结果不会立刻写回内存而是传给相邻的 PE 继续参与计算。以最经典的权重固定Weight Stationary数据流为例权重预先加载到每个 PE 的寄存器里输入数据从左侧流入部分和从上往下流动。这样一个输入数据进入阵列后会依次和多个权重相乘实现输入复用一个权重加载一次后会被多个输入数据反复使用实现权重复用。访存次数从 MNK 降到了 MK KN 量级直接砍掉了几个数量级。我打个比方。传统方式像是你做饭每切一次菜都要跑一趟冰箱拿食材切完再跑一趟放回去。脉动阵列像是把食材摆在一条流水线上每个工位只负责一道工序切好的菜直接传给下一个工位冰箱只在最开始和最末尾各开一次。省下来的全是来回跑腿的时间和体力。2.3 为什么是二维网格而不是一维流水线有人会问一维的流水线不也能实现数据复用吗确实可以但二维网格能同时实现两个方向的复用。一维流水线只能让数据在一个方向上复用比如输入复用或者权重复用但二维阵列可以同时做到行方向和列方向的复用把复用率再提升一个维度。具体来说在一个 NxN 的脉动阵列里每个时钟周期有 N 个输入数据从左侧进入N 个权重从上侧进入经过 N 个周期后NxN 个部分和从下方流出。整个过程中每个 PE 每个周期都在做一次乘加利用率接近 100%。这种规整的结构还有一个好处非常适合硬件实现布线规整时钟树好做功耗可控。当然二维阵列也不是没有代价。它的灵活性差只适合做规整的矩阵运算遇到稀疏矩阵或者不规则计算就抓瞎。而且阵列越大边缘的 PE 等待时间越长填充和排空的开销占比就越高。所以实际设计里阵列尺寸的选择是个权衡活后面我会专门讲怎么算这笔账。3. 核心原理拆解从单个PE到完整阵列的数据流动3.1 单个PE长什么样乘法器加加法器就够了一个最基础的 PE 结构简单到令人发指一个乘法器一个加法器三个寄存器分别存权重、输入、部分和。每个时钟周期PE 做三件事把输入的 activation 和存储的 weight 相乘把乘积加到累加器上然后把 activation 和部分和分别传给右边和下面的邻居。这里有个细节值得注意PE 里的权重寄存器是固定的加载一次后在整个计算过程中不变。这意味着权重不需要每个周期都从内存读取只需要在计算开始前加载一次。这就是权重固定的精髓——把最频繁使用的数据钉在计算单元里。我实测过一个 8 位定点乘加器在 28nm 工艺下大概占 0.0001 平方毫米功耗在微瓦级别。一个 256x256 的阵列就是 65536 个 PE面积大概 6.5 平方毫米功耗在瓦级别。这个密度和能效比是通用处理器很难做到的。3.2 数据怎么流动三种经典数据流的对比脉动阵列有三种经典的数据流方式每种对应不同的复用策略选哪种取决于你的计算模式。数据流类型固定数据流动数据复用重点适用场景权重固定WS权重输入、部分和权重复用权重矩阵较小的推理输出固定OS部分和输入、权重输出复用大矩阵乘、训练输入固定IS输入权重、部分和输入复用输入重复度高的场景权重固定是最常见的因为推理场景下权重是固定的输入是变化的。TPU v1 用的就是权重固定的脉动阵列。输出固定适合训练因为训练时权重会更新把部分和固定在 PE 里可以减少部分和的搬运。输入固定用得相对少但在某些特定卷积模式下有奇效。选择哪种数据流核心看你的计算里哪类数据复用率最高、哪类数据变化最频繁。变化最频繁的数据应该让它流动变化最少的数据应该让它固定。这个原则记住选型就不会错。3.3 填充与排空阵列利用率的最大杀手脉动阵列有个天生的缺陷数据从边界流入到填满整个阵列需要时间。一个 NxN 的阵列从第一个数据进入到最后一个 PE 开始工作需要 2N-1 个周期。这段时间里阵列的利用率是逐渐上升的前面很多 PE 在空转。计算结束后部分和从阵列流出也需要 N 个周期这段时间利用率逐渐下降。所以整个计算过程的有效利用率是有效利用率 总计算周期 / (填充周期 计算周期 排空周期)假设矩阵维度是 MxK 乘 KxN阵列是 NxN那么计算周期大约是 MKN / N^2 M*K/N填充和排空加起来是 3N 左右。当 M 和 K 远大于 N 时填充排空开销可以忽略。但如果矩阵很小比如 MKN那利用率可能只有 25% 左右。这就是为什么实际芯片里脉动阵列通常配一个较大的片上缓存把多个小矩阵拼成一个大矩阵来算减少填充排空的次数。我在做架构评估时通常会算一个“盈亏平衡点”当矩阵维度大于阵列维度的 3 倍以上时填充排空开销才降到 10% 以下。4. 实操建模手把手搭一个脉动阵列模拟器4.1 为什么建议自己写一个模拟器看再多论文不如自己写一遍。脉动阵列的时序行为很微妙光看结构图容易漏掉细节比如数据对齐、部分和累加的时机、边界 PE 的处理。我建议用 Python 写一个周期精确的模拟器不需要考虑硬件细节只模拟数据流动和计算跑几个小矩阵就能把原理吃透。下面是我自己用的一个简化版模拟器核心逻辑以权重固定、4x4 阵列为例import numpy as np def systolic_array(A, B, size4): # A: MxK 输入矩阵, B: KxN 权重矩阵 M, K A.shape K2, N B.shape assert K K2 # 初始化 PE 状态权重、累加器 weight np.zeros((size, size)) psum np.zeros((size, size)) # 加载权重到阵列简化直接映射 for i in range(min(size, K)): for j in range(min(size, N)): weight[i][j] B[i][j] # 模拟数据流动 result np.zeros((M, N)) for t in range(M K size): # 每个周期从左侧注入输入从上侧注入权重已预载 for i in range(size): for j in range(size): # 输入从左边进入向右流动 if j 0 and t - i 0 and t - i M: a_val A[t - i][i] if i K else 0 else: a_val 0 # 从左边邻居传来的值简化处理 # 乘加 psum[i][j] a_val * weight[i][j] # 部分和向下流动这里简化实际需要移位 return result这个代码是简化版真实模拟需要维护每个 PE 的输入寄存器和部分和寄存器每个周期做移位操作。但核心逻辑就是每个周期每个 PE 做一次乘加然后把数据往右和往下推。你把这个跑通再对照时序图看就能明白为什么填充需要 2N-1 个周期。4.2 参数计算阵列尺寸怎么选阵列尺寸的选择是个多目标优化问题涉及面积、功耗、利用率、灵活性。我一般按下面这个流程来算第一步确定目标算力。假设你要做 4 TOPS 的 INT8 推理每个 PE 每周期做一次 8 位乘加算两次操作乘和加那么每个 PE 每周期贡献 2 ops。在 1GHz 时钟下需要 4e12 / (2 * 1e9) 2000 个 PE。开方得到大约 45x45 的阵列取整到 48x48 或 64x64。第二步评估填充排空开销。假设典型卷积层矩阵维度是 256x256阵列 64x64计算周期是 256256/64 1024填充排空约 364 192开销占比 192/(1024192) 15.8%。如果阵列加到 128x128计算周期降到 512填充排空升到 384开销占比 384/(512384) 42.9%反而更差。所以阵列不是越大越好。第三步考虑面积和功耗约束。64x64 阵列约 4096 个 PE按每个 PE 0.0001 平方毫米算核心面积约 0.4 平方毫米加上缓存和互连整体在几个平方毫米量级比较适合边缘端芯片。128x128 就得上到十几平方毫米更适合云端。我自己的经验值是边缘端推理芯片阵列尺寸在 32x32 到 64x64 之间比较甜云端训练芯片128x128 到 256x256 更合适。当然这只是一般规律具体还得看你的工作负载特征。4.3 数据位宽与精度INT8 还是 FP16位宽直接影响 PE 的面积和功耗。一个 INT8 乘法器比 FP16 乘法器小好几倍功耗也低一个数量级。但精度损失是实打实的。我的建议是推理场景优先用 INT8配合量化训练精度损失通常能控制在 1% 以内训练场景必须用 FP16 或 BF16因为梯度对精度敏感。实际芯片里常见的是混合精度脉动阵列支持 INT8 和 FP16 两种模式通过配置切换。硬件上可以复用乘法器的部分电路面积增加不多但灵活性大大提升。TPU v2/v3 就是这么干的。5. 常见问题与排查技巧实录5.1 数据对齐错误最常见的翻车现场脉动阵列对数据对齐极其敏感。输入数据进入阵列的时机、权重加载的顺序、部分和流出的周期任何一个错位都会导致结果完全错误。我踩过最坑的一次是权重加载时行列搞反了结果算出来的矩阵是转置的但数值看起来又“差不多”排查了大半天才发现。排查技巧先用一个单位矩阵做测试。单位矩阵乘任何矩阵都等于原矩阵如果结果不对说明数据流有问题。再用一个全 1 矩阵结果应该是每个元素等于 K内积维度这样可以快速定位是累加逻辑错了还是数据对齐错了。注意仿真时一定要打印每个周期的 PE 状态包括输入寄存器、权重寄存器、部分和寄存器。时序问题光看最终结果很难定位必须看中间状态。5.2 部分和溢出定点数的隐形炸弹用定点数做累加时溢出是个大问题。一个 K256 的矩阵乘每个乘积是 INT8累加结果可能达到 256127127 ≈ 4百万远超 INT8 范围。所以部分和寄存器必须用更宽的位宽通常至少 24 位最好 32 位。我见过一个设计为了省面积部分和只用 16 位结果在某些层上精度暴跌。后来改成 24 位面积增加了不到 5%但精度完全恢复。这个 trade-off 非常划算千万别在这里省钱。5.3 阵列利用率低不一定是阵列的问题很多人发现实测算力远低于峰值第一反应是阵列设计有问题。但根据我的经验八成是软件调度的问题。比如矩阵维度太小、batch size 太小、或者数据没准备好导致阵列空转。排查思路先算理论利用率再对比实测。如果理论利用率就低那是工作负载的问题需要调整调度策略比如把多个小矩阵拼成大矩阵。如果理论利用率高但实测低那才是硬件问题查数据供给带宽、查时钟频率、查缓存命中率。下面这张表是我总结的常见问题速查现象可能原因排查方法解决思路结果完全错误数据对齐错误单位矩阵测试检查加载顺序和时序结果接近但有偏差部分和位宽不足检查累加器位宽加宽部分和寄存器算力远低于峰值调度问题或数据供给不足算理论利用率优化调度、加大batch功耗超预期时钟频率过高或阵列过大功耗仿真降频或缩小阵列面积超预算PE 设计过于复杂综合报告简化 PE、共享乘法器5.4 权重加载的带宽瓶颈权重固定数据流里权重虽然只加载一次但如果阵列很大加载时间也很可观。一个 256x256 的阵列权重有 65536 个按 128 位总线宽度需要 65536*8/128 4096 个周期。如果计算周期只有几千那权重加载时间就占了一半。解决办法有两个一是权重预加载在计算开始前提前加载好和上一层的计算重叠二是权重压缩利用稀疏性减少加载量。我在实际项目里通常会把权重加载和上一层的排空阶段重叠能省掉大部分开销。6. 脉动阵列的边界与替代方案6.1 它不适合什么场景脉动阵列不是万能的。它最适合规整的、密集的矩阵运算。遇到下面这些场景它就不灵了稀疏矩阵。脉动阵列的 PE 是固定连接的稀疏矩阵里大量零元素会导致大量无效计算利用率暴跌。虽然有一些稀疏脉动阵列的研究但硬件复杂度上升很多收益不一定划算。不规则计算。比如图神经网络里的聚合操作、动态形状的注意力这些计算模式不规整脉动阵列的固定数据流很难适配。小矩阵。前面算过矩阵维度小于阵列维度 3 倍时填充排空开销太大不如用更灵活的向量处理器。6.2 替代方案对比SIMT、向量机、存内计算GPU 用的 SIMT 架构灵活性最好能处理各种不规则计算但能效比脉动阵列低。向量机介于两者之间灵活性不错能效也还行但峰值算力不如脉动阵列。存内计算是新兴方向把计算单元嵌入内存阵列理论上能彻底消除访存瓶颈但目前工艺成熟度还不够良率和可靠性是问题。我的判断是未来几年云端训练还是 GPU 的天下但推理场景脉动阵列会占据很大份额尤其是对能效敏感的边缘端。存内计算如果工艺突破可能会改变格局但那是更远的事。6.3 一个实际案例从算法到阵列的映射最后分享一个我做过的实际案例。一个 3x3 卷积输入 56x56x64输出 56x56x128。直接做卷积计算量是 56561283364 ≈ 2.3 亿次乘加。我把卷积展开成矩阵乘im2col 之后输入矩阵是 (5656) x (3364) 3136 x 576权重矩阵是 576 x 128。用 64x64 的脉动阵列分块计算每块 64x64。实测下来阵列利用率在 78% 左右主要损失在填充排空和边界处理。如果 batch size 加到 4利用率能到 85% 以上。这个数字在业界算是不错的了很多论文里报的 90% 以上利用率要么是矩阵特别大要么是忽略了某些开销。我个人在实际操作中的体会是脉动阵列的设计难点不在阵列本身而在数据供给和调度。阵列结构是规整的画出来很漂亮但怎么让数据源源不断地喂进去、怎么让部分和及时排出来、怎么和上层缓存配合这些才是真正花时间的地方。我见过太多设计阵列算力标得很高但实际跑起来利用率不到 50%问题全在数据通路上。另外一个小技巧做架构评估时别只看峰值算力一定要算有效算力。有效算力 峰值算力 x 阵列利用率 x 数据供给率。这三个数乘起来才是你真正能拿到的性能。我一般会把这三个数分别列出来哪一项低就优化哪一项比盲目加大阵列有效得多。