ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

脉动阵列如何突破AI芯片算力瓶颈:从数据流到工程实现

脉动阵列如何突破AI芯片算力瓶颈:从数据流到工程实现 1. 为什么一颗AI芯片的算力天花板往往由“数据怎么流动”决定如果你拆过几块主流的AI加速卡或者翻过几篇讲矩阵乘法加速的论文大概率会反复撞见一个词——脉动阵列Systolic Array。它不是什么新鲜概念上世纪七八十年代就由H.T. Kung提出来了但这几年因为AI芯片的爆发又被反复拎出来讲。问题在于很多资料一上来就画一堆PEProcessing Element小方块讲数据怎么一波波推过去看完还是不知道它到底解决了什么痛点。我用一句话概括它的价值脉动阵列是一种让数据在计算单元之间“有节奏地流动”从而把矩阵乘法的数据复用率拉到极高的硬件结构。它要解决的核心矛盾是——AI计算里乘加运算本身很便宜真正贵的是把数据从内存搬到计算单元这件事。谁能减少搬运谁就能在同样的功耗和面积下堆出更高的有效算力。这篇文章我打算按一个从业者的视角把脉动阵列从“为什么这么设计”到“怎么落地实现”整条链路拆开讲。适合谁看如果你在做AI加速器的架构评估、写RTL、做算子映射或者只是想知道NPU里那些MAC阵列到底怎么组织的都能拿到能直接用的东西。我会尽量少堆公式多用类比和实际参数把那些文档里不会写的坑一并交代清楚。2. 脉动阵列到底在解决什么问题从矩阵乘法的数据困境说起2.1 一个卷积层到底搬了多少数据先算一笔账这样后面所有设计取舍才有依据。假设一个典型的卷积层输入特征图 56×56×256卷积核 3×3×256输出通道 256。这一层的乘加运算量MACs大约是 56×56×256×3×3×256 ≈ 1.85 G MACs也就是 3.7 G FLOPs 左右。现在看数据量。输入特征图 56×56×256 个元素按 INT8 算就是约 800 KB权重 3×3×256×256 约 590 KB。如果每个 MAC 都要重新从内存读一次操作数那数据搬运量会是运算量的好几倍。而现实是从 DRAM 读一个字节的能量比做一次 INT8 乘加高两到三个数量级。这就是所谓的内存墙——算力不是瓶颈喂数据才是。脉动阵列的第一个设计动机就出来了让同一个数据被尽可能多的计算单元复用减少对内存的重复访问。传统做法是每个 PE 自己去取数脉动阵列则让数据像血液一样在 PE 阵列里“脉动”着流过每个数据进入阵列后会被沿途的多个 PE 依次使用。2.2 三种数据流风格的取舍在讲脉动阵列之前得先把它和另外两种常见数据流放一起对比不然你没法理解它为什么长这样。数据流风格权重是否固定数据复用方式典型代表思路权重固定Weight Stationary是权重留在PE输入流动早期部分加速器输出固定Output Stationary否部分和留在PE累加很多通用矩阵加速器行固定 / 脉动Row Stationary / Systolic部分输入和权重都流动部分和流动TPU类脉动阵列脉动阵列本质上是权重和输入都参与流动的一种折中。它的精妙之处在于权重从左侧一列列推入输入从上方一行行推入两者在 PE 里相遇做乘加部分和则沿着对角线方向往下传。这样每个 PE 在一个周期里只跟相邻 PE 交换数据不需要全局广播连线短、时钟频率容易做高、功耗也低。注意脉动阵列不是万能的。它对规整的、大尺寸的矩阵乘法最友好遇到稀疏、不规则、动态shape的算子利用率会掉得很难看。这也是为什么很多芯片会同时配一个向量单元来兜底。2.3 为什么是“脉动”这个词“Systolic”这个词借用了生理学里的心脏收缩节律。数据像血液一样随着时钟节拍一波一波地泵入阵列每个周期都有新数据进来、旧数据流出整个阵列保持一种有规律的搏动。这个比喻其实很到位因为它强调了两件事一是节律性每个周期固定推进一格二是局部性数据只在相邻单元间传递。理解了这两点你就能明白为什么脉动阵列的 PE 之间几乎不需要复杂的路由网络。每个 PE 只需要三个方向的数据从左来的权重、从上来的输入、从左上来的部分和。控制逻辑简单到极致这也是它面积效率高的根本原因。3. 脉动阵列的核心结构拆解PE、连线和数据节拍3.1 单个PE里到底有什么一个最基础的脉动阵列 PE内部结构其实非常朴素一个乘法器、一个累加器、几个寄存器。以经典的输出固定型脉动阵列为例每个 PE 每个周期做一次acc a * b然后把a往右传、把b往下传、把acc往对角传。这里有个关键细节PE 里的寄存器不是可有可无的它们承担了“对齐”的职责。因为权重和输入进入阵列的时间点不同必须靠寄存器打拍让它们在正确的周期、正确的 PE 里相遇。如果时序对不齐算出来的结果就是错的。我见过不少自己写 RTL 的朋友功能仿真过了但时序对不上最后发现就是某个方向的延迟少打了一拍。3.2 一个 4×4 阵列的完整数据流推演光说结构太抽象我直接拿一个 4×4 的脉动阵列算一个 4×4 矩阵乘 4×4 矩阵把每个周期的数据位置推一遍。假设权重矩阵 W 从左侧输入激活矩阵 A 从上方输入部分和沿对角线向右下流动。设 W 的第 i 行第 j 列元素为 w[i][j]A 的第 i 行第 j 列元素为 a[i][j]。为了让大家看清楚我用一个简化的时序表周期从左侧推入从上方推入说明T1w[0][0]a[0][0]只有左上角PE工作T2w[0][1], w[1][0]a[0][1], a[1][0]对角线两个PE工作T3w[0][2], w[1][1], w[2][0]a[0][2], a[1][1], a[2][0]三个PE工作T4四个元素四个元素阵列填满T5三个元素三个元素开始排空T6两个元素两个元素继续排空T7一个元素一个元素最后一个PE工作可以看到阵列有一个**填充fill和排空drain**的过程。对于 N×N 的阵列算 N×N 的矩阵总周期数是 2N-1。这意味着阵列越大填充排空的开销占比越小利用率越高。但阵列不可能无限大因为面积、功耗和良率都会限制它。实际芯片里常见的规模是 128×128 到 256×256 这个量级。3.3 权重和激活的复用率到底有多高这是脉动阵列最值得吹的一点我用具体数字说明。在一个 N×N 的脉动阵列里每个权重元素进入阵列后会被同一列的所有 PE 依次使用也就是被复用 N 次每个激活元素进入后会被同一行的所有 PE 使用也是 N 次。而部分和则是在 PE 内部累加不需要反复读写内存。对比一下如果不用脉动阵列每个 MAC 都要从缓存里读两个操作数、写一个结果访存次数是运算次数的 3 倍。脉动阵列把这个比例压到了接近 1/N。当 N128 时访存压力直接降了两个数量级。这就是为什么同样的工艺和功耗预算脉动阵列能堆出更高的有效算力。实操心得复用率高不代表实际性能一定好。如果矩阵尺寸小于阵列尺寸或者算子形状不规则填充排空的开销会吃掉大部分收益。做架构评估时一定要拿真实模型的算子分布去跑别只看理论峰值。4. 从理论到芯片脉动阵列的工程实现要点4.1 阵列规模怎么定算一笔面积和功耗的账假设你要设计一个 128×128 的 INT8 脉动阵列每个 PE 包含一个 8 位乘法器、一个 32 位累加器、若干寄存器和控制逻辑。粗略估算单个 PE 的面积在 2000 到 4000 个等效门左右128×128 就是 16384 个 PE总面积相当可观。更关键的是功耗。每个周期整个阵列有 16384 个乘法器同时翻转动态功耗是巨大的。假设每个乘法器每次翻转消耗 0.1 pJ阵列跑在 1 GHz那光乘法器的动态功耗就是 16384 × 0.1 pJ × 1 GHz ≈ 1.6 W。这还没算寄存器和连线的功耗。所以实际芯片里阵列规模、频率、电压是要一起权衡的不是越大越好。我个人的经验是先用目标模型的算子尺寸分布确定一个“甜点区”再倒推阵列规模。如果模型里大部分矩阵乘法的维度在 256 到 512 之间那 128×128 或 256×256 的阵列比较合适如果维度普遍偏小硬堆大阵列反而浪费。4.2 数据位宽与精度INT8、FP16还是混合脉动阵列的 PE 是定点还是浮点直接决定了它的应用场景。INT8 的乘法器面积小、功耗低适合推理FP16 或 BF16 的乘法器复杂得多但能覆盖训练和部分高精度推理。现在主流的做法是混合精度阵列本身支持 INT8 和 FP16 两种模式通过配置切换。实现上INT8 模式下可以把两个 8 位乘法器拼成一个 16 位乘法器或者让 PE 在两种模式间复用部分硬件。这里有个坑累加器的位宽必须按最坏情况设计。INT8 乘加如果累加 256 次结果可能溢出 16 位所以累加器通常要 32 位。如果做 FP16累加还得考虑浮点对齐和舍入硬件复杂度会明显上升。4.3 片上缓存怎么配权重、激活和部分和的三角关系脉动阵列再能复用数据也得有人把数据喂进来。片上缓存的配置直接决定了阵列的利用率。通常会有三块缓存权重缓存、激活缓存、部分和缓存。权重缓存的带宽需求相对低因为权重在推理时是固定的可以提前加载好甚至常驻。激活缓存的带宽需求最高因为每个周期都要往阵列里推新数据。部分和缓存则要看阵列是输出固定还是行固定如果是行固定部分和需要在阵列内部流动缓存压力小一些。一个实用的配置思路是激活缓存的带宽至少要能支撑阵列满负荷运转。比如 128×128 的阵列每个周期需要 128 个激活元素和 128 个权重元素如果位宽是 8 位那激活缓存每周期要提供 128 字节也就是 1 TB/s 级别的带宽在 1 GHz 下。这个数字很吓人所以实际设计里会用多级缓存和预取来平滑。注意别把缓存带宽算得太理想。实际访问会有 bank 冲突、刷新开销、仲裁延迟有效带宽通常只有理论值的 60% 到 80%。做架构评估时留够余量。5. 脉动阵列的典型应用场景与真实案例拆解5.1 卷积神经网络里的隐式矩阵乘法卷积层是脉动阵列最典型的用武之地。虽然卷积看起来不是矩阵乘法但通过 im2colimage to column变换可以把卷积转成一个大矩阵乘法。具体做法是把输入特征图的每个滑动窗口展平成一列权重展平成一个矩阵然后做矩阵乘。这个变换的代价是输入数据被复制了多次内存占用会膨胀。比如 3×3 卷积im2col 后输入数据量会变成原来的 9 倍。所以实际芯片里往往不真的做 im2col而是用隐式矩阵乘法直接在硬件里按卷积的访问模式去取数让脉动阵列“以为”自己在做矩阵乘。这样既保留了脉动阵列的高复用率又避免了数据膨胀。5.2 全连接层和注意力机制全连接层本身就是矩阵乘法直接映射到脉动阵列就行没什么好说的。真正有意思的是注意力机制。Transformer 里的 Q、K、V 计算和注意力矩阵乘法本质上都是矩阵乘但它们的形状比较特殊序列长度可能很长头维度通常不大比如 64 或 128。这就带来一个问题如果头维度小于阵列尺寸阵列利用率会很低。解决办法是把多个头拼在一起做批量矩阵乘或者把序列维度切分后并行。我见过一些设计会把注意力计算拆成多个小矩阵乘然后让脉动阵列分时复用虽然利用率不是 100%但比硬等大矩阵要划算。5.3 一个真实芯片的架构切片拿一个公开资料比较多的 AI 加速器举例这里不点名只讲架构思路。它的核心是一个 256×256 的 INT8 脉动阵列权重从左侧按列推入激活从上方按行推入部分和沿对角线向右下流动。阵列周围配了权重缓存、激活缓存和累加器缓存通过一个专门的 DMA 引擎把数据从 DRAM 搬到片上。它的调度策略是权重预加载 激活流式输入。因为推理时权重是固定的可以提前把整个模型的权重分批加载到权重缓存里然后激活数据像流水一样流过阵列。这样权重缓存的带宽压力很小激活缓存成了主要瓶颈。为了缓解它用了双缓冲一块缓存给当前计算用另一块同时预取下一批数据。实测下来这种架构在 ResNet-50 上的阵列利用率能到 70% 以上但在一些形状不规则的模型上会掉到 40% 以下。所以后来很多芯片都加了向量单元和标量单元来兜底。6. 实操中容易踩的坑与排查技巧6.1 时序对不齐导致结果错误这是自己写 RTL 时最常见的问题。脉动阵列的每个方向都有延迟权重、激活、部分和的到达时间必须精确对齐。如果某个方向的寄存器少打了一拍结果就会错位。排查方法先做一个 2×2 的最小阵列用固定的测试向量跑功能仿真把每个周期的中间结果都打印出来。对照理论时序表看哪个 PE 的输出和预期不符。定位到具体 PE 后再检查它三个输入方向的延迟链。我一般会在 PE 里加一些调试寄存器把关键信号引出来这样波形一看就清楚。6.2 阵列利用率上不去的几个原因利用率低是脉动阵列的常见病原因通常有这几类现象可能原因排查方向小矩阵时利用率骤降填充排空开销占比高检查算子尺寸分布考虑分块策略某些层利用率特别低形状不规则无法映射看是否有向量单元兜底整体利用率都低缓存带宽不足测实际带宽看是否被DMA瓶颈卡住周期性掉利用率双缓冲切换开销检查预取和计算的重叠度我个人的经验是先看算子分布再看缓存带宽最后看调度策略。大部分利用率问题都能在这三步里定位到。6.3 权重加载的隐藏开销很多人只关注计算阶段的性能忽略了权重加载。如果模型很大权重加载的时间可能占总时间的 10% 到 20%。优化方法包括权重压缩稀疏化、量化、权重预取在计算当前层时预取下一层、权重常驻把频繁使用的权重留在片上。实操心得做性能建模时一定要把权重加载时间算进去。我见过一个设计计算部分的理论峰值很高但实际跑起来因为权重加载太慢端到端性能只有理论值的 50%。7. 脉动阵列的边界与演进方向7.1 它不擅长什么脉动阵列的软肋很明确稀疏、不规则、动态形状。遇到这些情况阵列里会有大量 PE 闲置利用率惨不忍睹。比如稀疏矩阵乘法如果非零元素分布随机脉动阵列的规则数据流根本没法高效处理。再比如动态形状的模型每次推理的矩阵尺寸都不一样硬件没法提前优化。所以现在的趋势是异构脉动阵列负责规整的大矩阵乘向量单元处理不规则计算标量单元做控制流。三者协同才能覆盖真实模型的全部算子。7.2 稀疏化和可重构的尝试针对稀疏问题学术界和工业界都在尝试可重构脉动阵列通过配置 PE 之间的连接关系让阵列适应不同的数据流模式。比如在稀疏模式下跳过零元素的乘加只计算非零部分。但可重构的代价是控制逻辑变复杂、面积变大实际收益要看稀疏度。如果稀疏度低于 50%可重构的收益可能还抵不上开销。另一个方向是分块脉动阵列把大阵列拆成多个小阵列每个小阵列独立调度。这样小矩阵可以只激活部分阵列避免填充排空浪费。代价是阵列间的数据共享变复杂需要额外的路由网络。7.3 我个人对脉动阵列的判断脉动阵列不是银弹但它在规整矩阵乘法上的效率优势短期内没有更好的替代方案。未来它大概率会以**“大阵列 小向量 可重构”**的组合形态存在而不是单独打天下。做架构设计时别指望一个脉动阵列解决所有问题关键是搞清楚你的目标模型里规整矩阵乘占多大比例然后据此决定阵列规模和配套单元。最后分享一个小技巧如果你在评估一个脉动阵列设计先别急着看峰值算力拿三个真实模型的算子分布去跑一遍利用率再算端到端性能。这个数字比任何理论峰值都有说服力。我自己踩过好几次坑都是被峰值算力忽悠了实际跑起来才发现利用率根本达不到。
返回列表