AI芯片算力瓶颈为何在数据搬运?脉动阵列原理与工程实践
1. 为什么一颗AI芯片的算力上限往往卡在“数据搬运”上如果你拆过几块主流的AI加速卡或者翻过几篇讲TPU、NPU架构的论文大概率会反复撞见一个词——脉动阵列Systolic Array。我第一次认真啃这个概念是在琢磨“为什么矩阵乘法在AI芯片里要被单独拎出来做硬件加速”的时候。当时我的直觉是矩阵乘法不就是乘加吗GPU里堆几千个CUDA核心不就完了后来自己动手算了一笔账才明白问题根本不在“算”而在“搬”。AI推理和训练里超过七成甚至九成的运算量都集中在矩阵乘法和卷积上而这两者本质上都能拆成大规模的乘累加MAC操作。一个典型的Transformer层权重矩阵动辄几千乘几千中间激活值也是同量级。如果按传统冯·诺依曼架构每次乘加都要从内存取两个数、算完再写回去那内存带宽会瞬间成为瓶颈。业内有个粗略的经验值做一次32位浮点乘加计算本身消耗的能量大约是几个皮焦耳而从片外DRAM搬一个32位数据进来能耗是它的几十倍甚至上百倍。也就是说算力还没跑满功耗和带宽先被数据搬运吃光了。脉动阵列就是冲着这个矛盾去的。它的核心思路非常朴素让数据像血液在血管里一样有节奏地流过一排排处理单元PE每个PE只负责一次乘加算完就把部分和往下一个PE推同时把权重或激活值横向或纵向传下去。数据一旦进入阵列就会被复用很多次而不是算一次取一次。这种“数据流动、计算就地”的模式把访存次数压到了极低同时让PE之间的连线变得规整非常适合做深亚微米工艺下的高密度布局。所以这篇内容我想从一个一线从业者的角度把脉动阵列这件事讲透。不是复述教科书定义而是把它为什么长这样、怎么算、怎么搭、怎么调、坑在哪这几件事串起来。适合正在做AI芯片架构评估的工程师、想入门加速器设计的学生也适合那些被“算力TOPS”参数绕晕、想搞清楚背后门道的技术爱好者。你不需要先精通数字电路我会尽量用生活化的类比把关键机制说清楚但该有的参数计算和结构细节一个都不会少。2. 脉动阵列到底“脉动”在哪从一次矩阵乘法说起2.1 用“流水线工厂”理解脉动阵列的基本原理先抛开硬件用一个工厂的类比来建立直觉。假设你要生产一批产品每件产品需要经过三道工序拧螺丝、装芯片、贴标签。传统做法是每个工人守着自己的工位产品在工位之间来回搬运工人经常在等零件。脉动阵列的做法更像是一条节拍严格的流水线产品每隔一个固定时钟周期就往前挪一站每个工位只做自己那一道工序做完立刻传给下一站同时新的半成品从上一站补进来。整条线一旦填满每个时钟周期都能吐出一个成品。脉动阵列里的“脉动”这个词就是形容这种数据像心跳一样有节奏地同步流动的状态。每个处理单元PE在每个时钟周期做一次乘加然后把结果传给邻居。关键在于数据不是随机访问的而是按照预先设计好的路径一步一步“挤”过去。权重可以预先加载并固定在PE里激活值从左侧流入部分和从上往下累积最后从底部流出完整结果。这种结构最直接的好处是数据复用率极高。以最经典的N×N脉动阵列做矩阵乘法为例一个权重值一旦加载进某个PE在整个计算过程中会被反复使用N次一个激活值从左边进入后会沿着对角线方向经过N个PE也被复用N次。相比每次乘加都去内存取数访存量下降了几个数量级。这也是为什么Google在第一代TPU里就用了256×256的脉动阵列把矩阵乘法的能效比做到了同期GPU的几倍。2.2 权重固定与数据流动两种主流数据流方式脉动阵列不是只有一种形态实际设计里最常见的是**权重固定Weight Stationary和输出固定Output Stationary**两种数据流偶尔也有输入固定的变体。理解它们的区别对后面做架构选型非常关键。权重固定的思路是把卷积核或权重矩阵的每个元素预先塞进对应PE的寄存器里之后在整个计算过程中不再移动。激活值从阵列左侧依次流入每经过一个PE就与该PE的权重做一次乘加部分和沿着垂直方向往下累加。这种方式的优势是权重访存几乎为零特别适合推理场景——因为推理时权重是固定的一次加载可以服务成千上万次计算。缺点是如果权重矩阵太大装不进阵列就需要分块tiling分块会带来额外的数据调度开销。输出固定则是让部分和留在PE里不动权重和激活值从两个方向流入。这种方式适合输出通道数很大的场景因为部分和不需要频繁搬运能省下不少寄存器读写功耗。实际芯片里很多设计会把两者结合比如在阵列边缘做权重预加载在阵列内部做部分和的定向流动。我个人的经验是做推理加速优先考虑权重固定做训练或动态权重场景则要仔细评估输出固定的收益。因为训练时权重每步都在更新权重固定的优势会被削弱反而部分和的搬运会成为新的瓶颈。2.3 一个4×4阵列的乘加过程拆解为了把原理落到实处我们拿一个最小的4×4脉动阵列手动走一遍矩阵乘法的过程。假设要计算 C A × B其中A和B都是4×4矩阵。我们把B的每一列作为权重预先加载到阵列的每一列PE中A的每一行作为激活值从左侧按行流入。在第一个时钟周期A的第0行第0个元素进入最左上角的PE(0,0)与B的第0列第0个元素相乘得到部分和。第二个时钟周期这个部分和往下传给PE(1,0)同时A的第0行第1个元素进入PE(0,1)A的第1行第0个元素进入PE(1,0)……以此类推。整个阵列像被斜着切了一刀数据沿着对角线波前wavefront推进。等到第7个时钟周期左右第一个完整的结果C(0,0)会从阵列右下角流出。之后每个时钟周期都能流出一个新的结果。整个4×4矩阵乘法理论上只需要大约44-17个周期的填充时间之后就是满速输出。这个填充时间pipeline fill是脉动阵列的一个固有开销阵列越大填充周期越长但相对于总计算量这个开销会被摊薄到可以忽略。这里有个容易被忽略的细节部分和的位宽会随着累加深度增加而增长。4×4阵列里一个结果最多累加4次位宽增长有限但如果是256×256阵列累加深度达到256如果还用定点数就必须预留足够的位宽防止溢出。实际设计里通常会用定点数做推理位宽选16位或24位并在阵列边缘做饱和处理或移位缩放。这个位宽选择直接影响到PE的面积和功耗是架构定义阶段就要拍板的事。3. 从纸面到硅片脉动阵列的硬件实现细节3.1 PE内部到底长什么样一个处理单元PE看起来简单但真到RTL实现的时候每个细节都会影响最终的能效和面积。最基础的PE包含三个部分一个乘法器、一个加法器、一个寄存器用于暂存部分和或权重。乘法器通常用布斯编码Booth Encoding加华莱士树Wallace Tree来做加法器则常用进位保留加法器Carry Save Adder来降低关键路径延迟。我参与过的一个16位定点PE设计里乘法器占了整个PE面积的六成以上。为了省面积我们把乘法器做成流水线结构两级流水这样关键路径能缩短但代价是增加了一级寄存器延迟多了一个周期。这个取舍在脉动阵列里要特别小心因为PE的延迟会直接影响整个阵列的同步节奏。如果某个PE比其他PE慢整个阵列的时钟频率就被它拖住了。另一个容易踩坑的地方是部分和的寄存器复位。在每一轮新的矩阵乘法开始前所有PE里的部分和寄存器必须清零否则上一轮的结果会污染下一轮。这个复位信号如果走线太长会导致复位偏斜reset skew部分PE提前清零、部分滞后结果就全乱了。我们的做法是在阵列内部做局部复位同步用一棵平衡的复位树来分发信号实测下来比全局复位稳得多。3.2 阵列规模怎么定算力、面积与利用率的三角博弈阵列规模是脉动阵列设计里最核心的决策之一。做大了峰值算力高但面积和功耗也上去了而且如果实际 workload 的矩阵维度不够大阵列利用率会惨不忍睹。做小了利用率高但峰值算力上不去可能满足不了目标模型的吞吐要求。这里有一个很实用的估算方法。假设你要跑一个典型的Transformer推理主要矩阵乘法的维度是M×K×N其中M是序列长度K是隐藏维度N是输出维度。如果阵列是S×S的那么一次能处理的块大小就是S×S。当M、K、N都远大于S时阵列利用率可以接近100%但如果某一维小于S比如M1自回归解码的常见情况那阵列里大部分PE就闲置了。我见过不少团队一开始拍脑袋定256×256结果发现实际部署的模型里大量矩阵的M维只有几十利用率长期在30%以下等于花了大面积买了个低效的算力。后来他们改成128×128配合更好的分块调度实际吞吐反而提升了。所以我的建议是先统计目标模型里矩阵维度的分布再决定阵列规模不要盲目追大。一般来说128×128到256×256是当前主流推理芯片的甜点区训练芯片可能会更大但也要配合高带宽内存和复杂的调度逻辑。3.3 数据加载与排空那些教科书不讲的时序细节脉动阵列的优雅之处在于计算过程规整但它的“外围”一点都不简单。数据怎么从内存加载到阵列边缘、怎么保证加载节奏和阵列内部节奏对齐、计算完怎么把结果排空这些才是工程实现里最耗精力的部分。以权重加载为例一个256×256的阵列如果每个PE的权重是16位那总共需要加载256×256×2字节128KB的权重。这些权重不可能在一个周期内全部塞进去必须分批次、按列或按行加载。加载过程中阵列要么暂停计算要么用双缓冲double buffering让计算和加载重叠。双缓冲的做法是准备两块权重存储区一块正在被阵列使用另一块在后台加载下一批权重加载完成后切换。这个切换的时机必须精确控制否则会出现数据竞争。排空drain同样讲究。当最后一批激活值进入阵列后还需要等待若干个周期让还在阵列里流动的部分和全部流出。这个排空周期数等于阵列的深度如果是N×N大约N个周期。如果排空期间就启动下一轮计算新数据会和旧数据在阵列里混在一起结果直接报废。我们的做法是维护一个状态机精确跟踪阵列的填充状态只有在完全排空后才允许下一轮加载。这个状态机的设计看似简单但边界条件特别多调试时经常在这里翻车。4. 实操推演手把手估算一个脉动阵列的算力与能效4.1 算力估算从PE数量到TOPS假设我们设计一个128×128的脉动阵列每个PE每个时钟周期完成一次16位定点乘加。芯片目标频率是1GHz。那么峰值算力可以这样算PE总数 128 × 128 16384个每个PE每周期做1次乘加即2次运算乘和加各算一次峰值算力 16384 × 2 × 1e9 32.768 TOPS万亿次运算每秒这个数字看起来很漂亮但实际有效算力要打折扣。折扣主要来自三个方面一是前面说的填充和排空开销如果处理的矩阵块很小这个开销占比会很高二是数据加载带宽限制如果内存供不上数据阵列会饿着三是调度效率分块策略不好会导致大量PE闲置。我一般会用有效算力 峰值算力 × 阵列利用率 × 带宽满足率来粗略估算。阵列利用率取决于矩阵维度与阵列尺寸的匹配度带宽满足率取决于内存带宽和计算带宽的比值。一个设计良好的推理芯片跑大矩阵时有效算力能达到峰值的70%到85%跑小矩阵或自回归解码时可能只有30%到50%。4.2 能效对比脉动阵列 vs 传统SIMD为了直观感受脉动阵列的能效优势我们拿它和传统的SIMD单指令多数据架构做个对比。SIMD架构里一条指令同时操作多个数据但数据仍然需要从寄存器堆或内存中取出算完再写回。每次乘加的数据搬运次数大约是2次读加1次写。脉动阵列里权重固定在PE内激活值流过时被复用部分和就地累加。以N×N阵列为例完成N×N×N的矩阵乘法权重只加载一次激活值每个元素被复用N次部分和只在PE间传递。粗略估算数据搬运次数能降到SIMD架构的1/N到1/N²量级。考虑到数据搬运的能耗远高于计算本身这个差距直接体现在能效比上。实际芯片数据也印证了这一点。早期用SIMD做矩阵乘法的加速器能效比通常在0.1到0.5 TOPS/W而采用脉动阵列的推理芯片能效比可以做到2到5 TOPS/W差距在十倍左右。当然这个对比不是绝对的SIMD在灵活性上有优势脉动阵列在通用性上偏弱选哪个要看具体场景。4.3 一个可复现的Python模拟脚本光说不练容易飘我写了一个简化的脉动阵列模拟脚本用Python把4×4矩阵乘法的数据流走一遍。这个脚本不涉及硬件时序但能帮你把数据流动的路径和部分和的累加过程看清楚。import numpy as np def systolic_array_matmul(A, B, size4): 模拟权重固定的脉动阵列矩阵乘法 A: 激活值矩阵 (size x size) B: 权重矩阵 (size x size) 返回: C A B # 初始化部分和阵列 psum np.zeros((size, size), dtypenp.int32) # 权重固定在每个PE中这里PE(i,j)固定B[i][j] weights B.copy() # 模拟波前推进总共需要 2*size-1 个周期 for t in range(2 * size - 1): # 每个周期激活值从左侧流入 for i in range(size): # 计算当前周期应该进入第i行的激活值列索引 col t - i if 0 col size: a_val A[i][col] # 该激活值沿着第i行向右流动 for j in range(size): # 激活值到达PE(i,j)时与权重相乘并累加 # 注意实际硬件中激活值是逐周期流动的 # 这里简化为同一周期内完成该行的传播 psum[i][j] a_val * weights[i][j] return psum # 测试 A np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], dtypenp.int32) B np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtypenp.int32) C_sim systolic_array_matmul(A, B) C_ref A B print(模拟结果:\n, C_sim) print(参考结果:\n, C_ref) print(是否一致:, np.array_equal(C_sim, C_ref))这个脚本为了可读性做了简化真实的脉动阵列里激活值在每个周期只移动一格部分和也是逐周期往下传。但通过这个脚本你能直观看到权重固定、激活值流动、部分和累加这三个核心动作是怎么配合的。如果你想更贴近硬件可以把内层循环改成按周期推进给每个PE加一个状态变量模拟流水线填充和排空的过程。5. 踩坑实录脉动阵列设计与部署中的常见问题5.1 阵列利用率低问题往往出在分块策略上阵列利用率低是脉动阵列落地时最常见的抱怨。很多人第一反应是“阵列太大了”但实际排查下来分块策略不合理才是主因。举个例子一个256×256的阵列如果软件调度把矩阵切成256×256的块那没问题但如果因为内存限制切成了64×64的块那每次只有六分之一的PE在工作利用率直接掉到不到20%。解决思路有两个方向。一是调整分块大小尽量匹配阵列尺寸这需要和内存子系统协同设计保证有足够的片上缓存来存放较大的块。二是做多块并行调度让多个小矩阵块同时进入阵列的不同区域填满空闲的PE。第二种做法对调度器要求更高但灵活性更好。我见过一个设计通过把四个64×64的块映射到256×256阵列的四个象限利用率从25%拉回到了80%以上。还有一个隐蔽的坑是边界块处理。当矩阵维度不是阵列尺寸的整数倍时最后一块会不满需要做padding。padding本身没问题但如果padding的数据参与了计算就会引入错误。我们的做法是在PE里加一个valid信号padding数据的valid为0乘加结果不累加。这个信号如果没处理好会出现结果偏差而且很难调试因为大部分块是对的只有边界块出错。5.2 时序收敛难关键路径往往在你想不到的地方脉动阵列的规整性让很多人以为时序收敛很容易实际恰恰相反。阵列越大时钟树越难做平衡PE之间的走线延迟差异会导致建立时间和保持时间违例。我们曾经在一个128×128阵列上遇到时序问题最后发现关键路径不在乘法器里而在部分和从PE(i,j)传到PE(i1,j)的那根线上。因为阵列是二维的垂直方向的走线要穿过很多行线延迟累积起来很可观。解决办法有几个。一是插入流水线寄存器把长走线打断代价是增加延迟周期。二是调整阵列的物理布局让垂直方向的走线尽量短比如把阵列做成扁平的长方形而不是正方形。三是降低时钟频率但这会牺牲算力通常是最后的选择。我们最后采用了流水线加布局优化的组合把频率从800MHz提到了1.2GHz代价是多了两个周期的延迟但整体吞吐反而提升了。5.3 常见问题速查表问题现象可能原因排查思路解决方向阵列利用率长期低于50%分块尺寸与阵列不匹配统计实际矩阵维度分布调整分块策略或阵列规模计算结果在边界块出错padding数据未屏蔽检查valid信号逻辑在PE内加valid门控时序违例集中在垂直走线阵列过大导致线延迟做静态时序分析定位路径插流水线或优化布局功耗高于预期部分和位宽过大分析累加深度和溢出概率缩位宽或做饱和处理加载权重时计算停顿无双缓冲或切换时机不对检查加载状态机引入双缓冲和精确切换复位后首轮结果错误复位偏斜导致部分PE未清零检查复位树平衡性局部复位同步这张表里的每一条都是我在实际项目里真实遇到过的。尤其是复位偏斜那条当时调了整整一周才定位到因为现象是“大部分时候对偶尔错”非常像随机故障最后用示波器抓复位信号才发现不同PE的复位到达时间差了将近一个周期。6. 脉动阵列的边界与演进它适合什么不适合什么6.1 适合的场景大矩阵、固定权重、高吞吐推理脉动阵列最舒服的场景是权重固定、矩阵维度大、吞吐要求高的推理任务。比如卷积神经网络里的卷积层权重在推理时不变卷积核可以展开成矩阵乘法维度通常很大正好喂饱阵列。再比如Transformer里的前馈网络和注意力矩阵乘法虽然注意力部分有动态性但整体上仍然是矩阵乘法为主脉动阵列能发挥很大作用。我个人的判断是只要你的模型以矩阵乘法为主且推理时权重固定脉动阵列就是当前能效比最高的选择之一。它的规整性还带来一个额外好处物理设计相对容易面积和功耗可预测性强适合做大规模量产芯片。6.2 不适合的场景稀疏、动态、小矩阵脉动阵列的软肋也很明显。首先是稀疏计算如果权重或激活值里有大量零脉动阵列仍然会按部就班地算零乘零也占一个周期效率很低。虽然有一些稀疏脉动阵列的研究通过跳过零来提升效率但硬件复杂度会大幅上升实际收益要看稀疏度有多高。其次是动态权重场景比如训练。训练时权重每步都更新权重固定的优势没了反而要频繁加载新权重加载开销可能吃掉计算收益。所以训练芯片往往不会纯用脉动阵列而是结合其他架构。最后是小矩阵或向量运算。自回归解码时很多矩阵乘法的M维只有1阵列里大量PE闲置。这种情况下用SIMD或向量单元反而更高效。所以现在很多AI芯片采用混合架构大矩阵走脉动阵列小运算走向量单元各取所长。6.3 后续可以扩展的方向如果你已经把基础脉动阵列跑通了接下来可以往几个方向深挖。一是稀疏化支持研究如何在PE里加零检测和跳过逻辑同时不破坏阵列的规整性。二是混合精度让阵列支持INT8、FP16甚至BF16的动态切换适应不同层的精度需求。三是三维堆叠把权重存储和计算阵列做垂直集成进一步缩短数据路径。四是可重构阵列让PE之间的连接可以根据算法动态调整兼顾灵活性和能效。我自己最近在琢磨的是脉动阵列与存内计算的结合。如果把权重直接存在PE的存储器里用模拟或数字存内计算的方式做乘加理论上能进一步降低功耗。不过这条路目前还在研究阶段工程化还有不少距离。如果你也在做类似探索欢迎交流踩过的坑和绕过的弯往往比论文里的漂亮曲线更有参考价值。