ARTICLE DETAIL

资讯详情

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

AI芯片算力瓶颈破解:脉动阵列设计与矩阵乘映射实战

AI芯片算力瓶颈破解:脉动阵列设计与矩阵乘映射实战 1. 为什么一颗AI芯片的算力天花板往往卡在“数据搬运”上如果你拆过几块主流的AI加速卡或者翻过几份推理芯片的白皮书会发现一个反复出现的词——脉动阵列Systolic Array。它不是什么新概念上世纪八十年代卡内基梅隆的H.T. Kung就提出了这套结构但直到TPU、昇腾、寒武纪这些AI芯片把它当成矩阵计算的主力引擎它才真正从论文走进硅片。简单说脉动阵列是一种让数据像心跳一样有节奏地在计算单元之间流动的阵列结构专门用来高效完成矩阵乘法——而矩阵乘法恰恰是深度学习里卷积、全连接、注意力机制最底层的运算形态。它解决的核心问题很朴素算力上去了但数据喂不饱计算单元。传统CPU靠缓存和总线反复搬运数据功耗和延迟都花在“路上”而不是“算上”。脉动阵列的思路是让数据在阵列里“流过”时就完成计算每个计算单元只跟邻居打交道数据复用率极高权重可以提前加载并驻留激活值像波浪一样一波波推过去。这样一来访存带宽压力骤降能效比大幅提升。这篇内容适合谁看如果你是芯片架构方向的学生、刚入行的数字IC设计工程师、做AI推理部署的算法工程师或者只是好奇“为什么TPU能效这么高”的技术爱好者都能从下面这些拆解里拿到可复现的思路。我会从整体设计逻辑讲到具体的PE单元、数据流、参数计算再补上实操中容易踩的坑尽量把这件事说透。2. 脉动阵列的整体设计思路与方案选型2.1 从“冯诺依曼瓶颈”说起为什么不用传统架构硬扛矩阵乘传统CPU做矩阵乘法是取一行、取一列、乘加、写回再取下一组。每次乘加都要访问寄存器或缓存数据在存储和计算之间来回跑。矩阵越大搬运次数越多功耗里可能有六七成都花在数据移动上而不是计算本身。这就是常说的“冯诺依曼瓶颈”——计算单元再快也架不住数据供不上。脉动阵列的破局点在于把计算单元排成二维网格让数据在网格里流动时被反复复用。举个例子一个N×N的阵列做矩阵乘权重矩阵可以预先加载到每个PE里驻留不动激活值从左侧一列列推入部分和从上往下累积。每个PE只做一次乘加然后把结果传给下一个PE。数据不是“取-算-存”循环而是“流过即算完”。这种结构天然匹配矩阵乘法的数据复用模式权重复用N次激活值也复用N次访存次数直接降一个数量级。2.2 三种主流数据流权重固定、输出固定、行固定怎么选脉动阵列不是只有一种玩法按数据驻留方式大致分三类权重固定Weight Stationary, WS权重预加载到PE里不动激活值横向流动部分和纵向累积。适合卷积层权重复用率高的场景TPU第一代就用的这种。输出固定Output Stationary, OS每个PE负责一个输出元素累加结果留在本地权重和激活值都流动。适合全连接层或大矩阵乘能减少部分和的搬运。行固定Row Stationary, RS折中方案一行权重固定激活值对角流动部分和横向传递。Eyeriss这类研究芯片常用能兼顾卷积的多种形状。选哪种取决于你的模型里哪类层占主导。卷积为主、权重复用高WS更省全连接或注意力矩阵大、输出复用高OS更合适。实际芯片里往往是混合的比如卷积用WS全连接切OS靠编译器调度。2.3 阵列尺寸怎么定算力、面积、利用率的三角平衡阵列尺寸N×N不是越大越好。N越大理论峰值算力越高N²个PE并行但面积和功耗也平方级增长。更关键的是利用率如果实际矩阵维度小于N大量PE会闲置。比如N128遇到64×64的矩阵利用率只有25%。实操中常见做法是选一个折中尺寸比如32×32、64×64、128×128再配合**分块Tiling**把大矩阵切成小块喂进去。分块大小要匹配阵列尺寸和片上缓存容量。我见过不少设计为了追峰值算力把阵列堆到256×256结果实际负载利用率不到40%能效反而输给128×128的方案。所以尺寸选择要回到真实工作负载的矩阵维度分布别只看理论峰值。3. 核心细节解析PE单元、数据流与缓存层级3.1 PE单元内部一个乘法器加一个累加器但细节决定成败每个PEProcessing Element的核心就是一个乘加单元MAC输入激活值a和权重w算a×w加上从上方传来的部分和输出到下方。听起来简单但工程上有几个关键点位宽选择INT8还是FP16INT8面积小、功耗低适合推理FP16精度高适合训练。很多芯片支持混合精度PE里做位宽可配置。流水线深度PE内部可以插寄存器做流水提高时钟频率但会增加延迟。通常2到3级流水比较常见。部分和位宽累加结果位宽要比输入宽防止溢出。INT8乘加部分和常用INT32留足余量。注意PE里的乘法器如果直接用综合工具生成面积可能偏大。手工优化过的乘法器在同样位宽下能省20%到30%面积但验证工作量也上去了看项目节奏取舍。3.2 数据流动的节奏为什么叫“脉动”“脉动”这个词很形象数据像血液一样随着时钟节拍一格一格往前推。每个时钟周期激活值向右移一格部分和向下移一格权重不动。这样每个PE在每个周期都在做有效计算没有空转。这种节奏带来的好处是控制逻辑极简。不需要复杂的地址生成和仲裁数据流向固定时钟一到就推进。对比传统SIMD架构里那些复杂的指令译码和寄存器堆访问脉动阵列的控制开销几乎可以忽略功耗自然低。但代价是灵活性差。数据流固定意味着只能高效处理特定形状的运算遇到非规则计算比如稀疏矩阵、动态形状就抓瞎。所以现代AI芯片往往在脉动阵列旁边配一些通用向量单元处理激活函数、归一化这些非矩阵操作。3.3 片上缓存层级权重缓存、激活缓存、部分和缓存怎么配脉动阵列再高效也得有人喂数据。片上缓存通常分三块权重缓存存预加载的权重容量要能放下一个分块的权重矩阵。比如64×64阵列INT8权重一个分块就是64×64×1字节4KB双缓冲就是8KB。激活缓存存输入激活值同样按分块大小配。通常用行缓存Line Buffer做卷积的滑动窗口复用。部分和缓存存中间累加结果尤其是分块累加时需要暂存。位宽大容量需求也大是面积大户。缓存带宽要匹配阵列吞吐。一个N×N阵列每周期需要N个激活值和N个权重WS模式下权重驻留只需激活所以激活缓存读带宽至少N字节/周期。这个数要跟SRAM的端口数和频率对齐不然阵列会饿死。4. 实操过程从矩阵乘到脉动阵列映射的完整推演4.1 用一个4×4矩阵乘例子手算一遍数据流假设要做C A × BA是4×4B是4×4用4×4权重固定脉动阵列。B作为权重预加载每个PE存B的一列。A的行从左侧推入。时钟周期1A[0][0]进入PE(0,0)与B[0][0]乘部分和暂存。 周期2A[0][1]进入PE(0,1)同时PE(0,0)的结果向下传给PE(1,0)A[0][0]向右传给PE(0,1)。以此类推。到第7个周期C[0][0]从PE(3,3)输出。整个4×4矩阵乘需要大约44-17个周期完成第一轮后续流水起来后每周期出一个结果。对比传统方式4×4矩阵乘需要64次乘加每次都要访存脉动阵列把访存次数压到了个位数。4.2 分块策略大矩阵怎么切成阵列能吃的块实际矩阵往往几百上千维阵列只有64×64。分块逻辑是把A按行切、B按列切每块64×64依次送入阵列部分和累加。具体步骤把A分成M/64行块B分成N/64列块。对每个(A行块, B列块)组合送入阵列算64×64的部分积。部分和累加到输出缓存。所有K维度块累加完得到最终C块。这里的关键是部分和的累加顺序。如果K维度分块多部分和要反复读写带宽压力大。优化方法是让部分和尽量留在片上用大容量部分和缓存或寄存器堆暂存减少DDR往返。4.3 参数计算阵列尺寸、频率、带宽的定量关系给一组实际参数算一下。假设64×64阵列频率1GHzINT8运算。峰值算力 64×64×2乘加算两次操作×1G 8.192 TOPS。激活带宽需求 64字节/周期 × 1G 64 GB/s。权重带宽需求 权重驻留只需加载时带宽平均下来很低。如果激活缓存用双端口SRAM位宽512bit64字节频率1GHz刚好满足。如果位宽只有256bit就得降频或缩小阵列不然阵列利用率掉一半。提示算力、带宽、面积这三个数要一起看。我见过有人只盯峰值算力结果带宽配不上实测性能只有峰值的30%。做架构评估时先把带宽算清楚再反推阵列尺寸。4.4 一个可参考的RTL模块划分如果真要动手写RTL模块可以这样分pe_array.v顶层阵列例化N×N个PE连线数据流。pe.v单个PE含乘法器、累加器、流水寄存器。weight_buffer.v权重缓存支持预加载和驻留。activation_buffer.v激活缓存行缓存结构。partial_sum_buffer.v部分和缓存大位宽SRAM。controller.v状态机控制加载、计算、输出三个阶段。验证时先用小阵列比如4×4跑通功能再扩到目标尺寸。直接上大阵列debug会非常痛苦。5. 常见问题与排查技巧实录5.1 阵列利用率低先查矩阵维度再查分块利用率低是最常见的问题。排查顺序看实际负载的矩阵维度分布。如果大量矩阵小于阵列尺寸利用率天然低。解决办法是选更小的阵列或者用多个小阵列并行。查分块策略。分块太小部分和搬运多分块太大缓存放不下。用profiling工具看每层的分块利用率和缓存命中率。查数据对齐。矩阵维度不是阵列尺寸整数倍时边缘块会浪费PE。可以用padding补齐但会增加计算量要权衡。5.2 部分和溢出位宽估算的实操方法部分和溢出会导致结果错误而且不容易发现。估算方法最大累加次数 K维度大小。INT8输入每次乘加最大绝对值127×127≈16129K1024时最大部分和约16129×1024≈1.65e7需要25位。所以INT3231位有效足够。如果K更大或者用FP16要重新算。注意有些设计为了省面积用INT24部分和K大时就会溢出。建议留足余量或者做饱和处理但饱和会引入误差推理精度可能掉点。5.3 时序不收敛流水线插入与时钟树优化大阵列的时序是老大难。64×64阵列数据要横穿64个PE组合逻辑延迟很大。解决办法PE内部插流水把乘法和累加分到两级频率能提上去。阵列内插流水每几个PE之间插一级寄存器切断长路径。时钟树优化阵列大了时钟偏斜难控用H-tree或网格时钟分布。代价是延迟增加但吞吐不变。脉动阵列本来就是流水架构多几级流水对吞吐影响不大。5.4 常见问题速查表问题现象可能原因排查方法解决方向利用率低于50%矩阵维度小、分块不当profiling看每层维度调整阵列尺寸或分块结果错误部分和溢出、位宽不足检查K维度和位宽加宽部分和或饱和处理频率上不去组合逻辑太长时序报告看关键路径插流水、优化乘法器带宽不够缓存位宽或端口不足算带宽需求对比实际加宽SRAM或降阵列尺寸功耗偏高数据翻转多、时钟树差功耗分析看翻转率优化数据编码、时钟门控5.5 几个踩过的坑权重加载阶段别忽略权重预加载也要时间如果频繁切换权重加载开销会吃掉收益。尽量让同一组权重多算几批数据。激活缓存的bank冲突行缓存如果bank划分不合理滑动窗口访问会冲突带宽掉一半。按卷积窗口大小设计bank数。验证要覆盖边界矩阵维度为1、K1、全零输入这些边界情况容易暴露控制逻辑bug。6. 脉动阵列的适用边界与扩展思路脉动阵列不是万能药。它在规则、密集的矩阵乘上效率极高但遇到稀疏、动态、非规则计算就力不从心。现在不少芯片走混合架构脉动阵列做主干矩阵乘旁边配向量单元做激活、归一化、池化再用编译器把算子映射到合适的单元上。这种思路比纯脉动阵列灵活比纯通用架构高效。扩展方向上有几个值得关注的点一是稀疏支持在PE里加零跳过逻辑稀疏矩阵能省算力二是可重构阵列数据流可切换WS/OS适应不同层三是近存计算把阵列和存储堆叠进一步降访存功耗。这些方向目前都有研究芯片在跑离大规模量产还有距离但思路值得跟踪。我个人在实际评估芯片时会先看它的阵列尺寸、数据流模式、缓存带宽这三个数基本就能判断它适合什么负载。脉动阵列的精髓不在单个PE多强而在整个阵列的节奏感和数据复用效率。把数据流理顺了能效自然就上来了。
返回列表