ARTICLE DETAIL

资讯详情

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

AI芯片脉动阵列:从矩阵乘法到TPU架构设计与工程实践

AI芯片脉动阵列:从矩阵乘法到TPU架构设计与工程实践 1. 从矩阵乘法说起为什么AI芯片需要脉动阵列搞AI芯片的人绕不开一个话题矩阵乘法怎么加速。不管是大模型里的注意力机制还是卷积神经网络里的特征提取底层计算几乎全是矩阵乘加。一个典型的Transformer层QKV投影、输出投影、FFN里的两层全连接全都是大规模矩阵乘法。问题在于通用CPU做矩阵乘法效率极低——取数、计算、存回大部分时间花在数据搬运上而不是计算本身。我第一次接触脉动阵列这个概念时脑子里冒出的第一个问题是为什么不让数据像流水一样穿过计算单元算完就走而不是反复存取这个直觉其实就抓住了脉动阵列的核心思想。它的本质是让数据在计算单元之间流动每个计算单元只做简单的乘加数据流过时顺便被计算掉。这种设计把数据复用做到了极致特别适合矩阵乘法这种规整的计算模式。脉动阵列Systolic Array这个名字来源于“systole”这个词原本是描述心脏收缩泵血的节律。在硬件架构里它指的是数据像心跳一样有节奏地流过阵列每个周期都有新的数据进入、旧的数据流出计算单元始终处于忙碌状态。这个概念最早由H.T. Kung在1982年提出当时是为了解决VLSI超大规模集成电路中计算与通信的平衡问题。但真正让它大放异彩的是Google在2016年发布的TPUTensor Processing Unit。为什么TPU选择脉动阵列而不是传统的SIMD或GPU架构核心原因在于能效比。GPU虽然算力强但它的架构是为图形渲染设计的有大量的寄存器文件和复杂的调度逻辑功耗很大一部分花在了指令调度和数据搬运上。而脉动阵列把控制逻辑简化到了极致——每个计算单元只需要知道“什么时候把数据往下传、什么时候把部分和往右传”不需要复杂的指令译码。这种极简的控制逻辑意味着更少的晶体管用于控制、更多的晶体管用于计算能效比自然就上去了。我实测过一些边缘端的AI加速方案发现一个规律当矩阵规模足够大且形状规整时脉动阵列的能效优势非常明显但当矩阵形状不规则、稀疏度高时它的利用率会急剧下降。这也是为什么后来的AI芯片架构往往采用脉动阵列加其他灵活计算单元的混合方案。理解这个边界比单纯知道“脉动阵列好”要重要得多。2. 拆开一个脉动阵列数据怎么流、计算怎么叠2.1 基本结构PE、连线和节拍一个典型的脉动阵列由二维网格状的处理单元PEProcessing Element组成。每个PE内部只有一个乘法器和一个累加器外加几个寄存器用于暂存数据。PE之间通过水平方向和垂直方向的连线连接数据沿着这些连线逐拍传递。以最经典的N×N脉动阵列做矩阵乘法CA×B为例。矩阵A的元素从左侧流入矩阵B的元素从上方流入。每个PE在每一拍做一件事把来自左边的A元素和来自上方的B元素相乘累加到自己的部分和寄存器里然后把A元素往右传、把B元素往下传。注意A元素在阵列中向右流动B元素向下流动而部分和则沿着对角线方向累积。这里有个关键细节为了让正确的元素在正确的时刻相遇A和B的输入需要做斜向排列skewed input。具体来说A矩阵的第i行需要延迟i拍输入B矩阵的第j列需要延迟j拍输入。这样当A[i][k]和B[k][j]在PE(i,j)相遇时它们恰好是同一拍到达的。这个斜向排列是脉动阵列能正确工作的前提也是硬件实现时最容易被忽略的地方。我见过一些开源的脉动阵列实现在仿真阶段结果总是对不上排查半天发现就是输入没有做skew。skew的本质是补偿数据在阵列中传播的延迟——A[i][k]从左边进入PE(i,j)需要经过j个PE所以它要比A[i][0]晚j拍出发同理B[k][j]从上方进入需要经过i个PE要晚i拍出发。这个逻辑用一句话概括就是让所有需要相乘的元素在同一时刻到达同一个PE。2.2 三种数据流模式权重固定、输出固定、行固定脉动阵列不是只有一种数据流方式。根据哪个数据在阵列中停留、哪个数据在流动可以分为几种典型模式数据流模式固定数据流动数据适用场景特点权重固定WS权重B激活A、部分和推理为主权重预加载后不动适合批量推理输出固定OS部分和C激活A、权重B训练/通用部分和留在PE内累加减少搬运行固定RS激活A的一行权重B、部分和特定卷积一行激活复用多次权重固定模式在推理芯片里最常见。因为推理时权重是固定的可以提前加载到PE的寄存器里然后让输入激活流过阵列。这样权重不需要反复搬运省了大量带宽。Google TPU用的就是权重固定加输出固定的混合方案——权重预加载部分和在PE内累积输入激活从左侧流入。输出固定模式则更适合训练场景。训练时权重需要更新部分和需要反复读写把部分和固定在PE内可以减少对内存的访问。但这也带来一个问题部分和寄存器需要足够的位宽来避免溢出。做INT8推理时部分和通常需要32位做FP16训练时部分和可能需要FP32。这个位宽设计直接影响到PE的面积和功耗。2.3 阵列尺寸怎么定不是越大越好脉动阵列的尺寸比如16×16、32×32、128×128是一个核心设计参数。直觉上阵列越大一次能处理的矩阵越大算力越高。但实际设计中阵列尺寸受限于几个因素第一是面积和功耗。一个128×128的脉动阵列有16384个PE每个PE哪怕只有几百个晶体管总数也是千万级别。而且阵列越大时钟树越复杂布线延迟越难控制。第二是利用率。如果实际矩阵是64×64但阵列是128×128那有一半的PE是闲置的。矩阵越小利用率越低。我见过一些芯片标称算力很高但实际跑小模型时利用率不到30%就是因为阵列尺寸和实际负载不匹配。第三是数据供给带宽。阵列越大每拍需要喂进去的数据越多。一个128×128的阵列每拍需要从左侧输入128个A元素、从上方输入128个B元素。如果内存带宽跟不上PE就会饿死。这就是所谓的内存墙问题——算力上去了但数据供不上。实际设计中阵列尺寸的选择需要在算力、面积、功耗、利用率之间做权衡。常见的做法是用适中的阵列尺寸如32×32或64×64配合高带宽的片上缓存而不是一味追求大阵列。另外很多芯片会支持阵列分区——把大阵列切成几个小阵列分别处理不同的矩阵块提高小矩阵的利用率。3. 脉动阵列在AI芯片里的真实落地形态3.1 TPU的脉动阵列设计细节Google TPU v1是最经典的脉动阵列落地案例。它采用了一个256×256的INT8脉动阵列峰值算力92 TOPS。这个阵列用的是权重固定加输出固定的混合模式权重预加载到PE中激活从左侧流入部分和在PE内累积。TPU v1的PE设计很精简每个PE有一个8位乘法器、一个32位累加器、几个寄存器。控制逻辑极其简单——每个PE只需要知道当前拍是加载权重还是计算以及数据往哪个方向传。这种极简设计让TPU v1的能效比达到了当时GPU的几十倍。但TPU v1也有明显的局限它只支持INT8推理不支持训练。而且256×256的阵列对于小批量推理来说利用率不高。后来的TPU v2/v3增加了浮点支持、支持训练、引入了更多的灵活性比如支持稀疏计算但核心的脉动阵列架构一直保留。我在分析TPU的架构时注意到一个细节TPU的权重加载是逐列进行的而不是一次性全部加载。这是因为256×256的权重矩阵有65536个元素如果一次性加载需要很大的带宽。逐列加载可以让权重加载和计算重叠减少等待时间。这个设计思路在实际芯片实现中很值得借鉴——把大操作拆成小操作让不同阶段重叠起来。3.2 边缘端AI芯片的脉动阵列变体边缘端芯片对功耗和面积的要求比数据中心严格得多所以脉动阵列的设计也需要调整。常见的做法包括缩小阵列尺寸。边缘端常见的阵列尺寸是8×8到32×32而不是数据中心的128×128以上。小阵列面积小、功耗低虽然峰值算力不高但对于边缘端的模型如MobileNet、TinyBERT来说够用。支持可变阵列尺寸。有些边缘芯片支持把大阵列配置成多个小阵列或者动态关闭部分PE来省电。比如一个16×16的阵列可以配置成4个8×8的阵列分别处理不同的计算任务。混合精度支持。边缘端模型往往需要INT8和FP16混合精度。脉动阵列的PE需要支持多种精度的乘加这增加了PE的复杂度但提高了灵活性。与DSP或CPU核的协同。纯脉动阵列只能做矩阵乘法但AI模型里还有激活函数、归一化、池化等操作。边缘芯片通常把脉动阵列和DSP或RISC-V核集成在一起脉动阵列做矩阵乘法其他核做后处理。我实际调试过一款边缘AI芯片它的脉动阵列是16×16的支持INT8和INT16。实测下来跑MobileNet V2时阵列利用率能到70%左右但跑一些自定义的小模型时利用率掉到30%以下。关键问题是模型形状和阵列尺寸的匹配度——如果模型的通道数不是16的倍数阵列就会有闲置。后来我们在模型部署时做了通道对齐把通道数补齐到16的倍数利用率明显提升。3.3 脉动阵列与存内计算的结合趋势最近几年存内计算Compute-in-Memory和脉动阵列的结合成为一个热门方向。传统脉动阵列的权重存在PE的寄存器里而存内计算把权重存在SRAM或RRAM阵列里计算直接在存储阵列里完成。这样省去了权重加载的功耗和延迟能效比可以再上一个台阶。但这种结合也带来新的挑战。存内计算的模拟特性导致计算精度不如数字电路而且编程模型和传统脉动阵列差异很大。目前这个方向还在探索阶段离大规模商用还有距离。不过从架构演进的逻辑来看脉动阵列的核心思想——数据流动、计算复用——在存内计算时代依然适用只是实现方式会发生变化。4. 设计脉动阵列时最容易踩的五个坑4.1 输入skew做错导致结果全错这是最经典的坑。前面提到过A矩阵的第i行需要延迟i拍输入B矩阵的第j列需要延迟j拍输入。如果skew做错了A[i][k]和B[k][j]就不会在PE(i,j)相遇算出来的结果完全是乱的。我见过一个开源实现作者把skew的方向搞反了——A矩阵延迟的是列而不是行。结果仿真出来的矩阵乘法结果看起来“有点像”但完全不对。排查这种问题的方法是先用一个3×3的小矩阵做仿真手动跟踪每一拍每个PE的输入和输出。把每一拍的中间结果打印出来和理论值对比很快就能定位到skew的问题。提示做脉动阵列仿真时建议先用小尺寸如4×4验证功能正确性再扩展到大规模。小尺寸下手动跟踪数据流是可行的大规模下只能靠断言和覆盖率验证。4.2 部分和位宽不够导致溢出部分和的位宽是一个容易被低估的问题。做INT8乘法时两个8位数相乘得到16位结果累加N次后需要log2(N)16位。如果阵列是256×256N256部分和需要16824位。但实际设计中为了保险通常会留更多余量用32位累加器。如果部分和位宽不够累加过程中会溢出结果完全错误。更麻烦的是这种溢出往往在特定输入下才出现仿真时不一定能覆盖到。建议在PE的累加器上加上溢出检测逻辑一旦溢出就置一个标志位方便调试。4.3 阵列利用率被小矩阵拖垮前面提到过阵列尺寸和矩阵形状不匹配会导致利用率下降。但实际中还有一个更隐蔽的问题批量大小batch size的影响。推理时如果batch size1矩阵乘法的另一个维度可能很小阵列的利用率会很低。解决这个问题的方法有几种一是批处理把多个请求攒在一起做推理二是阵列分区把大阵列切成小阵列处理小矩阵三是数据重排把多个小矩阵拼成一个大矩阵。具体用哪种方法取决于实际负载的特征。4.4 权重加载带宽成为瓶颈权重固定模式下权重需要预加载到PE中。如果权重矩阵很大加载时间可能很长影响整体性能。比如一个256×256的权重矩阵如果每个周期只能加载一行256个元素需要256个周期才能加载完。在这256个周期里阵列无法做计算。解决方法是权重加载和计算重叠。比如在计算当前矩阵块的同时预加载下一个矩阵块的权重。这需要双缓冲的权重寄存器增加了面积但提高了吞吐。另一种方法是权重压缩利用权重的稀疏性减少加载量。4.5 时钟频率上不去脉动阵列的时钟频率受限于PE之间的连线延迟。阵列越大最远两个PE之间的连线越长延迟越大时钟频率就越低。一个128×128的阵列如果每个PE的延迟是100ps最远路径的延迟可能超过10ns时钟频率只能跑到100MHz左右。提高时钟频率的方法包括插入流水线寄存器把长连线切成短段、优化PE布局让相邻PE的连线最短、使用更先进的工艺。但这些方法都有代价——流水线寄存器增加面积和延迟先进工艺增加成本。实际设计中需要在频率、面积、功耗之间做权衡。5. 从架构到落地脉动阵列的编程与验证5.1 怎么把模型映射到脉动阵列上把神经网络模型映射到脉动阵列上核心工作是矩阵分块tiling。一个大矩阵乘法需要切成多个小块逐块送入阵列计算。分块的大小取决于阵列尺寸和片上缓存大小。以一个全连接层为例权重矩阵是1024×1024输入是1024×1输出是1024×1。如果阵列是32×32那么权重矩阵需要切成32×32的块共32×321024个块。每个块计算完后部分和需要累加。这里的关键是部分和的存储和累加顺序——如果部分和存在片上缓存里需要足够的缓存来存放中间结果如果存在片外内存里带宽会成为瓶颈。实际映射时还需要考虑数据复用。输入向量在计算不同输出元素时会被复用权重矩阵在计算不同batch时会被复用。好的映射策略会尽量把可复用的数据留在片上减少片外访问。我实际做模型映射时的经验是先用一个简单的映射策略跑通再逐步优化。一开始不要追求极致的利用率先把功能做对。功能对了之后再通过调整分块大小、调整数据流顺序、增加双缓冲等方式提升性能。5.2 功能验证从单元测试到系统测试脉动阵列的验证是一个分层的过程PE级验证。先验证单个PE的功能——乘法、累加、数据传递是否正确。PE的测试用例应该覆盖各种输入组合包括边界值最大正数、最小负数、零。阵列级验证。把多个PE连成小阵列如4×4验证数据流是否正确、skew是否正确、部分和是否正确。这个阶段可以用随机矩阵做测试对比软件计算结果。系统级验证。把脉动阵列集成到完整芯片中验证与内存、DMA、控制器的交互。这个阶段需要跑真实的模型对比端到端的精度和性能。验证中最容易漏掉的是边界情况矩阵尺寸不是阵列尺寸整数倍时怎么处理输入数据有异常值如NaN、Inf时怎么处理这些情况在实际部署中一定会遇到验证时必须覆盖。5.3 性能调优从理论峰值到实际利用率脉动阵列的理论峰值算力很容易算PE数量×时钟频率×2乘加各算一次。但实际利用率往往远低于100%。影响利用率的因素包括数据供给带宽如果数据供不上PE会空闲矩阵形状不规则的矩阵导致PE闲置部分和搬运部分和读写占用时间权重加载权重加载期间阵列无法计算提升利用率的方法包括增加片上缓存减少片外访问、优化数据流让数据复用最大化、支持稀疏计算跳过零元素、动态阵列分区适应不同矩阵尺寸。我实测过一个64×64的脉动阵列理论峰值是64×64×2×500MHz4 TOPS。跑一个典型的卷积层时实际利用率在60%到75%之间波动。瓶颈主要在数据供给——片上缓存不够大部分数据需要从片外读取。后来把片上缓存加大了一倍利用率提升到了85%左右。6. 脉动阵列的边界与未来演进方向脉动阵列不是万能的。它的优势在于规整的矩阵乘法但对于非矩阵运算如激活函数、归一化、排序效率很低。而且它对矩阵形状敏感小矩阵和稀疏矩阵的利用率不高。这些边界决定了脉动阵列在AI芯片中的定位——它是矩阵计算的核心加速器但需要其他计算单元配合完成整个模型的计算。从演进方向来看几个趋势值得关注稀疏脉动阵列。利用权重的稀疏性跳过零元素的乘加可以大幅提升有效算力。但稀疏性会破坏数据流的规整性需要额外的控制逻辑。可重构脉动阵列。支持多种数据流模式根据不同的计算任务动态配置。这增加了灵活性但也增加了面积和功耗。三维脉动阵列。把二维阵列扩展到三维增加数据复用的维度。但三维集成的工艺复杂度和散热问题需要解决。与近存计算的结合。把部分权重存在离PE更近的存储里减少数据搬运。这需要新的存储器件和电路设计。我在实际项目中体会最深的一点是脉动阵列的设计没有“最优解”只有“最适合当前负载的解”。数据中心的大阵列、边缘端的小阵列、支持稀疏的变体阵列各有各的适用场景。理解负载特征比追求架构的“先进”更重要。
返回列表