AI 芯片这个话题聊架构、聊算力、聊制程的人很多但真正落到软硬件怎么配合着设计这一层的往往就只剩下几句软硬协同的口号。我做了几年加速器相关的活儿越来越觉得AI 芯片的软硬件设计本质上是一道翻译题——把 Transformer 里那些看似优雅的数学操作翻译成硅片上真实流动的数据和电流。而这道题里最核心的词汇就是脉动阵列和注意力机制。这篇内容我想聊的是当你手里有一块要跑 Transformer 的 AI 芯片时硬件侧为什么偏偏是脉动阵列成了矩阵乘法的宠儿软件侧的多头注意力、因果掩码、位置编码又该怎么迁就硬件的脾气。适合正在做加速器架构、算子映射、或者单纯想把 Transformer 从 PyTorch 落到硬件上的朋友。不管你是刚接触脉动阵列基本原理还是已经在调 Swin Transformer 的窗口注意力我都尽量把为什么这么设计讲透而不是只丢结论。1. 为什么矩阵乘法决定了 AI 芯片的硬件骨架1.1 从 Transformer 的计算分布说起要理解 AI 芯片为什么长成现在这个样子得先看清楚 Transformer 到底把算力花在了哪里。一个标准的 Transformer 层抛开 LayerNorm 和残差这些配菜主菜就两块自注意力和前馈网络。前馈网络是两个大矩阵乘中间夹一个激活自注意力则是 Q、K、V 三个投影矩阵乘加上 QK^T 和 softmax 后的加权求和。我实测统计过一个 base 级别的模型前馈网络的参数量和计算量大概占到整个层的三分之二自注意力占三分之一左右。但注意这个比例会随序列长度剧烈变化——自注意力的计算量是 O(n²·d) 的序列一长QK^T 这个矩阵乘就会迅速膨胀成主导项。所以硬件设计者面对的第一个现实是芯片的主要任务就是高效地做矩阵乘而且是两种形态的矩阵乘——一种是规整的权重矩阵乘前馈、投影一种是随序列长度变化的注意力矩阵乘。这就解释了为什么几乎所有主流 AI 加速器的核心都是一个巨大的矩阵乘法引擎。GPU 用 Tensor CoreTPU 用脉动阵列各家 NPU 用 MAC 阵列名字不同本质都是在回答同一个问题怎么让乘加运算的吞吐最大化同时把数据搬运的能耗压到最低。1.2 数据搬运才是真正的能耗大头很多人第一次做加速器会犯一个直觉错误以为算力越强越好拼命堆 MAC 单元。但真正做过芯片的人都知道在先进工艺下一次乘加运算的能耗可能只有从 DRAM 搬一个数进来能耗的几百分之一。这个数量级差距不是夸张是工艺决定的物理现实。我打个比方计算单元就像厨房里的灶台数据就像食材。你可以把灶台做得又多又猛但如果食材要从很远的大仓库一趟趟搬灶台大部分时间都在空烧。AI 芯片设计的核心矛盾从来不是算得够不够快而是喂得够不够饱。这个认知直接推导出两个设计原则。第一数据复用要最大化——同一个数据读进来要尽可能参与多次计算摊薄搬运成本。第二计算要有规律性——规整的访问模式才能做高效的预取和流水随机访问是性能杀手。脉动阵列之所以成为矩阵乘的首选恰恰是因为它同时满足了这两条数据在阵列里像血液一样脉动流动每个数据都被复用了很多次而且流动路径完全规整。1.3 脉动阵列基本原理让数据流过计算单元脉动阵列Systolic Array这个名字很形象。Systolic 来自心脏收缩指的是数据像心跳泵血一样有节奏地、一波一波地流过整个计算阵列。我第一次看这个概念的时候觉得有点绕后来用一个特别土的类比就通了想象一个工厂流水线每个工位上坐着一个工人PE处理单元原料从左边和上边同时流进来每个工人做完自己那一步就把半成品往右下传。具体到矩阵乘 C A × B脉动阵列的做法是把 A 的元素按行从左侧注入把 B 的元素按列从上侧注入每个 PE 内部有一个乘加单元和一个小寄存器。数据每拍向右、向下各移动一格PE 在每个节拍里做一次乘加。这样设计的好处是权重可以预先加载并驻留在 PE 里不动激活值像波浪一样流过每个激活值在流过的过程中和多个权重相乘复用率极高。这里有个关键细节值得展开脉动阵列的脉动节奏是全局同步的。所有 PE 在同一个时钟节拍下动作数据的位置在每个节拍都是确定的。这种确定性带来一个巨大的工程好处——不需要复杂的地址生成和仲裁逻辑控制通路极其简单面积和功耗都省下来了。相比之下如果用一个大的 SRAM 加一堆 MAC 靠地址索引取数光是地址计算和 bank 冲突处理就能吃掉可观的功耗。不过脉动阵列也不是没有代价。它的规整性意味着灵活性差遇到非规整的计算比如稀疏、动态形状就不好使。这也是为什么后来的架构很多采用脉动阵列 向量单元的混合结构大矩阵乘交给脉动阵列element-wise 操作、softmax、LayerNorm 这些交给灵活的向量单元。2. 把注意力机制映射到脉动阵列上的真实难点2.1 注意力机制的计算拆解与硬件视角注意力机制的公式大家都背得出来Attention(Q,K,V) softmax(QK^T / √d) V。但从硬件角度看这个公式里藏着好几个不友好的地方。先说 QK^T。这是一个 [n, d] × [d, n] 的矩阵乘输出 [n, n]。当序列长度 n 是 512、1024 甚至更长时这个 n×n 的中间矩阵会非常大。硬件上如果老老实实把整个 QK^T 算出来存下来再算 softmax再乘 V那中间结果的存储和搬运开销会非常吓人。这就是为什么FlashAttention 这类算法在硬件圈这么受追捧——它把 softmax 和矩阵乘融合在一起分块计算中间结果不落大存储直接在片上 SRAM 里滚动处理。再说 softmax。softmax 需要先求每行的最大值为了数值稳定再算指数再归一化。这三步都是跨列依赖的——算第 j 个输出需要知道整行的信息。而脉动阵列擅长的是规整的、无跨行依赖的矩阵乘。所以 softmax 通常不能直接塞进脉动阵列得交给向量单元或者专门的 SFU特殊功能单元来做。这个矩阵乘引擎 向量引擎的分工几乎是现代 AI 芯片的标配。2.2 多头注意力给数据布局出的难题多头注意力Multi-Head Attention在数学上就是把 d 维的 Q、K、V 切成 h 份每份独立做注意力最后拼接。听起来简单但落到硬件上数据布局layout会变成一个让人头疼的问题。假设 d768h12每个头 64 维。在 PyTorch 里张量通常是 [batch, seq, d] 的布局。做多头的时候要 reshape 成 [batch, seq, h, d/h] 再 transpose 成 [batch, h, seq, d/h]。这个 transpose 在软件里就是一次内存重排但在硬件上如果布局不对会导致脉动阵列取数时访问不连续bank 冲突严重性能直接腰斩。我的经验是在硬件友好的实现里往往会在数据进入芯片之前就把布局调整好让每个头的数据在内存里是连续的。这样脉动阵列处理每个头的时候取数都是规整的连续块。代价是软件侧要做一次预处理但换来的是硬件侧稳定的高吞吐这笔账很划算。还有一个坑多头注意力的 h 个头之间是独立的理论上可以并行。但脉动阵列通常是一个大阵列怎么把 h 个头映射上去常见做法有两种。一种是时间上串行一个头一个头地算阵列利用率高但延迟大另一种是空间上切分把阵列分成 h 个小块并行算延迟低但每块的规模小、复用率下降。选哪种取决于你的芯片是追求吞吐还是追求延迟没有标准答案。2.3 因果掩码与位置编码的硬件处理自回归生成任务里注意力要加因果掩码causal mask让每个位置只能看到自己和之前的位置。这个掩码在数学上就是把 QK^T 矩阵的上三角部分置成负无穷。硬件上怎么处理最直接的做法是在 softmax 之前把对应位置的分数改掉。但这里有个效率问题如果老老实实算完整个 QK^T 再掩码那上三角那部分计算就白做了浪费了将近一半的算力。更聪明的做法是在分块计算时就跳过被完全掩码的块。比如 FlashAttention 的分块策略里对于完全在上三角的块直接跳过不算对于跨越对角线的块才做部分掩码。这样能把因果注意力的实际计算量降到接近一半。位置编码这块现在主流分两派绝对位置编码可学习或正弦和相对位置编码如 RoPE。绝对位置编码通常是在输入 embedding 上加一个位置向量这个加法是 element-wise 的交给向量单元就行没什么难度。RoPE 稍微麻烦点它是对 Q、K 做旋转操作涉及成对的乘加硬件上需要专门的旋转单元或者用向量单元模拟。我见过一些实现为了省事把 RoPE 放到软件侧预处理但这样会增加数据搬运得不偿失能放到片上做的尽量放片上。3. 软件栈怎么迁就硬件从算子到指令的翻译3.1 算子融合减少搬运的第一手段前面反复强调搬运是能耗大头那软件侧最有效的优化就是算子融合——把多个算子合并成一个中间结果不落 DRAM直接在片上流转。举个最常见的例子矩阵乘 偏置 激活GELU/ReLU。如果分开做矩阵乘的结果要写回 DRAM激活再读回来一来一回就是两次搬运。融合之后矩阵乘的输出直接进激活单元结果才写回。这一个融合就能省掉一次完整的读写。在 Transformer 里融合的机会特别多。LayerNorm 可以和前面的残差加法融合softmax 可以和 QK^T 融合就是 FlashAttention 的思路QKV 三个投影可以合并成一个大矩阵乘因为输入是同一个权重拼起来算一次更高效。我做过一个统计一个优化良好的 Transformer 推理实现算子数量能比朴素实现少一半以上省下来的全是搬运。但融合不是无脑合。融合的边界要考虑片上 SRAM 的容量——中间结果得放得下。如果为了融合硬塞导致 SRAM 溢出反而要 spill 到 DRAM那就适得其反了。所以融合策略要和芯片的 SRAM 大小、脉动阵列的 tile 尺寸一起考虑这是个联合优化问题。3.2 分块Tiling策略与脉动阵列尺寸的匹配脉动阵列的尺寸是固定的比如 128×128 或者 256×256。但实际的矩阵乘维度往往比这个大得多。所以必须做分块把大矩阵切成阵列能处理的小块一块一块地喂进去。分块策略直接决定了数据复用率。经典的矩阵乘分块有三种复用维度A 的行复用、B 的列复用、以及输出累加的复用。好的分块能让每个数据块被复用尽可能多次。这里有个经验公式如果阵列是 N×N那么分块时让内层维度接近 N复用率最高。因为脉动阵列在处理一个 N×N 块时输入数据流过的路径最长复用最充分。我踩过的一个坑是分块尺寸没和阵列尺寸对齐。比如阵列是 128×128但分块用了 100×100结果每次都有 28 行/列的空转利用率只有 61%。后来改成 128 的整数倍利用率直接上到 95% 以上。分块尺寸一定要是阵列尺寸的整数倍或者约数这个细节看起来小实测影响巨大。还有一个和注意力相关的分块技巧QK^T 的输出是 n×n如果 n 很大这个输出矩阵本身就要分块。分块之后softmax 的跨列依赖就变成了跨块依赖需要在线online计算 softmax 的归一化因子。这就是 FlashAttention 在线 softmax 的核心——维护一个 running max 和 running sum每处理一个新块就更新。这个算法在硬件上特别友好因为它把大矩阵的存储需求降到了 O(n) 而不是 O(n²)。3.3 指令调度与流水线填充硬件有了算子也映射好了最后一公里是指令调度。脉动阵列有个特点从数据注入到结果流出中间有固定的延迟流水线深度。如果调度不好阵列会出现气泡——前面一批数据算完了后面一批还没喂进来阵列空转。解决气泡的核心是双缓冲double buffering准备两块输入缓冲一块在被阵列消费的同时另一块在从内存加载下一批数据。这样加载和计算重叠阵列几乎不停。这个技术在硬件里是老生常谈但真正做好不容易因为要精确计算加载时间和计算时间的匹配。如果加载比计算慢那瓶颈就在加载得优化访存如果计算比加载慢那阵列就是瓶颈得考虑加宽数据通路。对于 Transformer 这种层数多、结构重复的网络还有一个优化是层间流水。第 L 层的输出直接作为第 L1 层的输入中间不落 DRAM。这要求软件栈能把整个网络或者至少几层作为一个整体来调度对编译器的要求比较高。我见过做得好的实现能把连续几层的中间激活全部保持在片上吞吐提升非常明显。4. 视觉 Transformer 带来的新硬件挑战4.1 从 Swin Transformer 看窗口注意力的局部性前面聊的都是标准 Transformer注意力是全局的。但视觉领域流行的 Swin Transformer 引入了窗口注意力——把图像切成不重叠的窗口注意力只在窗口内做。这个设计对硬件其实是个好消息因为窗口内的序列长度短QK^T 的规模小片上就能放下不需要复杂的分块和在线 softmax。但 Swin 有个移位窗口shifted window的机制窗口会偏移半个窗口大小导致窗口边界被切断需要重新划分。这个操作在硬件上是个麻烦因为它涉及跨窗口的数据交换访问模式不规整。我的经验是移位窗口这一步最好在软件侧或者用专门的 shuffle 单元处理不要指望脉动阵列去做它不擅长这个。Swin 的分层结构hierarchical也值得注意。它像 CNN 一样逐级下采样特征图越来越小、通道越来越多。这意味着不同层的最优硬件配置可能不同——浅层特征图大但通道少深层特征图小但通道多。理想的硬件应该能动态调整但固定阵列做不到只能取折中。这也是为什么很多视觉芯片会同时配备卷积引擎和注意力引擎让不同层用不同的单元。4.2 通道注意力与空间注意力的硬件差异视觉领域还有一大类注意力是通道注意力如 SE 模块和空间注意力。这两类和 Transformer 的注意力机制在硬件上差别很大值得单独说。通道注意力的做法是对每个通道做全局池化得到一个通道描述向量过一个小 MLP 算出每个通道的权重再乘回原特征图。这里的全局池化是跨空间维度的归约硬件上就是一个求和树很规整。MLP 很小向量单元就能搞定。整体来说通道注意力对硬件很友好开销不大。空间注意力则是在空间维度上算注意力比如对特征图的每个位置算一个权重。这个和 Transformer 注意力更像但维度通常小一些。视觉领域还专门有高度方向、宽度方向注意力这种设计把二维注意力拆成两个一维的分别沿 H 和 W 方向做。这种拆解的好处是计算量从 O(H²W²) 降到 O(H²W HW²)硬件上更容易处理而且两个方向的注意力可以流水并行。我个人的观察是视觉注意力的硬件实现比 NLP 注意力更碎片化因为视觉任务对延迟敏感、对精度容忍度高所以各种轻量化、局部化的注意力变体特别多。做视觉芯片的软件栈得有足够的灵活性去支持这些五花八门的算子不能只盯着标准多头注意力。4.3 轻量 Transformer 的硬件取舍轻量 Transformer如 MobileViT、EfficientFormer 这类为了在边缘设备上跑做了大量精简减少头数、降低维度、用深度可分离卷积替代部分注意力、甚至用线性注意力近似。这些精简对硬件设计影响很大。线性注意力是个典型例子。它把 softmax(QK^T)V 重新结合成 Q(K^T V)利用矩阵乘结合律把复杂度从 O(n²) 降到 O(n)。数学上很漂亮但硬件上要注意K^T V 是一个 [d, n] × [n, d] 的乘输出 [d, d]这个 d×d 的中间矩阵和序列长度无关可以常驻片上。对于长序列场景这个改动能极大降低存储压力。但代价是精度会掉一些而且 Q(K^T V) 的计算模式和标准注意力不同脉动阵列的调度要重新设计。边缘芯片还有个硬约束是功耗和面积。脉动阵列虽然高效但面积不小。轻量模型往往用更小的阵列甚至用 SIMD 向量单元代替脉动阵列。这时候软件栈的重点就从榨干阵列吞吐变成减少访存次数和提高并行度。我做过一个边缘端的对比同样一个轻量 Transformer用 64×64 小阵列比用 256×256 大阵列虽然峰值算力低很多但因为模型本身小、访存少实际能效反而更高。硬件规模要匹配模型规模大马拉小车在边缘场景是浪费。5. 实操中那些文档不会告诉你的经验5.1 数值精度从 FP32 到 INT8 的踩坑记录Transformer 对数值精度比 CNN 敏感这是我踩过最深的坑之一。CNN 量化到 INT8 通常掉点很少但 Transformer 直接 INT8 量化注意力分数容易溢出softmax 之后分布畸变精度掉得厉害。我的经验是分而治之。矩阵乘部分可以大胆用 INT8因为它是线性累加动态范围可控。但softmax 的输入也就是 QK^T 的结果要格外小心通常需要保留更高的精度或者做精细的缩放。LayerNorm 也是敏感点它的均值和方差计算涉及平方和容易溢出一般保留 FP16 或 BF16。还有一个细节是 QK^T 里的那个 1/√d 缩放。如果放到 softmax 之前做能有效控制分数范围减少溢出风险。有些实现为了省事把这个缩放放到 softmax 之后结果就是 softmax 输入范围过大量化时精度损失严重。这个缩放的位置对量化友好度影响很大建议放在 softmax 之前。5.2 调试工具链怎么定位性能瓶颈做软硬件协同最怕的是性能不达预期却不知道卡在哪。我总结了一套排查顺序基本能覆盖大部分情况。第一步看阵列利用率。如果利用率低说明是喂数不够瓶颈在访存或者调度。第二步看访存带宽。如果带宽打满了那就是数据搬运的问题得考虑提高复用率或者融合算子。第三步看流水线气泡。如果阵列有周期性空转多半是双缓冲没做好或者分块尺寸不匹配。第四步看指令开销。如果计算单元利用率高但整体吞吐还是低可能是控制指令太多得考虑用更大的指令粒度或者硬件循环。这套顺序的逻辑是从大到小、从粗到细先确认是不是根本性的资源瓶颈再逐步深入到调度细节。我见过不少人一上来就抠指令级优化结果发现根本问题是分块尺寸没对齐白忙活。5.3 从模型到芯片的协同设计建议最后分享几条我在实际项目里总结的协同设计原则都是血泪换来的。第一软件侧要尽早介入硬件设计。不要等硬件流片了才想软件怎么适配那时候很多决策已经改不了了。理想情况是算法、软件、硬件三方从第一天就坐在一起算法知道硬件的约束硬件知道算法的趋势。第二为未来的模型留余量。Transformer 的变体层出不穷今天支持多头注意力明天可能就要支持各种稀疏、线性、局部注意力。硬件如果能留一些可编程的灵活性比如可配置的阵列分块、可编程的向量单元生命周期会长很多。第三把最频繁的操作做成硬件把最多变的操作留给软件。矩阵乘最频繁做成脉动阵列softmax、激活、各种注意力变体最多变留给向量单元和可编程逻辑。这个分工原则能兼顾效率和灵活性。第四一定要做端到端的 profiling。不要只看单个算子的性能要看整个模型跑下来的实际表现。我见过单算子都很快但整体很慢的情况问题出在算子之间的数据搬运和同步上。端到端看才能发现真正的瓶颈。说到底AI 芯片的软硬件设计没有银弹。脉动阵列是个好工具但它只解决矩阵乘这一件事注意力机制的各种变体需要软件栈用融合、分块、调度这些手段去适配。真正做得好的团队都是把算法、软件、硬件当成一个整体来打磨的。我个人的体会是理解硬件的脾气比记住多少种注意力变体更重要——因为变体会过时但数据搬运是瓶颈规整性带来效率这些底层规律会一直成立。 
长河君 · 2026/10/8 20:41:28 · 阅读 1.2k AI芯片软硬件协同设计：脉动阵列与Transformer注意力机制映射实战
AI 芯片这个话题聊架构、聊算力、聊制程的人很多但真正落到软硬件怎么配合着设计这一层的往往就只剩下几句软硬协同的口号。我做了几年加速器相关的活儿越来越觉得AI 芯片的软硬件设计本质上是一道翻译题——把 Transformer 里那些看似优雅的数学操作翻译成硅片上真实流动的数据和电流。而这道题里最核心的词汇就是脉动阵列和注意力机制。这篇内容我想聊的是当你手里有一块要跑 Transformer 的 AI 芯片时硬件侧为什么偏偏是脉动阵列成了矩阵乘法的宠儿软件侧的多头注意力、因果掩码、位置编码又该怎么迁就硬件的脾气。适合正在做加速器架构、算子映射、或者单纯想把 Transformer 从 PyTorch 落到硬件上的朋友。不管你是刚接触脉动阵列基本原理还是已经在调 Swin Transformer 的窗口注意力我都尽量把为什么这么设计讲透而不是只丢结论。1. 为什么矩阵乘法决定了 AI 芯片的硬件骨架1.1 从 Transformer 的计算分布说起要理解 AI 芯片为什么长成现在这个样子得先看清楚 Transformer 到底把算力花在了哪里。一个标准的 Transformer 层抛开 LayerNorm 和残差这些配菜主菜就两块自注意力和前馈网络。前馈网络是两个大矩阵乘中间夹一个激活自注意力则是 Q、K、V 三个投影矩阵乘加上 QK^T 和 softmax 后的加权求和。我实测统计过一个 base 级别的模型前馈网络的参数量和计算量大概占到整个层的三分之二自注意力占三分之一左右。但注意这个比例会随序列长度剧烈变化——自注意力的计算量是 O(n²·d) 的序列一长QK^T 这个矩阵乘就会迅速膨胀成主导项。所以硬件设计者面对的第一个现实是芯片的主要任务就是高效地做矩阵乘而且是两种形态的矩阵乘——一种是规整的权重矩阵乘前馈、投影一种是随序列长度变化的注意力矩阵乘。这就解释了为什么几乎所有主流 AI 加速器的核心都是一个巨大的矩阵乘法引擎。GPU 用 Tensor CoreTPU 用脉动阵列各家 NPU 用 MAC 阵列名字不同本质都是在回答同一个问题怎么让乘加运算的吞吐最大化同时把数据搬运的能耗压到最低。1.2 数据搬运才是真正的能耗大头很多人第一次做加速器会犯一个直觉错误以为算力越强越好拼命堆 MAC 单元。但真正做过芯片的人都知道在先进工艺下一次乘加运算的能耗可能只有从 DRAM 搬一个数进来能耗的几百分之一。这个数量级差距不是夸张是工艺决定的物理现实。我打个比方计算单元就像厨房里的灶台数据就像食材。你可以把灶台做得又多又猛但如果食材要从很远的大仓库一趟趟搬灶台大部分时间都在空烧。AI 芯片设计的核心矛盾从来不是算得够不够快而是喂得够不够饱。这个认知直接推导出两个设计原则。第一数据复用要最大化——同一个数据读进来要尽可能参与多次计算摊薄搬运成本。第二计算要有规律性——规整的访问模式才能做高效的预取和流水随机访问是性能杀手。脉动阵列之所以成为矩阵乘的首选恰恰是因为它同时满足了这两条数据在阵列里像血液一样脉动流动每个数据都被复用了很多次而且流动路径完全规整。1.3 脉动阵列基本原理让数据流过计算单元脉动阵列Systolic Array这个名字很形象。Systolic 来自心脏收缩指的是数据像心跳泵血一样有节奏地、一波一波地流过整个计算阵列。我第一次看这个概念的时候觉得有点绕后来用一个特别土的类比就通了想象一个工厂流水线每个工位上坐着一个工人PE处理单元原料从左边和上边同时流进来每个工人做完自己那一步就把半成品往右下传。具体到矩阵乘 C A × B脉动阵列的做法是把 A 的元素按行从左侧注入把 B 的元素按列从上侧注入每个 PE 内部有一个乘加单元和一个小寄存器。数据每拍向右、向下各移动一格PE 在每个节拍里做一次乘加。这样设计的好处是权重可以预先加载并驻留在 PE 里不动激活值像波浪一样流过每个激活值在流过的过程中和多个权重相乘复用率极高。这里有个关键细节值得展开脉动阵列的脉动节奏是全局同步的。所有 PE 在同一个时钟节拍下动作数据的位置在每个节拍都是确定的。这种确定性带来一个巨大的工程好处——不需要复杂的地址生成和仲裁逻辑控制通路极其简单面积和功耗都省下来了。相比之下如果用一个大的 SRAM 加一堆 MAC 靠地址索引取数光是地址计算和 bank 冲突处理就能吃掉可观的功耗。不过脉动阵列也不是没有代价。它的规整性意味着灵活性差遇到非规整的计算比如稀疏、动态形状就不好使。这也是为什么后来的架构很多采用脉动阵列 向量单元的混合结构大矩阵乘交给脉动阵列element-wise 操作、softmax、LayerNorm 这些交给灵活的向量单元。2. 把注意力机制映射到脉动阵列上的真实难点2.1 注意力机制的计算拆解与硬件视角注意力机制的公式大家都背得出来Attention(Q,K,V) softmax(QK^T / √d) V。但从硬件角度看这个公式里藏着好几个不友好的地方。先说 QK^T。这是一个 [n, d] × [d, n] 的矩阵乘输出 [n, n]。当序列长度 n 是 512、1024 甚至更长时这个 n×n 的中间矩阵会非常大。硬件上如果老老实实把整个 QK^T 算出来存下来再算 softmax再乘 V那中间结果的存储和搬运开销会非常吓人。这就是为什么FlashAttention 这类算法在硬件圈这么受追捧——它把 softmax 和矩阵乘融合在一起分块计算中间结果不落大存储直接在片上 SRAM 里滚动处理。再说 softmax。softmax 需要先求每行的最大值为了数值稳定再算指数再归一化。这三步都是跨列依赖的——算第 j 个输出需要知道整行的信息。而脉动阵列擅长的是规整的、无跨行依赖的矩阵乘。所以 softmax 通常不能直接塞进脉动阵列得交给向量单元或者专门的 SFU特殊功能单元来做。这个矩阵乘引擎 向量引擎的分工几乎是现代 AI 芯片的标配。2.2 多头注意力给数据布局出的难题多头注意力Multi-Head Attention在数学上就是把 d 维的 Q、K、V 切成 h 份每份独立做注意力最后拼接。听起来简单但落到硬件上数据布局layout会变成一个让人头疼的问题。假设 d768h12每个头 64 维。在 PyTorch 里张量通常是 [batch, seq, d] 的布局。做多头的时候要 reshape 成 [batch, seq, h, d/h] 再 transpose 成 [batch, h, seq, d/h]。这个 transpose 在软件里就是一次内存重排但在硬件上如果布局不对会导致脉动阵列取数时访问不连续bank 冲突严重性能直接腰斩。我的经验是在硬件友好的实现里往往会在数据进入芯片之前就把布局调整好让每个头的数据在内存里是连续的。这样脉动阵列处理每个头的时候取数都是规整的连续块。代价是软件侧要做一次预处理但换来的是硬件侧稳定的高吞吐这笔账很划算。还有一个坑多头注意力的 h 个头之间是独立的理论上可以并行。但脉动阵列通常是一个大阵列怎么把 h 个头映射上去常见做法有两种。一种是时间上串行一个头一个头地算阵列利用率高但延迟大另一种是空间上切分把阵列分成 h 个小块并行算延迟低但每块的规模小、复用率下降。选哪种取决于你的芯片是追求吞吐还是追求延迟没有标准答案。2.3 因果掩码与位置编码的硬件处理自回归生成任务里注意力要加因果掩码causal mask让每个位置只能看到自己和之前的位置。这个掩码在数学上就是把 QK^T 矩阵的上三角部分置成负无穷。硬件上怎么处理最直接的做法是在 softmax 之前把对应位置的分数改掉。但这里有个效率问题如果老老实实算完整个 QK^T 再掩码那上三角那部分计算就白做了浪费了将近一半的算力。更聪明的做法是在分块计算时就跳过被完全掩码的块。比如 FlashAttention 的分块策略里对于完全在上三角的块直接跳过不算对于跨越对角线的块才做部分掩码。这样能把因果注意力的实际计算量降到接近一半。位置编码这块现在主流分两派绝对位置编码可学习或正弦和相对位置编码如 RoPE。绝对位置编码通常是在输入 embedding 上加一个位置向量这个加法是 element-wise 的交给向量单元就行没什么难度。RoPE 稍微麻烦点它是对 Q、K 做旋转操作涉及成对的乘加硬件上需要专门的旋转单元或者用向量单元模拟。我见过一些实现为了省事把 RoPE 放到软件侧预处理但这样会增加数据搬运得不偿失能放到片上做的尽量放片上。3. 软件栈怎么迁就硬件从算子到指令的翻译3.1 算子融合减少搬运的第一手段前面反复强调搬运是能耗大头那软件侧最有效的优化就是算子融合——把多个算子合并成一个中间结果不落 DRAM直接在片上流转。举个最常见的例子矩阵乘 偏置 激活GELU/ReLU。如果分开做矩阵乘的结果要写回 DRAM激活再读回来一来一回就是两次搬运。融合之后矩阵乘的输出直接进激活单元结果才写回。这一个融合就能省掉一次完整的读写。在 Transformer 里融合的机会特别多。LayerNorm 可以和前面的残差加法融合softmax 可以和 QK^T 融合就是 FlashAttention 的思路QKV 三个投影可以合并成一个大矩阵乘因为输入是同一个权重拼起来算一次更高效。我做过一个统计一个优化良好的 Transformer 推理实现算子数量能比朴素实现少一半以上省下来的全是搬运。但融合不是无脑合。融合的边界要考虑片上 SRAM 的容量——中间结果得放得下。如果为了融合硬塞导致 SRAM 溢出反而要 spill 到 DRAM那就适得其反了。所以融合策略要和芯片的 SRAM 大小、脉动阵列的 tile 尺寸一起考虑这是个联合优化问题。3.2 分块Tiling策略与脉动阵列尺寸的匹配脉动阵列的尺寸是固定的比如 128×128 或者 256×256。但实际的矩阵乘维度往往比这个大得多。所以必须做分块把大矩阵切成阵列能处理的小块一块一块地喂进去。分块策略直接决定了数据复用率。经典的矩阵乘分块有三种复用维度A 的行复用、B 的列复用、以及输出累加的复用。好的分块能让每个数据块被复用尽可能多次。这里有个经验公式如果阵列是 N×N那么分块时让内层维度接近 N复用率最高。因为脉动阵列在处理一个 N×N 块时输入数据流过的路径最长复用最充分。我踩过的一个坑是分块尺寸没和阵列尺寸对齐。比如阵列是 128×128但分块用了 100×100结果每次都有 28 行/列的空转利用率只有 61%。后来改成 128 的整数倍利用率直接上到 95% 以上。分块尺寸一定要是阵列尺寸的整数倍或者约数这个细节看起来小实测影响巨大。还有一个和注意力相关的分块技巧QK^T 的输出是 n×n如果 n 很大这个输出矩阵本身就要分块。分块之后softmax 的跨列依赖就变成了跨块依赖需要在线online计算 softmax 的归一化因子。这就是 FlashAttention 在线 softmax 的核心——维护一个 running max 和 running sum每处理一个新块就更新。这个算法在硬件上特别友好因为它把大矩阵的存储需求降到了 O(n) 而不是 O(n²)。3.3 指令调度与流水线填充硬件有了算子也映射好了最后一公里是指令调度。脉动阵列有个特点从数据注入到结果流出中间有固定的延迟流水线深度。如果调度不好阵列会出现气泡——前面一批数据算完了后面一批还没喂进来阵列空转。解决气泡的核心是双缓冲double buffering准备两块输入缓冲一块在被阵列消费的同时另一块在从内存加载下一批数据。这样加载和计算重叠阵列几乎不停。这个技术在硬件里是老生常谈但真正做好不容易因为要精确计算加载时间和计算时间的匹配。如果加载比计算慢那瓶颈就在加载得优化访存如果计算比加载慢那阵列就是瓶颈得考虑加宽数据通路。对于 Transformer 这种层数多、结构重复的网络还有一个优化是层间流水。第 L 层的输出直接作为第 L1 层的输入中间不落 DRAM。这要求软件栈能把整个网络或者至少几层作为一个整体来调度对编译器的要求比较高。我见过做得好的实现能把连续几层的中间激活全部保持在片上吞吐提升非常明显。4. 视觉 Transformer 带来的新硬件挑战4.1 从 Swin Transformer 看窗口注意力的局部性前面聊的都是标准 Transformer注意力是全局的。但视觉领域流行的 Swin Transformer 引入了窗口注意力——把图像切成不重叠的窗口注意力只在窗口内做。这个设计对硬件其实是个好消息因为窗口内的序列长度短QK^T 的规模小片上就能放下不需要复杂的分块和在线 softmax。但 Swin 有个移位窗口shifted window的机制窗口会偏移半个窗口大小导致窗口边界被切断需要重新划分。这个操作在硬件上是个麻烦因为它涉及跨窗口的数据交换访问模式不规整。我的经验是移位窗口这一步最好在软件侧或者用专门的 shuffle 单元处理不要指望脉动阵列去做它不擅长这个。Swin 的分层结构hierarchical也值得注意。它像 CNN 一样逐级下采样特征图越来越小、通道越来越多。这意味着不同层的最优硬件配置可能不同——浅层特征图大但通道少深层特征图小但通道多。理想的硬件应该能动态调整但固定阵列做不到只能取折中。这也是为什么很多视觉芯片会同时配备卷积引擎和注意力引擎让不同层用不同的单元。4.2 通道注意力与空间注意力的硬件差异视觉领域还有一大类注意力是通道注意力如 SE 模块和空间注意力。这两类和 Transformer 的注意力机制在硬件上差别很大值得单独说。通道注意力的做法是对每个通道做全局池化得到一个通道描述向量过一个小 MLP 算出每个通道的权重再乘回原特征图。这里的全局池化是跨空间维度的归约硬件上就是一个求和树很规整。MLP 很小向量单元就能搞定。整体来说通道注意力对硬件很友好开销不大。空间注意力则是在空间维度上算注意力比如对特征图的每个位置算一个权重。这个和 Transformer 注意力更像但维度通常小一些。视觉领域还专门有高度方向、宽度方向注意力这种设计把二维注意力拆成两个一维的分别沿 H 和 W 方向做。这种拆解的好处是计算量从 O(H²W²) 降到 O(H²W HW²)硬件上更容易处理而且两个方向的注意力可以流水并行。我个人的观察是视觉注意力的硬件实现比 NLP 注意力更碎片化因为视觉任务对延迟敏感、对精度容忍度高所以各种轻量化、局部化的注意力变体特别多。做视觉芯片的软件栈得有足够的灵活性去支持这些五花八门的算子不能只盯着标准多头注意力。4.3 轻量 Transformer 的硬件取舍轻量 Transformer如 MobileViT、EfficientFormer 这类为了在边缘设备上跑做了大量精简减少头数、降低维度、用深度可分离卷积替代部分注意力、甚至用线性注意力近似。这些精简对硬件设计影响很大。线性注意力是个典型例子。它把 softmax(QK^T)V 重新结合成 Q(K^T V)利用矩阵乘结合律把复杂度从 O(n²) 降到 O(n)。数学上很漂亮但硬件上要注意K^T V 是一个 [d, n] × [n, d] 的乘输出 [d, d]这个 d×d 的中间矩阵和序列长度无关可以常驻片上。对于长序列场景这个改动能极大降低存储压力。但代价是精度会掉一些而且 Q(K^T V) 的计算模式和标准注意力不同脉动阵列的调度要重新设计。边缘芯片还有个硬约束是功耗和面积。脉动阵列虽然高效但面积不小。轻量模型往往用更小的阵列甚至用 SIMD 向量单元代替脉动阵列。这时候软件栈的重点就从榨干阵列吞吐变成减少访存次数和提高并行度。我做过一个边缘端的对比同样一个轻量 Transformer用 64×64 小阵列比用 256×256 大阵列虽然峰值算力低很多但因为模型本身小、访存少实际能效反而更高。硬件规模要匹配模型规模大马拉小车在边缘场景是浪费。5. 实操中那些文档不会告诉你的经验5.1 数值精度从 FP32 到 INT8 的踩坑记录Transformer 对数值精度比 CNN 敏感这是我踩过最深的坑之一。CNN 量化到 INT8 通常掉点很少但 Transformer 直接 INT8 量化注意力分数容易溢出softmax 之后分布畸变精度掉得厉害。我的经验是分而治之。矩阵乘部分可以大胆用 INT8因为它是线性累加动态范围可控。但softmax 的输入也就是 QK^T 的结果要格外小心通常需要保留更高的精度或者做精细的缩放。LayerNorm 也是敏感点它的均值和方差计算涉及平方和容易溢出一般保留 FP16 或 BF16。还有一个细节是 QK^T 里的那个 1/√d 缩放。如果放到 softmax 之前做能有效控制分数范围减少溢出风险。有些实现为了省事把这个缩放放到 softmax 之后结果就是 softmax 输入范围过大量化时精度损失严重。这个缩放的位置对量化友好度影响很大建议放在 softmax 之前。5.2 调试工具链怎么定位性能瓶颈做软硬件协同最怕的是性能不达预期却不知道卡在哪。我总结了一套排查顺序基本能覆盖大部分情况。第一步看阵列利用率。如果利用率低说明是喂数不够瓶颈在访存或者调度。第二步看访存带宽。如果带宽打满了那就是数据搬运的问题得考虑提高复用率或者融合算子。第三步看流水线气泡。如果阵列有周期性空转多半是双缓冲没做好或者分块尺寸不匹配。第四步看指令开销。如果计算单元利用率高但整体吞吐还是低可能是控制指令太多得考虑用更大的指令粒度或者硬件循环。这套顺序的逻辑是从大到小、从粗到细先确认是不是根本性的资源瓶颈再逐步深入到调度细节。我见过不少人一上来就抠指令级优化结果发现根本问题是分块尺寸没对齐白忙活。5.3 从模型到芯片的协同设计建议最后分享几条我在实际项目里总结的协同设计原则都是血泪换来的。第一软件侧要尽早介入硬件设计。不要等硬件流片了才想软件怎么适配那时候很多决策已经改不了了。理想情况是算法、软件、硬件三方从第一天就坐在一起算法知道硬件的约束硬件知道算法的趋势。第二为未来的模型留余量。Transformer 的变体层出不穷今天支持多头注意力明天可能就要支持各种稀疏、线性、局部注意力。硬件如果能留一些可编程的灵活性比如可配置的阵列分块、可编程的向量单元生命周期会长很多。第三把最频繁的操作做成硬件把最多变的操作留给软件。矩阵乘最频繁做成脉动阵列softmax、激活、各种注意力变体最多变留给向量单元和可编程逻辑。这个分工原则能兼顾效率和灵活性。第四一定要做端到端的 profiling。不要只看单个算子的性能要看整个模型跑下来的实际表现。我见过单算子都很快但整体很慢的情况问题出在算子之间的数据搬运和同步上。端到端看才能发现真正的瓶颈。说到底AI 芯片的软硬件设计没有银弹。脉动阵列是个好工具但它只解决矩阵乘这一件事注意力机制的各种变体需要软件栈用融合、分块、调度这些手段去适配。真正做得好的团队都是把算法、软件、硬件当成一个整体来打磨的。我个人的体会是理解硬件的脾气比记住多少种注意力变体更重要——因为变体会过时但数据搬运是瓶颈规整性带来效率这些底层规律会一直成立。