ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

脉动阵列原理与AI芯片设计实战:从矩阵乘法到TPU架构

脉动阵列原理与AI芯片设计实战:从矩阵乘法到TPU架构 1. 从一次面试翻车说起为什么脉动阵列值得单独拎出来讲几年前我去面试一家做AI加速器的公司技术负责人问了我一个问题“如果让你从零设计一个矩阵乘法加速器你会怎么组织数据流”我当时脑子里第一反应是“并行展开、多PE堆叠”然后就开始讲怎么把乘加单元铺开、怎么调度。对方听完沉默了几秒说了一句让我记到现在的话“你描述的是一个SIMD阵列不是脉动阵列。你知道两者最本质的区别在哪吗”那次面试我挂了但这个问题我回去查了很久。后来我才明白脉动阵列Systolic Array之所以在AI芯片领域被反复提起不是因为它“快”而是因为它解决了一个更根本的问题在算力、带宽和功耗三者之间找到那个微妙的平衡点。尤其是当Transformer类模型把矩阵乘法推到极致的时候脉动阵列几乎成了绕不开的架构选择。这篇文章我想把脉动阵列这件事从头到尾讲清楚。不管你是刚接触AI芯片的学生还是正在做加速器设计的工程师或者只是好奇“为什么TPU要用脉动阵列”的技术爱好者我都会从基本原理讲到实操细节从数据流设计讲到踩坑经验。文章会涉及不少硬件层面的东西但我会尽量用生活化的类比来解释保证你不需要有流片经验也能看懂。先给一个最直观的理解脉动阵列就像一条工厂流水线数据像血液一样在阵列中“脉动”流动每个计算单元只做一件小事但所有单元同时在工作。它的核心思想是让数据复用最大化让访存最小化。这个思路在矩阵乘法这种“数据量大、计算密集”的场景里效果极其明显。2. 脉动阵列到底在解决什么问题2.1 从矩阵乘法说起为什么它成了AI芯片的命门任何做过深度学习推理的人都知道卷积、全连接、注意力机制底层拆开来看全是矩阵乘法。一个典型的Transformer层QKV投影、注意力打分、输出投影、FFN两层加起来就是四次大矩阵乘。参数量动辄几十亿每次推理要做的乘加运算MAC数量是天文数字。问题在于通用处理器CPU和图形处理器GPU在处理这种任务时都有各自的瓶颈。CPU的ALU数量有限大部分时间花在取指、译码、访存上GPU虽然并行度高但它的寄存器文件和共享内存之间的数据搬运功耗往往比计算本身还大。有研究数据表明在65nm工艺下一次32位浮点乘加的能耗大约是3.1pJ而从DRAM取一个32位数据的能耗是它的几百倍。也就是说算得再快喂不进去数据也是白搭。脉动阵列的思路就是既然访存这么贵那我就让数据在芯片内部多流动、少往返。数据从边缘流入经过一个个计算单元每个单元顺手做一次乘加然后把结果往下一个单元传。数据像波浪一样在阵列中推进每个周期都有新的数据进入、旧的数据流出整个阵列始终处于工作状态。2.2 和SIMD阵列的本质区别数据复用方式不同很多人会把脉动阵列和SIMD阵列搞混包括当年的我。两者的PE处理单元都是二维排列都做并行计算但数据流动方式完全不同。SIMD阵列里每个PE通常从共享内存或寄存器文件里读取操作数计算完再写回去。数据是“广播”式的所有PE在同一时刻拿到相同的数据然后各自和不同的权重相乘。这种方式的问题在于每次计算都要访问一次存储带宽压力大功耗也高。脉动阵列里数据是“流动”的。以经典的输出驻留Output Stationary或权重驻留Weight Stationary数据流为例权重预先加载到PE中并保持不动激活值从左侧流入部分和从上往下累加。每个PE只负责一次乘加然后把激活值传给右边的PE把部分和传给下面的PE。数据在阵列中“脉动”前进每个周期只移动一步。这个区别带来的效果是巨大的。在SIMD阵列中如果阵列是N×N每个周期需要从内存读取2N个操作数而在脉动阵列中每个周期只需要从边缘输入N个激活值和N个权重权重可以预加载内部的数据搬运全部在PE之间完成。访存次数从O(N²)降到了O(N)这就是脉动阵列最大的价值。2.3 为什么AI芯片特别适合脉动阵列脉动阵列不是新东西上世纪80年代CMU的H.T. Kung就提出了这个概念最早用于信号处理和线性代数。但它在AI芯片时代才真正爆发原因有几个。第一AI计算以矩阵乘法为主计算模式规整非常适合脉动阵列这种“规则数据流”架构。第二AI推理对延迟有一定容忍度但对吞吐量和能效比要求极高脉动阵列的高数据复用率正好契合。第三现代AI芯片的片上存储SRAM容量越来越大可以把权重和部分激活值缓存在片上进一步减少对DRAM的访问。Google的TPU v1是最著名的脉动阵列商用案例256×256的MXU矩阵乘法单元峰值算力92 TOPS功耗只有40W左右。这个能效比在当时的GPU面前是碾压级的。后来TPU v2、v3、v4不断迭代脉动阵列的规模和数据流设计也在进化但核心思想没变。3. 脉动阵列的核心原理拆解3.1 一个PE长什么样乘加单元的最小结构脉动阵列的基本单元是PEProcessing Element每个PE的核心就是一个乘加器MAC加几个寄存器。最简单的PE结构包括一个乘法器、一个加法器、一个用于暂存权重的寄存器、一个用于暂存部分和的寄存器以及若干用于数据传递的寄存器。以权重驻留型脉动阵列为例每个PE的权重在计算开始前加载好计算过程中保持不变。激活值从左侧输入进入PE后与权重相乘乘积加上从上方传来的部分和结果传给下方PE同时激活值直接传给右侧PE。整个过程只需要一个时钟周期PE内部没有复杂的控制逻辑也不需要访存。这种设计的精妙之处在于PE的结构极其简单可以大规模复制。一个256×256的阵列就是65536个PE每个PE只有几百个门电路整体面积和功耗都可控。相比之下如果每个PE都配一套独立的访存和控制逻辑面积和功耗会爆炸。3.2 数据流动的三种经典模式脉动阵列的数据流设计是核心中的核心不同的数据流模式决定了权重、激活值、部分和谁驻留、谁流动。常见的有三种权重驻留Weight Stationary, WS权重预加载到PE中保持不动激活值从左侧流入、从右侧流出部分和从上往下累加。这种模式适合权重复用率高的场景比如卷积层中同一个卷积核在多个位置滑动。输出驻留Output Stationary, OS部分和留在PE中累加权重和激活值从不同方向流入。这种模式适合输出通道数较多的场景可以减少部分和的搬运。行驻留Row Stationary, RS这是Eyeriss提出的一种数据流结合了WS和OS的优点在行方向上做权重驻留在列方向上做输出驻留进一步优化了能耗。实际芯片设计中往往不是纯粹的一种模式而是根据网络结构做混合设计。比如TPU的MXU采用的是权重驻留输出驻留的混合方式权重在阵列中保持部分和在列方向累加。3.3 数据复用率怎么算一个具体的例子假设我们有一个N×N的脉动阵列要计算两个N×N矩阵的乘法。在权重驻留模式下权重矩阵的每一行加载到阵列的每一行PE中激活矩阵的每一列从左侧流入。第一个周期激活矩阵的第一列进入阵列与权重矩阵的第一列相乘得到部分和。第二个周期第一列激活值向右移动一格第二列激活值进入同时第一列的部分和向下移动一格。以此类推经过2N-1个周期后所有计算完成。在这个过程中每个权重被复用了N次因为激活矩阵有N列每个激活值也被复用了N次因为权重矩阵有N行。总访存次数是2N²加载权重和激活而计算量是N³次乘加。计算访存比达到了N/2当N256时这个比值是128。也就是说每从内存取一个数据可以做128次乘加。这个比例在SIMD架构中很难达到。4. 动手设计一个简单的脉动阵列4.1 用Python模拟4×4脉动阵列的行为在写RTL之前我习惯先用Python把数据流模拟一遍确认逻辑正确。下面是一个4×4权重驻留脉动阵列的简化模拟计算两个4×4矩阵的乘法。import numpy as np def systolic_array(A, B, N4): A: 激活矩阵 (N x N) B: 权重矩阵 (N x N) 返回: C A B # 权重驻留: 每个PE保存B的一行 # PE[i][j] 保存 B[i][j] pe_weight [[B[i][j] for j in range(N)] for i in range(N)] # 部分和寄存器, 初始为0 psum [[0 for _ in range(N)] for _ in range(N)] # 激活值寄存器, 每个PE有一个 act_reg [[0 for _ in range(N)] for _ in range(N)] # 总共需要 2N-1 个周期 total_cycles 2 * N - 1 results [[0 for _ in range(N)] for _ in range(N)] for cycle in range(total_cycles): # 从左侧输入激活值 for i in range(N): col_idx cycle - i if 0 col_idx N: act_reg[i][0] A[i][col_idx] else: act_reg[i][0] 0 # 每个PE做乘加 new_psum [[0 for _ in range(N)] for _ in range(N)] new_act [[0 for _ in range(N)] for _ in range(N)] for i in range(N): for j in range(N): # 乘加 new_psum[i][j] psum[i][j] act_reg[i][j] * pe_weight[i][j] # 激活值向右传递 if j N - 1: new_act[i][j1] act_reg[i][j] # 部分和向下传递 for i in range(N-1, 0, -1): for j in range(N): new_psum[i][j] new_psum[i-1][j] if i 0 else new_psum[i][j] # 更新寄存器 psum new_psum act_reg new_act # 记录最下面一行的部分和作为输出 for j in range(N): if cycle N - 1: results[N-1][j] psum[N-1][j] return np.array(results) # 测试 A np.random.randint(0, 5, (4, 4)) B np.random.randint(0, 5, (4, 4)) C_systolic systolic_array(A, B) C_ref A B print(脉动阵列结果:\n, C_systolic) print(参考结果:\n, C_ref) print(是否一致:, np.array_equal(C_systolic, C_ref))这段代码虽然简化了很多细节比如没有处理部分和的精确传递时序但能帮你直观理解数据在阵列中的流动方式。实际RTL设计中每个PE的寄存器更新、数据传递的时序都需要精确控制。4.2 关键参数怎么定阵列规模、数据位宽、时钟频率设计脉动阵列时有几个参数需要权衡阵列规模N×NN越大数据复用率越高但面积和功耗也越大。TPU v1选了256×256是因为当时28nm工艺下这个规模的面积和功耗在可接受范围内同时能覆盖大多数矩阵乘法的维度。如果N太小比如16×16复用率不够访存瓶颈依然存在如果N太大比如512×512面积和功耗会急剧上升而且很多矩阵乘法的维度可能填不满阵列利用率下降。数据位宽AI推理通常用INT8或FP16训练用FP32或BF16。位宽越窄PE的面积和功耗越小同样面积下可以放更多PE。TPU v1用的是INT8后来TPU v2支持了FP32和BF16。实际选择要看目标场景如果只做推理INT8足够如果要兼顾训练需要支持浮点。时钟频率脉动阵列的关键路径通常比较短一个乘加加几个寄存器所以可以跑比较高的频率。但频率越高功耗越大。TPU v1的时钟频率是700MHz不算高但配合256×256的阵列峰值算力已经很高了。实际设计中频率和阵列规模需要联合优化找到能效比最高的点。4.3 从Python到RTL一个PE的Verilog实现下面是一个简单的PE的Verilog代码实现了权重驻留、激活值右传、部分和下传的功能。module pe #( parameter DATA_WIDTH 8, parameter ACC_WIDTH 32 )( input wire clk, input wire rst_n, input wire load_weight, input wire [DATA_WIDTH-1:0] weight_in, input wire [DATA_WIDTH-1:0] act_in, input wire [ACC_WIDTH-1:0] psum_in, output reg [DATA_WIDTH-1:0] act_out, output reg [ACC_WIDTH-1:0] psum_out ); reg [DATA_WIDTH-1:0] weight_reg; reg [DATA_WIDTH-1:0] act_reg; reg [ACC_WIDTH-1:0] psum_reg; always (posedge clk or negedge rst_n) begin if (!rst_n) begin weight_reg 0; act_reg 0; psum_reg 0; act_out 0; psum_out 0; end else begin if (load_weight) begin weight_reg weight_in; end act_reg act_in; psum_reg psum_in act_reg * weight_reg; act_out act_reg; psum_out psum_reg; end end endmodule这个PE每个周期做一次乘加激活值延迟一拍输出部分和延迟一拍输出。实际设计中可能需要流水线更深来满足时序也可能需要处理权重加载和计算的重叠。5. 脉动阵列在AI芯片中的实际应用5.1 TPU的MXU最经典的商用案例Google TPU v1的MXU是一个256×256的脉动阵列支持INT8乘加峰值算力92 TOPS。它的数据流是权重驻留型权重从片上SRAM加载到阵列中激活值从左侧流入部分和从下方流出。MXU的设计有几个值得注意的点。第一它用了双缓冲double buffering来隐藏权重加载的时间当一组权重在计算时下一组权重已经在后台加载。第二它的部分和累加器位宽是32位防止INT8乘加溢出。第三它的控制逻辑非常简化因为脉动阵列的数据流是固定的不需要复杂的调度。TPU v1的实测性能在推理任务上比同时期的GPU和CPU高出一个数量级能效比更是碾压。这直接证明了脉动阵列在AI推理场景中的价值。5.2 其他AI芯片的脉动阵列变体除了TPU很多AI芯片也采用了脉动阵列或类似的数据流架构。比如华为昇腾的达芬奇架构核心是一个3D Cube单元可以看作脉动阵列的变体支持更灵活的数据流配置。寒武纪的MLU系列用了类似的思路但在数据流和存储层次上做了自己的优化。还有一些针对特定场景的芯片比如边缘端的AI加速器会用小规模的脉动阵列比如32×32或64×64配合低功耗设计在手机、摄像头、IoT设备上做实时推理。5.3 脉动阵列的局限性和适用边界脉动阵列不是万能的。它的优势在于规则的大矩阵乘法但如果计算模式不规则比如稀疏矩阵、动态形状的模型脉动阵列的利用率会大幅下降。因为阵列是固定的如果矩阵维度填不满很多PE就闲置了。另外脉动阵列对片上存储的要求很高。权重和激活值需要在阵列附近有足够的SRAM来缓冲否则数据供应不上阵列就会饿死。TPU v1配了28MB的片上SRAM这在当时是很大的面积开销。还有一个问题是灵活性。脉动阵列的数据流是硬件固定的如果要支持不同的计算模式比如卷积、反卷积、注意力需要额外的控制逻辑或者多次映射。相比之下GPU的SIMT架构更灵活但能效比不如脉动阵列。6. 实操中容易踩的坑和排查技巧6.1 数据对齐和边界处理在实际部署脉动阵列时矩阵维度往往不是阵列规模的整数倍。比如阵列是256×256但某个矩阵是100×100这时候需要做padding。padding的方式会影响计算结果的正确性和阵列利用率。我踩过的一个坑是padding时用了0填充但部分和的累加逻辑没有正确处理边界导致结果偏大。后来发现是因为padding的0也参与了乘加虽然乘积是0但部分和的传递时序错了导致某些PE在错误的周期累加了无效数据。解决办法是在边界周期屏蔽无效的部分和传递或者用mask信号控制。6.2 权重加载和计算的重叠权重加载需要时间如果加载和计算串行阵列的利用率会下降。常见的做法是双缓冲用两块权重缓存一块在计算时另一块在加载。但双缓冲的控制逻辑比较复杂容易出现时序冲突。我的经验是权重加载的带宽要匹配阵列的消耗速度。如果阵列每个周期消耗256个权重加载带宽至少要达到这个数否则计算会停顿。实际设计中往往用DMA把权重从DRAM搬到片上SRAM再从SRAM加载到PE这个链路的带宽要仔细计算。6.3 部分和溢出的排查INT8乘加的结果是16位但累加多次后可能溢出。比如256次累加最大结果是256×255×255≈16.6M需要25位才能表示。如果累加器只有16位就会溢出。排查溢出的方法是在仿真时监控累加器的最高位如果出现翻转说明溢出了。解决办法是增加累加器位宽或者在累加过程中做饱和处理。TPU v1用了32位累加器就是为了防止溢出。6.4 常见问题速查表问题现象可能原因排查方法解决方案计算结果偏大padding边界处理错误检查边界周期的部分和传递加mask信号屏蔽无效累加阵列利用率低矩阵维度不匹配统计PE活跃周期占比调整阵列规模或做矩阵分块时序不收敛关键路径太长静态时序分析插入流水线寄存器功耗过高时钟频率太高或阵列太大功耗仿真降低频率或缩小阵列数据供应不上片上SRAM带宽不足带宽仿真增加SRAM或优化数据流7. 一些个人体会和后续可以深挖的方向脉动阵列这个东西入门容易精通难。基本原理几个小时就能看懂但真正要做好一个可用的设计需要反复迭代数据流、存储层次、控制逻辑。我在实际项目中最大的体会是不要追求理论上的最优而要追求工程上的平衡。阵列规模、位宽、频率、存储容量每一个参数都会影响最终的性能和功耗没有银弹。如果后续要继续深入我建议从几个方向入手。一是研究不同的数据流模式比如Eyeriss的行驻留看看在特定网络结构下能不能比权重驻留更优。二是探索稀疏脉动阵列针对剪枝后的模型做优化跳过零值计算。三是结合近存计算Near-Memory Computing的思路把部分计算搬到存储旁边进一步减少数据搬运。最后分享一个小技巧在写RTL之前一定要用Python或C把数据流模拟清楚确认每个周期的数据位置和计算结果。直接写RTL调试的代价太高仿真一次可能要几分钟而Python模拟几秒钟就能跑完。这个习惯帮我省了很多时间。
返回列表