从矩阵乘法到脉动阵列：AI芯片软硬件协同设计核心逻辑

1. 从矩阵乘法到脉动阵列AI芯片软硬件协同设计的核心逻辑搞AI芯片设计的人都有一个共识算力不是堆出来的是“流”出来的。这个“流”字说的就是数据在计算单元之间的流动方式。我接触过不少做芯片架构的朋友大家一开始都容易陷入一个误区——觉得只要把MAC乘累加单元阵列铺得足够大峰值算力标得足够高芯片就能跑得快。实际一上板测试发现跑Transformer模型的时候利用率连30%都不到功耗还高得离谱。问题出在哪就出在数据搬运上。AI芯片的软硬件设计本质上是在解决一个矛盾计算单元想要数据源源不断地送过来而存储器和总线带宽是有限的。你阵列越大这个矛盾就越尖锐。脉动阵列Systolic Array之所以成为TPU等主流AI加速器的核心架构就是因为它用一种极其优雅的方式缓解了这个矛盾——让数据像心跳一样有节奏地在阵列中流动每个周期只移动一跳计算和搬运重叠进行。这一篇我打算把脉动阵列的基本原理、它跟Transformer里矩阵乘法的映射关系、以及软硬件设计上那些“教科书不会写但实际会踩”的坑从头到尾捋一遍。适合谁看如果你正在做AI加速器的架构设计、算子映射、或者用FPGA/ASIC实现Transformer推理这篇内容应该能帮你少走一些弯路。如果你只是对AI芯片怎么跑Transformer感兴趣我也会尽量用生活化的类比把原理讲清楚。2. 脉动阵列基本原理为什么数据要“跳着走”2.1 从矩阵乘法说起AI计算的绝对核心Transformer也好CNN也好说到底绝大部分计算量都落在矩阵乘法上。一个注意力头做一次QK^T就是两个矩阵相乘FFN层两个全连接也是矩阵乘法。所以AI芯片设计的第一性问题就是怎么高效地做矩阵乘法。矩阵乘法C A × B假设A是M×KB是K×N那C就是M×N。朴素实现需要M×N×K次乘累加每次乘累加都要从内存里读A的元素和B的元素。如果A和B都放在片外DRAM那数据搬运量是M×K K×N计算量是M×N×K。计算访存比大概是O(min(M,N,K))。当矩阵维度大的时候这个比值很高理论上计算密集但实际中因为缓存容量有限数据反复换入换出有效带宽就成了瓶颈。我经常用一个类比来解释这件事矩阵乘法就像做一桌菜计算单元是厨师数据是食材。如果每个厨师做每道菜都要自己去仓库拿食材那大部分时间都花在路上了。脉动阵列的思路是让食材在厨师之间传递每个厨师从上一个厨师手里接过食材处理完再传给下一个。这样仓库只需要在开头送一次食材后面就是流水线作业。2.2 脉动阵列的核心机制节奏同步的数据流动脉动阵列这个名字来源于“systole”也就是心脏收缩的意思。它的核心特征是数据以固定的节拍在PEProcessing Element阵列中流动每个PE在每个时钟周期从相邻PE接收数据、执行计算、再把结果传给下一个PE。一个典型的N×N脉动阵列做矩阵乘法时矩阵A的元素从左侧流入矩阵B的元素从上方流入两者在PE中相遇并完成乘累加。关键在于A的每一行元素进入阵列的时间是错开的——第i行要延迟i个周期进入。B的每一列也类似第j列延迟j个周期。这样当A[i][k]和B[k][j]在PE(i,j)相遇时正好是它们应该相乘的时刻。这个“错开”的设计是脉动阵列的精髓。它不需要全局广播不需要复杂的地址译码每个PE只需要跟左右邻居打交道。数据复用率极高A的每个元素进入阵列后会被N个PE依次使用B的每个元素同理。这就是为什么脉动阵列的能效比通常远高于基于总线的架构。2.3 权重固定 vs 输出固定两种数据流的取舍实际设计脉动阵列时数据流的选择会直接影响硬件复杂度和利用率。常见的有两种权重固定Weight Stationary每个PE里存一个权重值输入数据在阵列中流动。这种方案适合权重可以提前加载、且推理时权重不变的场景。Transformer的FFN层权重在推理时是固定的用这种方案很合适。缺点是权重加载需要额外的周期且如果权重矩阵太大阵列装不下就要分块。输出固定Output Stationary每个PE负责累加一个输出元素输入数据从两个方向流入。这种方案适合卷积神经网络中卷积核滑动的场景因为输出特征图的每个点需要累加多个乘积。在Transformer的注意力计算中QK^T的结果需要经过softmax再跟V相乘输出固定的方案可以减少中间结果的搬运。我实测下来的经验是对于Transformer类模型混合方案往往最优——QK^T用输出固定FFN用权重固定。但这取决于你的阵列大小和片上缓存容量没有银弹。3. Transformer在脉动阵列上的映射从注意力到FFN3.1 注意力机制的矩阵乘法拆解Transformer的注意力计算可以拆成几个矩阵乘法Q X × W_QK X × W_KV X × W_VS Q × K^T / sqrt(d_k)P softmax(S)O P × V其中第1步和第4步是标准的矩阵乘法可以直接映射到脉动阵列。第2步Q×K^T稍微特殊因为K需要转置。在硬件上转置可以通过改变数据流入方向来实现——让K从上方流入时按列顺序而不是按行顺序。第3步softmax是非线性操作通常在单独的向量单元里做不占用脉动阵列。这里有个设计要点softmax需要知道整行的最大值来做数值稳定所以S矩阵的一行必须完整算完才能做softmax。这意味着脉动阵列算完一行后要暂停等softmax处理完再继续。这个“气泡”会降低阵列利用率。我的做法是把softmax和后续的P×V做流水线重叠。当脉动阵列在算S的第i1行时向量单元在处理第i行的softmax。这样气泡就被隐藏了。实现上需要双缓冲的S矩阵存储面积换效率值得。3.2 多头注意力的并行化策略多头注意力有h个头每个头独立做注意力计算。在脉动阵列上有两种并行方式头间并行把阵列分成h个子阵列每个子阵列处理一个头。优点是控制简单缺点是每个子阵列规模小数据复用率降低。头内并行所有头共享同一个大阵列按时间分片处理。优点是阵列利用率高缺点是需要存储所有头的Q、K、V片上缓存压力大。我倾向于头内并行因为Transformer的注意力头数通常不多8到16个分阵列会导致每个阵列太小。但如果你做的是边缘端芯片阵列本身就不大那头间并行反而更实际。3.3 FFN层的映射与优化FFN层是两个全连接先升维再降维。以BERT-base为例hidden size是768 intermediate size是3072。第一个全连接是768×3072的权重矩阵第二个是3072×768。在脉动阵列上第一个全连接的输出是3072维如果阵列是256×256那需要分12个块来计算。每个块算完后结果要经过GELU激活再送入第二个全连接。这里的关键是GELU是非线性的必须在两个全连接之间插入向量单元。我踩过的一个坑是GELU的计算延迟没有隐藏好导致脉动阵列在等激活结果时闲置。后来改成把GELU的计算拆成多个小段跟矩阵乘法的尾段重叠利用率从62%提升到了81%。具体做法是当脉动阵列算完输出块的前半部分时向量单元就开始对这部分做GELU不用等整个块算完。4. 软硬件协同设计的关键细节4.1 数据位宽的选择精度与面积的平衡AI芯片设计里数据位宽是个绕不开的决策。FP32精度高但面积大、功耗高INT8面积小但精度损失可能影响模型准确率。Transformer模型对精度的敏感度因层而异注意力层的QK^T对精度比较敏感FFN层相对鲁棒。我的经验是采用混合精度QK^T用FP16或BF16FFN用INT8。这样在保持模型准确率的同时把FFN的面积和功耗降下来。具体实现上脉动阵列的PE需要支持可配置的位宽——乘法器用FP16和INT8双模累加器用FP32保证累加精度。这里有个计算假设阵列是256×256FP16乘法器面积约是INT8的4倍。如果全部用FP16阵列面积会很大。混合精度后FFN部分的PE可以缩小整体面积能省30%左右。4.2 片上缓存的层次设计脉动阵列的利用率高度依赖数据供给。如果数据供不上阵列再大也是白搭。片上缓存通常分三级寄存器文件每个PE内部的寄存器存权重或部分和。阵列级缓存紧挨着阵列的SRAM存当前计算块的输入数据。全局缓存芯片级的SRAM存整个层的输入输出。关键参数是阵列级缓存的容量。太小了数据供不上太大了面积浪费。我的估算方法是阵列级缓存至少要能存下阵列一次计算所需的所有输入数据。对于N×N阵列做K维的矩阵乘法需要存N×K的A块和K×N的B块。如果N256K768那需要256×768×2字节×2 ≈ 786KB。这个容量不小所以实际中K会分块比如K64那缓存需求就降到64KB左右。4.3 控制逻辑的简化让数据自己“找路”脉动阵列的一个巨大优势是控制简单。因为数据流动是规整的不需要复杂的地址生成和仲裁逻辑。每个PE只需要知道什么时候从左边收数据、什么时候从上边收数据、什么时候把结果往右/往下传。但简化控制的前提是数据排布要配合好。如果输入数据的排布跟阵列的数据流不匹配就需要额外的重排逻辑反而增加复杂度。我的做法是在编译器层面就把数据排布做好让硬件只管流动。编译器根据阵列大小和数据流方案自动生成数据的加载顺序和地址序列。这样硬件控制逻辑可以做到极简面积和功耗都省。5. 实操过程从零搭建一个脉动阵列模拟器5.1 环境准备与工具选型要验证脉动阵列的设计不一定非要流片。用Python或C写一个周期精确的模拟器就能验证数据流、计算正确性和利用率。我推荐用Python因为开发快配合NumPy做参考模型很方便。需要准备的环境Python 3.8NumPy用于参考矩阵乘法Matplotlib用于可视化数据流和利用率如果你要做RTL级验证那就需要Verilog/VHDL和仿真器。但前期架构探索Python模拟器足够了。5.2 脉动阵列模拟器的核心代码下面是一个权重固定的脉动阵列模拟器核心逻辑。假设阵列是N×N权重预加载到PE中输入数据从左侧流入部分和从上方流入、下方流出。import numpy as np class PE: def __init__(self): self.weight 0 self.partial_sum 0 def compute(self, input_data): self.partial_sum input_data * self.weight return self.partial_sum class SystolicArray: def __init__(self, size): self.size size self.pe_array [[PE() for _ in range(size)] for _ in range(size)] def load_weights(self, weight_matrix): # weight_matrix: N x N for i in range(self.size): for j in range(self.size): self.pe_array[i][j].weight weight_matrix[i][j] def run(self, input_matrix): # input_matrix: M x N, M是时间步数 M input_matrix.shape[0] results [] for t in range(M self.size - 1): # 每个周期数据从左侧流入 for i in range(self.size): if t - i 0 and t - i M: data input_matrix[t - i][i] if i input_matrix.shape[1] else 0 else: data 0 # 数据在行内向右流动 for j in range(self.size): if j 0: self.pe_array[i][j].compute(data) else: self.pe_array[i][j].compute(self.pe_array[i][j-1].partial_sum) # 收集最后一列的部分和 if t self.size - 1: result_row [self.pe_array[i][self.size-1].partial_sum for i in range(self.size)] results.append(result_row) # 重置部分和 for i in range(self.size): for j in range(self.size): self.pe_array[i][j].partial_sum 0 return np.array(results)这段代码是简化版实际模拟器还需要处理数据加载的时序、部分和的累加清零、以及多块矩阵乘法的拼接。但核心逻辑就是每个周期数据在阵列中移动一跳PE完成一次乘累加。5.3 用模拟器验证Transformer的注意力计算拿一个简单的注意力计算来验证假设hidden size4头数1序列长度3。Q、K、V都是3×4的矩阵。QK^T是3×3再跟V相乘得到3×4的输出。用上面的模拟器先把K^T作为权重加载到阵列4×4的阵列实际用3×3然后让Q从左侧流入。跑完后得到S矩阵。再手动做softmax然后把P作为输入V作为权重再跑一次得到输出。对比NumPy的直接计算结果如果一致说明模拟器逻辑正确。我实测下来这个模拟器跑一个完整的BERT-base注意力层序列长度128hidden size 768大概需要几秒钟足够做架构探索了。5.4 利用率分析与瓶颈定位模拟器跑起来后可以统计每个周期的PE活跃数。理想情况下稳态时所有PE都在工作。但实际中数据加载、softmax、GELU等阶段会有气泡。我通常关注三个指标峰值利用率稳态时的PE活跃比例反映阵列设计的上限。平均利用率整个计算过程的PE活跃比例反映实际效率。气泡占比非计算周期占总周期的比例反映控制逻辑的优化空间。如果平均利用率低于50%那就要检查是数据供给跟不上还是非线性操作太频繁还是分块策略不合理。我的经验是Transformer类模型在脉动阵列上的平均利用率做到70%以上就算不错了再往上每提升一个点都需要精细的软硬件协同优化。6. 常见问题与排查技巧实录6.1 阵列利用率上不去先查这三个地方问题一数据加载带宽不足。脉动阵列算得快但数据从DRAM搬到片上缓存的速度跟不上。排查方法算一下计算访存比如果每字节数据只做了不到10次乘累加那大概率是带宽瓶颈。解决方法是增加片上缓存容量或者用数据压缩减少搬运量。问题二分块策略不合理。矩阵太大阵列装不下需要分块。如果分块后每块的K维太小那数据复用率就低。比如K16那每个数据元素只被用了16次远低于理论值。解决方法是增大K维分块或者用双缓冲隐藏加载延迟。问题三非线性操作打断流水线。softmax、GELU、LayerNorm这些操作如果串行执行会让阵列闲置。解决方法是把这些操作跟矩阵乘法重叠用流水线隐藏延迟。6.2 精度损失排查从FP32到INT8的坑量化到INT8后模型准确率下降是常见问题。排查思路先看是哪一层对精度最敏感。通常注意力层的QK^T最敏感FFN层最鲁棒。对敏感层保留FP16只对鲁棒层用INT8。如果INT8后准确率还是掉检查量化校准集是否覆盖了实际数据分布。校准集太小或分布偏差大会导致量化参数不准。我踩过的一个坑是量化校准用了随机数据结果实际推理时准确率暴跌。后来换成用真实数据集的子集做校准准确率就恢复了。6.3 时序收敛问题脉动阵列的时钟频率上不去脉动阵列的时钟频率受限于PE内部的乘法器延迟和布线延迟。如果频率上不去先检查乘法器是不是组合逻辑太长可以插入流水线寄存器把乘法拆成两级。阵列的布线是不是太长可以缩小阵列规模或者用H-tree结构做时钟分布。部分和的累加路径是不是关键路径可以把累加器拆成两级先算部分和再累加。我的经验是256×256的脉动阵列在16nm工艺下跑1GHz左右比较现实。再往上就要牺牲面积或功耗。6.4 常见问题速查表问题现象可能原因排查方法解决思路阵列利用率低于50%数据带宽不足计算访存比增大片上缓存数据压缩精度下降超过1%量化位宽不够逐层敏感度分析混合精度敏感层用FP16时钟频率上不去关键路径太长时序分析插入流水线缩小阵列功耗高于预期数据搬运太频繁功耗分析增加数据复用降低电压面积超标阵列太大或缓存太大面积报告分时复用缩小阵列7. 从架构到部署软硬件协同的完整链路7.1 编译器在软硬件协同中的角色脉动阵列的硬件是固定的但Transformer模型千变万化。怎么让硬件高效执行不同的模型靠编译器。编译器的任务是把模型的计算图映射到阵列上决定数据怎么分块、怎么排布、怎么调度。一个好的AI芯片编译器需要做几件事算子融合把多个小算子合并成一个大算子减少数据搬运。比如把QK^T、softmax、P×V融合成一个注意力算子。内存分配决定每个中间结果存在哪里什么时候释放。目标是让片上缓存刚好够用不多不少。流水线调度把矩阵乘法和非线性操作重叠隐藏延迟。我见过不少团队硬件做得好但编译器拉胯导致实际性能只有峰值的20%。软硬件协同设计编译器至少占一半的功劳。7.2 部署时的性能调优经验模型部署到AI芯片上后性能调优是个细活。我的经验是批处理大小批处理太小阵列利用率低批处理太大延迟高。找到平衡点很关键。对于Transformer推理批处理1到8通常比较合适。序列长度长序列的注意力计算是O(n^2)的序列长度翻倍计算量翻四倍。如果芯片的阵列不够大长序列会严重拖慢速度。可以考虑用稀疏注意力或分块注意力来降低计算量。权重预加载如果模型权重不变可以在初始化时一次性加载到片上避免每次推理都重新加载。这能省不少时间和功耗。7.3 一个实际案例BERT-base在脉动阵列上的部署拿BERT-base举例12层hidden size 76812个头序列长度128。在256×256的脉动阵列上部署每层的QKV投影3个768×768的矩阵乘法每个需要分块成768/2563个K维块共9个块。注意力计算12个头每个头64维QK^T是128×64×128P×V是128×128×64。FFN768×3072和3072×768各需要12个K维块。总共算下来一层大概需要几百个阵列周期。12层就是几千个周期。如果阵列跑1GHz那一次推理大概几微秒。加上数据加载和softmax的开销实际可能在几十微秒量级。这个性能对于实时推理够用了但还有优化空间。比如把QKV投影合并成一个大矩阵乘法减少分块开销或者用Winograd算法减少乘法次数。这些优化需要编译器和硬件配合不是单方面能解决的。8. 一些个人体会做AI芯片的软硬件设计最深的体会是不要孤立地看硬件或软件。脉动阵列的硬件设计再精妙如果编译器映射不好性能也出不来。反过来编译器再聪明硬件不支持某些数据流也白搭。我刚开始做这块的时候花了很多时间在硬件微架构上觉得把PE做小、把阵列做大就行。后来发现真正决定性能的是数据怎么流动、怎么复用。脉动阵列之所以经典就是因为它把数据流动这件事做到了极致——每个数据元素都被充分利用每个周期都在做有用的计算。如果你也在做类似的设计我的建议是先写一个周期精确的模拟器把数据流和利用率搞清楚再动手写RTL。模拟器上跑通了RTL实现就是体力活。模拟器上没跑通RTL调死也调不出来。另外Transformer模型还在快速演进新的变体层出不穷。脉动阵列的架构要有一定的灵活性能适应不同的矩阵维度和数据流。太专用的设计可能明年就过时了。保持可配置、可扩展比追求极致的峰值算力更重要。