ARTICLE DETAIL

资讯详情

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

AI芯片软硬件协同设计:脉动阵列与Transformer映射实战

AI芯片软硬件协同设计:脉动阵列与Transformer映射实战 1. AI 芯片软硬件协同设计的核心命题1.1 为什么单看硬件或软件都不够做 AI 芯片这行的人有个共识硬件和软件分开设计最后大概率要返工。我见过太多团队硬件组把脉动阵列的面积和频率优化到极致软件组把算子调度写得漂漂亮亮结果一联调发现数据搬运成了瓶颈算力利用率连 30% 都不到。这不是谁的水平问题是协同设计这件事本身没做透。AI 芯片和传统 CPU、GPU 最大的区别在于它的计算模式高度可预测。Transformer 里的矩阵乘法、卷积里的滑窗累加这些操作的访存规律是相对固定的。正因为可预测才有机会把硬件结构做得极度专用化也正因为专用化软件必须提前知道硬件的脾气否则写出来的调度根本喂不饱计算单元。这一篇我聚焦在**脉动阵列Systolic Array**这个结构上讲清楚它为什么适合做矩阵乘法、Transformer 这类负载在它上面怎么映射、软硬件接口怎么设计以及实际调试中会踩哪些坑。适合做芯片架构的、做编译器后端的、以及想搞明白 AI 加速器到底怎么回事的读者。1.2 从矩阵乘法说起AI 计算的本质先把问题简化。Transformer 的核心是什么注意力机制。注意力机制的核心是什么三个矩阵乘法Q 乘 K 的转置、softmax 之后乘 V、以及前面的线性投影。再往外看前馈网络 FFN 也是两个大矩阵乘法。整个模型跑一遍90% 以上的计算量都压在矩阵乘法上。矩阵乘法为什么这么重要因为它有极高的计算密度和数据复用潜力。一个 M×K 的矩阵乘 K×N 的矩阵输出 M×N总共 2×M×N×K 次浮点运算。如果 M、N、K 都是 1024那就是 20 亿次运算。而输入数据只有 2×1024×1024 个元素。也就是说每个输入元素平均要被复用 1024 次。这种复用潜力正是专用硬件能大幅超越通用处理器的根本原因。但复用不是自动发生的。通用 CPU 靠缓存层级来碰运气GPU 靠大量线程和共享内存来手工管理。脉动阵列的思路更激进把复用直接做进数据流里让数据在计算单元之间流动时自然被多次使用不需要反复从内存里取。2. 脉动阵列结构、原理与设计取舍2.1 脉动阵列到底脉动在哪里脉动阵列这个名字来自它的工作方式数据像心跳一样有节奏地从一个计算单元流向相邻的计算单元。每个计算单元PEProcessing Element只做一件小事——乘加运算然后把结果往下一个单元传。想象一个二维网格每个格子是一个 PE。矩阵 A 的元素从左侧一行行流入矩阵 B 的元素从上方一列列流入。每个 PE 在某个时刻收到一个 A 的元素和一个 B 的元素做一次乘加把 A 往右传、把 B 往下传。经过若干节拍后结果矩阵 C 的元素就会从底部或右侧流出。这种结构的精妙之处在于每个数据元素只从内存读取一次然后在阵列内部被复用多次。对于 N×N 的阵列一个输入元素最多可以被复用 N 次。这直接减少了对内存带宽的需求而内存带宽恰恰是 AI 芯片最稀缺的资源。2.2 为什么是脉动阵列而不是别的做矩阵乘法的硬件结构有很多选择。最直接的是一维乘加树把所有乘法结果加起来。但这样每个周期都要读大量数据带宽压力巨大。另一种是广播结构把数据广播给所有 PE但广播线的扇出和功耗随规模增长很快。脉动阵列的优势在于它的局部连接性。每个 PE 只和相邻的 PE 通信连线短、延迟低、功耗小。而且数据流动是规则的控制逻辑可以做得非常简单。TPU v1 就是靠这个结构在 28nm 工艺下做到了 92 TOPS 的峰值算力功耗只有 40W。但脉动阵列不是没有代价。它的规则性意味着灵活性差。遇到非矩阵乘法的操作比如 softmax、LayerNorm、激活函数脉动阵列就帮不上忙得靠旁边的向量单元来处理。而且脉动阵列对矩阵的维度很敏感如果矩阵尺寸和阵列尺寸不匹配利用率会掉得很厉害。2.3 阵列尺寸怎么定一个实际的权衡假设我们要设计一个 N×N 的脉动阵列N 取多大合适从算力角度看N 越大峰值算力越高。N×N 个 PE每个 PE 每周期做一次乘加峰值就是 2×N² 次浮点运算每周期。N128 时1GHz 频率下峰值约 32 TFLOPS。但 N 越大三个问题越突出。第一是面积PE 数量按 N² 增长N 翻倍面积翻四倍。第二是利用率实际矩阵的维度往往不是 N 的整数倍边缘填充会浪费算力。第三是延迟数据穿过整个阵列需要 2N 个周期N 太大时流水线填充和排空的代价不可忽略。我的经验是N 取 64 到 256 之间比较合理。具体取值要看目标负载的矩阵维度分布。如果主要跑 Transformer 的注意力序列长度通常在 128 到 2048 之间N128 或 256 能覆盖大部分情况。如果做边缘推理矩阵小N32 或 64 更划算。注意阵列尺寸不是越大越好。我见过一个团队为了追求峰值算力指标把阵列做到 512×512结果实际负载下利用率只有 15%功耗还高得离谱。峰值算力是给外行看的有效算力才是给自己用的。3. Transformer 在脉动阵列上的映射3.1 注意力机制的矩阵乘法拆解要把 Transformer 映射到脉动阵列上先得把它的计算图拆成矩阵乘法的形式。以单头注意力为例输入 X 的形状是 [seq_len, d_model]。经过三个线性投影得到 Q、K、V形状都是 [seq_len, d_k]。然后计算 attention scoreS Q × K^T形状 [seq_len, seq_len]。接着 softmax 得到 P再算输出 O P × V形状 [seq_len, d_k]。这里面有两个大的矩阵乘法。第一个是 Q×K^T计算量是 2×seq_len²×d_k。第二个是 P×V计算量是 2×seq_len²×d_k。当 seq_len 很大时这两步的计算量占主导。映射到脉动阵列时Q×K^T 可以这样安排Q 从左侧流入K^T 从上方流入。但 K^T 不是现成的K 是 [seq_len, d_k]需要转置。转置这个操作在硬件上要么靠额外的转置单元要么靠改变数据流入的顺序。实际设计中通常让 K 按列流入等效于 K^T 按行流入省掉显式转置。3.2 数据流设计权重 stationary 还是输出 stationary脉动阵列有三种经典的数据流权重固定Weight Stationary、输出固定Output Stationary、行固定Row Stationary。权重固定是把权重预加载到 PE 里输入数据流过时直接乘。这种方式适合权重复用率高的场景比如卷积。但 Transformer 的权重是动态的Q、K、V 都是输入算出来的没法预加载。输出固定是让部分和留在 PE 里累加输入和权重都流动。这种方式适合输出维度大的场景。Transformer 的注意力矩阵 seq_len×seq_len 通常很大输出固定的数据流比较合适。实际设计中很多加速器采用混合策略。比如 Google 的 TPU 用的是权重固定加输出固定的变体而一些学术界的加速器如 Eyeriss 用的是行固定。选择哪种数据流取决于目标负载的访存模式和复用机会。3.3 一个具体的映射示例假设 seq_len128d_k64阵列是 128×128。第一步算 S Q×K^T。把 Q 的每一行作为一个向量从阵列左侧流入。K 的每一行作为向量从上方流入。经过 128 个周期后阵列的每一列输出一个 S 的元素。但这里有个问题S 是 128×128而阵列也是 128×128刚好匹配。如果 seq_len 是 100就得填充到 128浪费 28% 的算力。第二步算 O P×V。P 是 softmax 的结果需要先算出来存到缓冲区再从左侧流入。V 从上方流入。这一步的输出是 128×64只用了阵列的一半。整个流程下来阵列的利用率大概在 60% 到 70% 之间。剩下的算力浪费在填充、softmax、LayerNorm 这些非矩阵乘法操作上。这也是为什么实际芯片的有效算力往往只有峰值的几分之一。4. 软硬件接口设计编译器要看到什么4.1 指令集该暴露多少硬件细节软硬件协同设计的一个核心问题是指令集应该暴露多少硬件细节给编译器暴露太少编译器没法做优化生成的代码效率低。暴露太多编译器复杂度爆炸而且硬件升级后软件不兼容。我的观点是指令集应该暴露数据流的方向和粒度但不暴露具体的 PE 阵列尺寸。比如可以提供矩阵乘法指令指定输入矩阵的地址、形状、数据流方向但不告诉编译器阵列是 128×128 还是 256×256。这样编译器只需要关心矩阵乘法的逻辑硬件可以自由调整阵列规模。但这里有个矛盾如果编译器不知道阵列尺寸就没法做分块tiling。分块是矩阵乘法优化的关键块大小选得不对数据复用率会差很多。所以实际设计中通常会在指令集里暴露一个推荐块大小的参数或者提供一个查询接口让编译器问硬件。4.2 数据搬运的调度策略脉动阵列的算力再强数据喂不上也是白搭。数据搬运的调度是软硬件协同设计里最容易被低估的部分。考虑一个场景要算一个 1024×1024 的矩阵乘法阵列是 128×128。需要把大矩阵切成 8×8 个小块每个小块 128×128。如果按行优先的顺序搬运每次搬一块 A 和一块 B算完再搬下一块。但这样 A 的每一行块会被复用 8 次B 的每一列块也会被复用 8 次。如果能在片上缓存里留住这些块就能减少 8 倍的内存访问。这就是分块加缓存的策略。编译器需要知道片上缓存有多大、有几级、每级的带宽是多少。这些信息通常通过一个硬件描述文件暴露给编译器而不是硬编码在指令集里。4.3 一个实际的编译器流程以一个 Transformer 层为例编译器大概要做这些事图优化把计算图里的算子融合比如把 LayerNorm 和线性投影融合减少中间结果的写回。算子映射把矩阵乘法映射到脉动阵列把 softmax 映射到向量单元。分块决策根据矩阵形状和片上缓存大小决定分块策略。调度生成生成数据搬运和计算的时序尽量让计算和搬运重叠。指令生成把调度翻译成硬件指令。这个流程里第 3 步和第 4 步最依赖硬件信息。如果硬件提供了灵活的 DMA 和双缓冲机制编译器就能做更激进的流水线优化。如果硬件只有简单的同步搬运编译器就只能保守调度。5. 实操中的坑与排查技巧5.1 利用率上不去的常见原因实际调试中脉动阵列利用率低是最常见的问题。我整理了一个排查表现象可能原因排查方法利用率低于 30%矩阵维度与阵列不匹配检查实际矩阵形状分布看填充浪费多少利用率波动大数据搬运成为瓶颈用性能计数器看 DMA 带宽和计算单元的等待周期小矩阵利用率极低流水线填充排空开销大看阵列深度考虑小矩阵走向量单元大矩阵利用率也低片上缓存不够反复读内存检查分块策略看缓存命中率我踩过最坑的一次是矩阵维度明明是 128 的整数倍利用率却只有 40%。查了半天发现是数据在内存里的排布不是连续的DMA 每次只能搬一小段带宽利用率极低。后来改了数据排布格式利用率直接上到 75%。5.2 数值精度问题AI 芯片通常用低精度计算比如 FP16 或 INT8。但低精度会带来数值问题尤其是在 Transformer 里。注意力分数经过 softmax 之前数值范围可能很大。如果用 FP16指数运算容易溢出。常见的做法是在 softmax 之前减去最大值但这需要额外的归约操作。如果硬件不支持快速归约这一步会成为瓶颈。另一个坑是累加精度。脉动阵列里的部分和如果也用 FP16 累加误差会累积。通常的做法是部分和用 FP32最后再转回 FP16。但这会增加 PE 的面积和功耗。怎么权衡要看目标模型的精度要求。5.3 调试工具链的搭建没有好的调试工具软硬件协同设计就是盲人摸象。我建议至少搭这三样第一是周期级模拟器。能精确模拟每个周期的数据流动看到每个 PE 在干什么。这对定位流水线冲突和气泡非常有用。第二是性能计数器。在硬件里埋点统计计算单元利用率、DMA 带宽、缓存命中率等指标。这些数据是优化的依据。第三是可视化工具。把矩阵乘法的数据流画出来直观看到哪些 PE 在忙、哪些在等。人眼对图形的敏感度远高于数字。提示调试工具链的搭建要趁早。我见过团队硬件都快流片了才开始搭调试环境结果发现一堆问题但没法定位只能靠猜。6. 从 Transformer 反推硬件设计6.1 模型趋势对硬件的影响Transformer 的架构还在快速演进。从最初的 encoder-decoder到 GPT 系列的 decoder-only再到各种稀疏注意力、线性注意力计算模式一直在变。这对硬件设计提出了挑战如果硬件过于专用模型一变就废了。比如如果硬件专门为稠密注意力优化遇到稀疏注意力就发挥不出来。如果专门为固定序列长度优化遇到变长序列就抓瞎。我的看法是硬件应该抓住不变的部分把变化的部分留给软件。矩阵乘法是不变的所以脉动阵列值得做。但注意力的具体形式会变所以不要把注意力机制硬编码进硬件而是提供灵活的向量单元和可编程的数据搬运。6.2 可重构性的代价为了提高灵活性有些加速器引入了可重构的数据通路。比如PE 之间的连接可以通过配置改变支持不同的数据流。但可重构性是有代价的。可重构的开关会增加面积和功耗配置本身也需要时间和存储。我算过一笔账一个 128×128 的阵列如果每个 PE 有 4 个可配置的输入来源光配置位就是 128×128×232768 位。每次切换数据流都要重新配置这个开销在小矩阵场景下可能超过计算本身。所以可重构性要用在刀刃上。我的建议是只在数据流的顶层做可重构比如支持权重固定和输出固定两种模式但 PE 内部的微架构保持固定。这样既能覆盖主要场景又不会让面积和功耗失控。6.3 一个实际的架构选择如果让我现在设计一个面向 Transformer 推理的 AI 芯片我会这么选计算核心用 128×128 的脉动阵列支持 FP16 输入和 FP32 累加。旁边配一个 512 宽的向量单元处理 softmax、LayerNorm、激活函数。片上缓存用 4MB 的 SRAM分成两级一级靠近阵列做双缓冲二级做全局共享。DMA 支持多维搬运和转置减少数据重排的开销。指令集方面提供矩阵乘法、向量运算、数据搬运三类指令。矩阵乘法指令指定输入地址、形状、数据流模式。编译器负责分块和调度硬件提供性能计数器反馈。这个配置不算激进但胜在均衡。峰值算力大概在 32 TFLOPS 左右实际跑 Transformer 推理利用率能到 60% 以上有效算力约 20 TFLOPS。对于边缘服务器场景这个算力足够跑 BERT 级别的模型。7. 写在最后的一些个人体会做 AI 芯片软硬件协同设计这些年最大的体会是不要追求单项指标的极致要追求系统的均衡。峰值算力再高数据喂不上就是废铁。阵列再大利用率上不去就是浪费面积。精度再低模型跑不准就是白搭。另一个体会是软件团队要早介入。硬件设计阶段就要让编译器团队参与一起定义指令集和硬件接口。等硬件定型了再让软件适配往往发现接口设计不合理但已经改不动了。最后分享一个小技巧在硬件仿真阶段用真实的模型跑一遍把每一层的矩阵形状、数据量、访存模式都记录下来。这些数据是硬件参数选择的依据。我见过太多团队拍脑袋定阵列尺寸和缓存大小结果流片回来发现和实际负载不匹配。花一周时间做负载分析能省下半年的返工。这个方向后续还可以往几个方向扩展一是稀疏矩阵乘法的硬件支持二是动态形状的编译优化三是多芯片互联的调度策略。每一个都够写一篇长文有机会再聊。
返回列表