1. 为什么一颗AI芯片里塞得下几千个乘法器却喂不饱数据第一次接触AI芯片架构的人十有八九会掉进同一个思维陷阱算力等于乘法器数量乘以频率。按这个逻辑堆乘法器就行了堆得越多算力越猛。但真正做过加速器设计的人都知道事情远没有这么简单。我刚开始研究神经网络加速器那会儿也天真地以为把PE阵列做大就能线性提升吞吐。后来跑仿真才发现当阵列规模超过某个阈值之后算力增长曲线急剧变平——不是乘法器不够快是数据根本送不进来。这个问题在业内有个很形象的说法叫内存墙。计算单元和存储单元之间的带宽成了整个系统的瓶颈。脉动阵列Systolic Array就是在这个背景下被重新推到台前的。它不是什么新概念早在1982年H.T. Kung就在CMU提出了这个结构。但直到Google TPU v1横空出世大家才真正意识到这个几十年前的老思路恰好是解开AI推理算力瓶颈的那把钥匙。这篇文章我打算把脉动阵列从原理到实现彻底拆一遍。不管你是正在做加速器架构选型的芯片工程师还是想搞清楚TPU内部到底怎么跑起来的算法同学又或者只是对AI芯片感兴趣想入个门我都会尽量用大白话把这件事讲透。核心关键词就两个AI芯片和脉动阵列。我会解释它为什么能解决数据复用问题、怎么设计一个基本的脉动阵列、实际部署时会踩哪些坑以及它到底适合什么场景、不适合什么场景。2. 脉动阵列到底在解决什么问题2.1 从矩阵乘法说起AI推理的本质就是不停地做MAC不管你跑的是CNN、Transformer还是什么别的网络结构拆到最底层绝大部分计算都是矩阵乘法或者卷积。而矩阵乘法和卷积拆到最底层就是乘加运算MACMultiply-Accumulate。一个N×N的矩阵乘一个N×N的矩阵需要N³次MAC。一个大语言模型的推理过程动辄需要几十亿甚至上百亿次MAC。这就带来一个很现实的问题如果每个MAC都要从内存里读两个操作数、写一个结果那内存带宽根本扛不住。打个比方你有一个超级快的厨师切菜速度是每秒一千刀但食材每隔十秒才送来一批那这个厨师大部分时间都在站着等。传统CPU和GPU的架构就面临这个困境——计算单元和寄存器堆、缓存之间的数据搬运消耗了大量时间和功耗。业内有个经验数据在28nm工艺下一次32位浮点乘法的能耗大约是3.7pJ而一次从DRAM读取32位数据的能耗大约是640pJ差了将近200倍。也就是说数据搬运的代价远远大于计算本身的代价。这就是为什么AI芯片设计的核心命题不是怎么算得更快而是怎么让数据少跑路。2.2 数据复用脉动阵列的核心思想脉动阵列的核心思想可以用一句话概括让数据在计算单元之间流动起来每个数据被复用多次而不是每次都从内存重新读取。想象一个工厂流水线。传统做法是每个工位都去仓库领原料做完自己的工序再把半成品放回仓库下一个工位再去仓库领。脉动阵列的做法是原料从流水线一端进入每经过一个工位就被加工一次同时半成品直接传给下一个工位不需要回仓库。数据在阵列中像血液一样脉动流动每个计算单元只负责一小部分工作但数据在整个流动过程中被反复利用。具体到矩阵乘法C A × B脉动阵列的做法是A矩阵的元素从左侧流入B矩阵的元素从上方流入每个PEProcessing Element在本地做一次MAC然后把A的元素向右传、B的元素向下传。这样每个A元素和每个B元素都会被多个PE使用数据复用率极高。2.3 和传统架构的对比为什么不用GPU那套你可能会问GPU不也是并行计算吗为什么不用GPU的架构来做AI芯片GPU的架构本质上是SIMT单指令多线程它通过大量的线程切换来隐藏内存延迟。每个线程有自己的寄存器和上下文当某个线程等待数据时调度器切换到另一个线程继续计算。这套机制在图形渲染和通用并行计算上非常有效但用在AI推理上有两个问题一是线程调度和寄存器管理的开销大二是数据复用主要靠共享内存和缓存效率不如脉动阵列直接。脉动阵列则是一种空间换时间的思路。它把数据复用直接做进了硬件连线里不需要复杂的调度逻辑每个周期数据自动流动到下一个PE。控制逻辑极其简单功耗也低得多。这也是为什么TPU v1能做到那么高的能效比——它的核心就是一个256×256的脉动阵列峰值算力92 TOPS功耗只有40W左右。3. 拆开一个脉动阵列从PE到阵列的完整设计3.1 PE长什么样一个乘法器加一个累加器就够了脉动阵列的基本单元叫PEProcessing Element结构简单到令人发指。一个典型的PE包含一个乘法器一个累加器若干寄存器用于暂存输入数据数据传递通路每个时钟周期PE做三件事接收来自左边和上边的数据做一次乘加运算把数据传给右边和下边。就这么简单。没有指令译码没有分支预测没有乱序执行。所有的控制逻辑都是固定的数据流。这种极简设计带来的好处是单个PE的面积和功耗可以做到非常小从而在同样的芯片面积下塞进更多的计算单元。TPU v1的256×256阵列就是65536个PE每个PE每个周期做一次8位整数的MAC。3.2 权重固定 vs 权重流动两种数据流的取舍设计脉动阵列时第一个要做的架构决策是哪些数据固定哪些数据流动最常见的方案是权重固定Weight Stationary。每个PE预先加载一个权重值输入数据激活值从左侧流入部分和从上方流入或向下流出。这种方案的好处是权重只需要加载一次之后在整个计算过程中保持不变大大减少了权重数据的搬运。对于推理场景权重是固定的这个方案非常合适。另一种方案是输出固定Output Stationary部分和留在PE本地累加权重和激活值都流动。这种方案适合卷积运算中输出特征图需要反复累加的场景。还有**行固定Row Stationary**等混合方案在不同的维度上做权衡。选择哪种方案取决于你的网络结构中权重和激活值的复用模式。我个人的经验是对于以矩阵乘法为主的Transformer类模型权重固定方案实现简单、效率高对于卷积层较多的CNN可能需要根据卷积核大小和通道数做更细致的分析。3.3 阵列尺寸怎么定不是越大越好阵列尺寸是脉动阵列设计中最关键的参数之一。尺寸越大数据复用率越高但灵活性越差利用率也越容易下降。假设你设计了一个N×N的阵列那么它天然适合处理N×N的矩阵乘法。如果你的实际矩阵维度是M×K其中M和K都小于N那阵列就会有大量PE闲置。比如一个128×128的阵列处理64×64的矩阵利用率只有25%。TPU v1选择256×256是因为Google在分析了自己的推理负载后发现大部分矩阵乘法的维度都在256附近。但这个选择也有代价当遇到维度较小的运算时利用率会明显下降。后来的TPU版本引入了多个较小的阵列或者可重构的阵列结构来缓解这个问题。我的建议是先分析你的目标负载中矩阵维度的分布选择覆盖最多运算的那个尺寸。如果负载比较杂可以考虑用多个小阵列拼接或者设计支持分块计算的机制。3.4 数据流控制怎么让数据在正确的时间到达正确的位置脉动阵列最精妙的地方在于数据流的同步。数据从阵列边缘流入每个周期移动一格经过若干周期后到达目标PE。这要求所有数据在时间上精确对齐。举个例子在权重固定的方案中激活值从左侧流入每个周期向右移动一格。第一行PE在第1个周期收到第一个激活值第二行PE在第2个周期收到同一个激活值因为它需要先经过第一行。这意味着不同行的PE处理的是同一个激活值的不同部分需要配合不同的权重。这种时间上的错位skew是脉动阵列设计中最容易出错的地方。我在早期做仿真的时候经常因为skew设置不对导致结果完全错误。后来养成了一个习惯先用一个小规模的4×4阵列手工推演一遍数据流确认每个周期每个PE的状态都正确再扩展到大规模阵列。4. 从零搭建一个脉动阵列实操过程记录4.1 用Python模拟一个4×4脉动阵列在写RTL之前我强烈建议先用高级语言把数据流模拟一遍。Python的numpy虽然不能直接模拟周期级行为但可以用循环来模拟每个周期的状态变化。下面是一个简化版的4×4权重固定脉动阵列模拟import numpy as np def systolic_array_4x4(A, B): 模拟4x4权重固定脉动阵列 A: 4x4 激活值矩阵 B: 4x4 权重矩阵 返回: C A B N 4 # 初始化PE状态 weights B.copy() # 每个PE固定一个权重 psum np.zeros((N, N)) # 部分和 result np.zeros((N, N)) # 激活值需要skew第i行延迟i个周期进入 # 权重需要skew第j列延迟j个周期进入 total_cycles 3 * N # 流水线填满需要的时间 for cycle in range(total_cycles): # 计算每个PE的输入 for i in range(N): for j in range(N): # 激活值从左边来第i行延迟i个周期 a_idx cycle - i if 0 a_idx N: a_val A[i][a_idx] else: a_val 0 # 权重从上面来第j列延迟j个周期 b_idx cycle - j if 0 b_idx N: b_val B[b_idx][j] else: b_val 0 # MAC操作 psum[i][j] a_val * b_val # 每N个周期一行结果计算完成 if cycle 2 * N - 1: row (cycle - (2 * N - 1)) // N if row N: result[row] psum[row] return result # 测试 A np.random.randint(0, 10, (4, 4)) B np.random.randint(0, 10, (4, 4)) C_systolic systolic_array_4x4(A, B) C_ref A B print(脉动阵列结果:\n, C_systolic) print(参考结果:\n, C_ref) print(是否一致:, np.allclose(C_systolic, C_ref))这段代码虽然简化了很多细节比如没有模拟PE之间的数据传递延迟但足以帮你理解数据流的基本逻辑。跑一遍之后你会对skew和流水线填充有直观的感受。4.2 RTL实现的关键模块如果你要真正流片接下来需要用Verilog或VHDL写RTL。一个基本的脉动阵列包含以下模块PE模块核心是乘法器和累加器。累加器需要支持清零和保持以便在计算完成后读出结果。乘法器的位宽需要根据你的精度需求来定——8位整数乘法比32位浮点乘法在面积和功耗上都有数量级的优势。数据分发模块负责把输入数据按照正确的时序送到阵列边缘。这通常是一个带有移位寄存器的FIFO结构每个输出端口有不同的延迟。控制模块生成全局的时钟使能、复位、累加器清零等信号。脉动阵列的控制逻辑非常简单基本上就是一个状态机在控制计算的开始和结束。结果收集模块从阵列边缘收集计算结果。由于结果是流水线输出的需要按照正确的顺序缓存和输出。我在实际项目中遇到的一个坑是累加器的位宽不够导致溢出。8位输入相乘得到16位结果累加N次后需要log2(N)位的额外位宽。如果阵列是256×256累加256次那累加器至少需要16824位。这个细节在架构设计阶段很容易被忽略等到仿真发现结果不对再改就麻烦了。4.3 时序收敛与物理设计注意事项脉动阵列的时序收敛是一个大挑战尤其是当阵列规模较大的时候。数据需要在单个时钟周期内从一个PE传到相邻PE如果阵列的物理布局不合理走线延迟可能超过时钟周期。解决这个问题的常用方法有两个一是插入流水线寄存器把长走线打断代价是增加一个周期的延迟二是优化物理布局让PE按照数据流方向排列减少走线长度。我在一个128×128的项目中最初没有做布局规划综合出来的时序差了将近30%。后来把PE按照行列对齐排列走线长度大幅缩短时序才勉强收敛。如果阵列再大可能就需要考虑三维堆叠或者多芯片互联的方案了。5. 脉动阵列的局限性与适用边界5.1 灵活性差为什么它做不了所有事脉动阵列最大的问题就是不够灵活。它的数据流是硬连线的一旦流片就无法更改。如果你的模型结构发生了变化比如从CNN换成了Transformer或者卷积核大小变了脉动阵列可能就无法高效处理。相比之下GPU或者FPGA的可重构性要好得多。GPU可以通过编程来适应不同的计算模式FPGA可以重新配置电路结构。脉动阵列则是一招鲜在它擅长的矩阵乘法上效率极高但遇到不匹配的运算就抓瞎。这也是为什么现在很多AI芯片采用了可重构脉动阵列的设计——在脉动阵列的基础上增加一些灵活性比如支持不同的数据流模式切换或者允许PE之间跳过某些连接。但可重构性带来的面积和功耗开销需要在设计时仔细权衡。5.2 稀疏性处理脉动阵列的天然短板另一个大问题是稀疏性。神经网络剪枝之后权重矩阵中会有大量的零。理想情况下零值不应该参与计算这样可以节省大量的计算资源和功耗。但脉动阵列的每个PE每个周期都在做MAC不管输入是不是零。处理稀疏性需要额外的逻辑来检测零值并跳过计算这会增加控制复杂度破坏脉动阵列规整的结构。目前学术界有很多关于稀疏脉动阵列的研究比如在PE中增加零值检测电路或者设计非零值压缩传输机制但都还没有形成统一的高效方案。5.3 什么场景最适合脉动阵列根据我的经验脉动阵列最适合以下场景推理为主权重固定数据流模式确定不需要频繁切换计算模式矩阵乘法占主导Transformer、全连接层、大卷积核的CNN对能效比要求极高边缘设备、移动端、数据中心的大规模推理集群批量大小固定或变化不大批量大小变化会导致阵列利用率波动不太适合的场景包括训练需要反向传播和梯度更新数据流模式复杂多变小批量或实时推理批量太小会导致阵列利用率极低稀疏模型无法有效利用稀疏性多模态混合负载需要频繁切换计算模式6. 常见问题与排查技巧实录6.1 仿真结果不对先检查这三个地方脉动阵列仿真出错是家常便饭我总结了一个排查顺序第一检查skew设置。这是最常见的错误来源。激活值和权重的延迟是否和PE的位置匹配第i行第j列的PE应该在哪个周期收到哪个数据手工推演一遍小规模阵列确认时序正确。第二检查累加器位宽。前面提到过累加次数多了之后位宽不够会溢出。用Python或者MATLAB先算一下理论最大值确认位宽足够。第三检查边界条件。阵列边缘的PE没有来自左边或上边的输入这些输入应该置零还是保持上一个值不同的设计选择会导致不同的结果。6.2 利用率上不去看看你的矩阵维度如果发现阵列利用率很低先别急着改架构看看你的实际矩阵维度是多少。一个256×256的阵列处理128×128的矩阵利用率只有25%。这时候可以考虑矩阵分块把大矩阵切成小块分多次计算提高利用率多阵列并行用多个小阵列同时处理不同的矩阵块动态阵列划分根据矩阵维度动态配置阵列的工作模式6.3 功耗比预期高可能是数据搬运没优化好脉动阵列的功耗主要来自三个方面计算、数据搬运、时钟树。如果功耗比预期高大概率是数据搬运的问题。检查一下输入数据的位宽是否合理8位整数和32位浮点的功耗差好几倍数据在阵列边缘的缓存是否足够频繁访问外部内存会大幅增加功耗时钟频率是否过高降低频率可以显著降低功耗但会影响吞吐6.4 常见问题速查表问题现象可能原因排查方法解决方案仿真结果全零累加器未清零或数据未加载检查控制信号时序确认复位和使能信号结果部分正确skew设置错误手工推演小阵列调整延迟参数结果溢出累加器位宽不足计算理论最大值增加位宽利用率低矩阵维度不匹配统计负载维度分布分块或调整阵列尺寸时序不收敛走线延迟过大查看布局布线报告插入流水线寄存器功耗过高数据搬运频繁分析内存访问模式增加本地缓存7. 脉动阵列的演进与我的几点体会7.1 从TPU v1到最新架构脉动阵列的进化路线TPU v1的脉动阵列是纯粹的权重固定结构256×256的规模只支持8位整数运算。到了TPU v2和v3Google引入了浮点支持和更大的阵列规模同时增加了对训练的支持。TPU v4则进一步优化了互联和并行能力支持多芯片协同计算。除了Google很多其他AI芯片也采用了脉动阵列或者类似的结构。比如一些边缘推理芯片采用小规模的脉动阵列配合可重构的数据流来适应不同的网络层。学术界也在探索稀疏脉动阵列、可重构脉动阵列、三维脉动阵列等新方向。7.2 我踩过的几个坑第一个坑是低估了skew的复杂度。我最初以为只要把数据延迟相应的周期数就行了后来发现当阵列支持多种数据流模式时skew的逻辑会变得非常复杂需要仔细设计状态机。第二个坑是忽略了累加器的复位时序。累加器需要在正确的时间清零否则会把上一轮的结果带进来。这个bug在仿真中很难发现因为大部分时候结果看起来是对的只有在特定输入下才会出错。第三个坑是物理设计阶段没有提前规划。脉动阵列的规整性对布局布线很友好但如果不在早期考虑物理约束后期可能会遇到严重的时序和拥塞问题。7.3 给正在做架构选型的同行几句实在话如果你正在考虑用脉动阵列做AI加速器我的建议是先搞清楚你的负载特征。如果你的负载以大规模矩阵乘法为主权重固定批量大小稳定那脉动阵列是非常好的选择。如果你的负载比较杂或者需要支持训练那可能需要考虑更灵活的架构。不要盲目追求大阵列。阵列尺寸要和你的实际负载匹配太大的阵列在遇到小矩阵时利用率会很低。可以先从128×128或者64×64开始根据实际负载数据再调整。仿真验证要趁早。在写RTL之前用Python或者C把数据流模拟一遍确认时序和结果都正确。这会帮你省下大量的调试时间。物理设计要提前介入。在架构设计阶段就考虑布局布线的约束避免后期出现无法收敛的情况。脉动阵列不是一个新东西但它在AI芯片领域的价值正在被重新发现。它的简洁、高效、低功耗恰好契合了AI推理的核心需求。当然它也不是万能的理解它的边界和局限才能做出正确的架构决策。