1. 从 TPUv1 说起AI 芯片软硬件协同设计的核心命题2016 年 Google 公开 TPUv1 的时候整个工业界第一次看到了一颗完全为神经网络推理而生的 ASIC 长什么样。它没有缓存层级、没有分支预测、没有乱序执行有的只是一块 256×256 的脉动阵列Systolic Array、一块 24MB 的片上统一缓冲Unified Buffer以及一套极其克制的 CISC 指令集。很多人第一次看论文会觉得这东西怎么这么简单但真正做过芯片架构的人会明白TPUv1 的简单是刻意为之——它把通用处理器里所有为灵活性付出的面积和功耗代价全部砍掉换成了纯粹的矩阵乘加吞吐。我接触 AI 芯片设计这些年最大的体会是AI 芯片的软硬件设计从来不是先做硬件再写编译器的串行流程而是一个从算法特征反推数据流、从数据流反推微架构、从微架构反推指令集和编译策略的闭环。TPUv1 之所以能在 2015 年就做到 92 TOPS 的 INT8 峰值算力、功耗只有 40W核心原因就是它把矩阵乘法占绝对主导这个算法事实吃透了然后围绕这个事实把整个软硬件栈重新设计了一遍。这篇文章我想把 AI 芯片软硬件设计里最核心的几条线索串起来讲脉动阵列到底怎么工作、为什么 Transformer 这类模型对硬件提出了新要求、指令集和编译器如何配合数据流、以及在实际项目中怎么权衡各种参数。内容会偏工程视角尽量把为什么这么设计讲清楚而不是只罗列结论。2. 脉动阵列AI 芯片数据流设计的基石2.1 脉动阵列到底在脉动什么脉动阵列这个概念其实很老1978 年 H.T. Kung 就提出了最初用于信号处理。它的核心思想是让数据像心跳一样有节奏地流过一个个处理单元PE每个 PE 只做最简单的乘加运算数据在流动过程中被反复复用。我用一个最直观的类比来解释。假设你要算一个 3×3 矩阵乘以一个 3×1 向量的结果传统做法是每个输出元素单独算需要把矩阵的每一行都从内存里读一遍。而脉动阵列的做法是把矩阵的元素预先钉在每个 PE 上不动让向量元素从左边一个个流进来每经过一个 PE 就做一次乘加结果从下方流出。这样向量元素只读一次就被复用了 3 次。TPUv1 的 256×256 脉动阵列就是这个思路的极致放大。权重Weight从上方加载并驻留在 PE 阵列中激活值Activation从左侧流入部分和Partial Sum从上方往下累积。每个时钟周期256 个激活值同时进入阵列经过 256 个周期后一整块 256×256 的矩阵乘法结果就从下方流出。整个过程没有一次访存是浪费的。2.2 为什么是脉动阵列而不是别的做 AI 加速器数据流设计主流方案其实有好几种权重固定Weight Stationary、输出固定Output Stationary、行固定Row Stationary以及脉动阵列这种可以看作权重固定输出流动的混合方案。选哪种取决于你的目标模型的计算特征。我列一个简单的对比表这是我在实际选型时常用的判断框架数据流方案核心思想优势劣势典型适用场景权重固定权重驻留 PE激活流动权重复用率高适合大 batch 推理激活复用低全连接层为主的老式网络输出固定部分和驻留 PE减少部分和搬运功耗权重和激活复用都一般通用卷积加速行固定行列混合复用卷积复用率最优控制逻辑复杂卷积神经网络脉动阵列权重驻留数据节拍流动复用率极高控制简单灵活性差尺寸固定矩阵乘法密集的模型TPUv1 选脉动阵列本质是因为它服务的模型2015 年前后的 MLP、LSTM、早期 CNN里矩阵乘法占了 90% 以上的计算量。脉动阵列把权重复用做到了极致——一个权重加载进 PE 后会被 256 个不同的激活值复用这意味着权重访存功耗被摊薄了 256 倍。在芯片功耗里访存功耗往往是计算功耗的几倍甚至十几倍这个复用带来的收益是决定性的。2.3 脉动阵列的尺寸怎么定这是实际项目里最常被问到的问题为什么 TPUv1 是 256×256而不是 128×128 或 512×512尺寸选择本质是一个面积、功耗、利用率三者的权衡。阵列越大单次能处理的矩阵块越大权重复用率越高但面积和功耗也线性增长而且当实际矩阵维度小于阵列尺寸时利用率会急剧下降。我做过一个粗略的估算假设 PE 是一个 8 位乘法器加一个 32 位累加器在 28nm 工艺下单个 PE 面积大约 0.001 平方毫米。256×256 就是 65536 个 PE光 PE 阵列面积就约 65 平方毫米加上缓冲和互连整颗芯片面积会到 300 平方毫米以上。如果换成 512×512PE 数量翻四倍面积直接爆掉良率也会很难看。另一个关键约束是矩阵维度的匹配度。主流模型的隐藏层维度通常是 256、512、768、1024 这类 2 的幂次256×256 的阵列能整除大部分维度利用率比较稳。如果选 300×300 这种非 2 幂次编译器处理边界时会非常痛苦。实操心得阵列尺寸不要盲目追大。我见过一些团队为了刷峰值算力指标把阵列做到 512×512结果实际模型跑下来利用率只有 30%因为大部分层的维度根本填不满阵列反而浪费了面积。阵列尺寸应该匹配你的目标模型的主力维度而不是越大越好。3. Transformer 给 AI 芯片带来的新挑战3.1 从矩阵乘法到注意力机制如果说 TPUv1 时代的模型是矩阵乘法为主那 Transformer 时代的模型就是矩阵乘法注意力大量逐元素操作的混合体。这个变化对硬件设计的影响是深远的。Transformer 的核心是自注意力Self-Attention计算过程可以拆成几步输入经过三个线性投影得到 Q、K、V 三个矩阵然后计算 QK^T 得到注意力分数经过 Softmax 归一化再和 V 相乘得到输出。这里面 QK^T 和注意力分数乘 V 都是矩阵乘法适合脉动阵列但 Softmax 涉及指数运算、求最大值、求和、除法这些是逐元素操作脉动阵列处理起来效率很低。更麻烦的是Transformer 的计算模式是动态的。序列长度决定了矩阵的维度而序列长度在不同任务、不同输入下是变化的。一个为固定 256×256 矩阵优化的脉动阵列遇到序列长度 512 或 1024 时要么分块处理导致利用率下降要么需要重新配置数据流。3.2 注意力机制的硬件友好化改造实际做芯片时我们不会傻乎乎地按原始公式实现注意力而是会做一系列硬件友好化改造。这些改造在软件层面看是数学等价变换在硬件层面看是把不规则计算变成规则计算。第一个改造是分块注意力Tiled Attention。把 Q、K、V 按序列维度切成小块每次只计算一个小块的注意力块内的中间结果放在片上缓冲里避免频繁访问外部内存。这个思路和 FlashAttention 的工程实现是一致的只不过 FlashAttention 是在 GPU 上做我们在芯片上做的是更底层的版本。第二个改造是在线 Softmax。标准 Softmax 需要先求全局最大值再求指数和这要求把整行注意力分数都存下来。在线 Softmax 的做法是边计算边更新最大值和指数和用一次遍历完成大幅减少片上存储需求。这个技巧在芯片设计里几乎是标配因为片上 SRAM 太贵了。第三个改造是算子融合。把 QK^T、Softmax、乘 V 这三个操作融合成一个流水线中间结果不落外部内存。这样做的收益非常大——原本三次访存变成一次功耗能降一半以上。3.3 逐元素操作的硬件加速Softmax 里的指数运算、LayerNorm 里的开方和除法这些在通用 CPU 上就是几条指令的事但在 AI 芯片上需要专门的硬件单元。常见的做法是设计一个向量处理单元VPU专门处理这类逐元素操作和脉动阵列并行工作。VPU 的设计要点是吞吐和精度的平衡。指数运算如果用查找表实现精度和表大小直接相关如果用多项式逼近需要多个乘法器和加法器。实际项目中我们通常会用分段线性逼近加查找表修正的方案在精度损失可控的前提下把面积压下来。注意事项Transformer 模型对数值精度比传统 CNN 敏感得多。我踩过的坑是把 Softmax 的指数运算精度砍得太狠导致注意力分布失真模型准确率直接掉了好几个点。逐元素操作的精度预算要单独评估不能和矩阵乘法用同一套标准。4. 指令集与编译器的协同设计4.1 为什么 AI 芯片的指令集要窄而深TPUv1 的指令集只有十来条指令包括矩阵乘法、激活、归一化、数据搬运等。这种窄而深的设计和通用 CPU 的宽而浅形成鲜明对比。原因很简单AI 芯片的指令集不需要表达所有计算只需要表达神经网络里的那几类计算然后把每类计算做到极致。以矩阵乘法指令为例TPUv1 的一条矩阵乘法指令可以指定输入地址、权重地址、输出地址和累加标志硬件会自动完成整个 256×256 阵列的运算。如果用通用 CPU 的思路这需要成千上万条指令才能完成。指令越少取指和译码的功耗就越低这对功耗敏感的推理芯片至关重要。但窄而深也有代价灵活性差。如果模型里出现了指令集不支持的算子就只能回退到 CPU 或者拆成多条指令模拟性能会断崖式下跌。所以指令集设计的一个核心原则是覆盖目标模型的主力算子边缘算子用组合指令或软件模拟兜底。4.2 编译器如何把模型映射到硬件编译器是 AI 芯片软硬件栈里最容易被低估的一环。硬件做得再好如果编译器映射得差实际性能可能只有峰值的 20%。我见过太多项目硬件指标很漂亮但因为没有好的编译器客户根本用不起来。编译器的核心工作可以拆成几步图优化、算子融合、内存分配、指令调度。图优化阶段会做常量折叠、死代码消除、算子替换等。比如把连续的 Conv-BN-ReLU 融合成一个算子减少中间结果的访存。算子融合是最影响性能的一步。以 Transformer 为例一个注意力层涉及十几个算子如果每个算子都单独执行中间结果要反复读写外部内存。编译器需要识别出可以融合的算子序列把它们合并成一个超级算子让中间结果留在片上。内存分配决定了数据在片上缓冲和外部内存之间怎么搬。AI 芯片的片上缓冲通常只有几 MB 到几十 MB而模型的权重和激活可能有几百 MB必须做精细的分块和调度。好的内存分配策略能把访存次数降到理论下限。指令调度则是把融合后的算子映射到具体的硬件单元安排执行顺序尽量让脉动阵列和 VPU 并行工作避免硬件空转。4.3 一个实际的编译优化案例我拿一个具体的例子说明编译器优化的威力。假设我们要在一个 256×256 脉动阵列的芯片上跑一个隐藏层维度 768 的 Transformer 层。如果不做任何优化编译器把 768×768 的权重矩阵直接丢给阵列阵列一次只能处理 256×256需要分 9 块处理每块之间还要把部分和写回内存再读出来访存开销巨大。优化后的做法是把权重矩阵按 256 分块激活值也按 256 分块让阵列连续处理多个块部分和留在阵列的累加器里不写回。这样 9 次矩阵乘法变成了一次流水线访存次数从 18 次降到 2 次性能提升接近 5 倍。这个优化的关键在于编译器要理解硬件的累加器行为知道部分和可以在阵列里累积而不必写回。这就要求编译器和硬件团队紧密协作硬件提供能力编译器负责利用。实操心得编译器和硬件团队一定要坐在一起设计。我经历过硬件团队先做完硬件、编译器团队再接手的情况结果编译器发现硬件缺少某些关键能力比如部分和累加、灵活的地址生成只能绕路实现性能损失惨重。软硬件协同设计不是口号是必须落到接口定义和早期评审上的。5. 实操从零搭建一个脉动阵列仿真模型5.1 为什么建议先做仿真在流片之前验证架构设计是否合理最经济的方式是做行为级仿真。我通常会用 Python 或 C 写一个脉动阵列的周期级仿真模型输入真实的模型权重和激活观察阵列的利用率、访存次数、计算周期数等指标。这个仿真模型不需要精确到门级但需要准确反映数据流和时序。比如每个 PE 的乘加延迟、数据在阵列里的传播周期、缓冲的读写带宽这些都要建模。5.2 仿真模型的核心结构我用 Python 写一个简化版的脉动阵列仿真核心结构如下import numpy as np class PE: def __init__(self): self.weight 0 self.acc 0 def compute(self, activation): self.acc self.weight * activation return self.acc class SystolicArray: def __init__(self, rows, cols): self.rows rows self.cols cols self.pe [[PE() for _ in range(cols)] for _ in range(rows)] def load_weight(self, weight_matrix): # 权重从上方加载驻留在 PE 中 for i in range(self.rows): for j in range(self.cols): self.pe[i][j].weight weight_matrix[i][j] def run(self, activation_vector): # 激活从左侧流入每个周期推进一列 results [] for cycle in range(self.cols self.rows - 1): for i in range(self.rows): col cycle - i if 0 col self.cols: self.pe[i][col].compute(activation_vector[i]) # 收集最右侧 PE 的输出 if cycle self.cols - 1: results.append(self.pe[cycle - self.cols 1][self.cols - 1].acc) return results这个模型虽然简化但能反映脉动阵列的核心时序权重加载一次激活逐周期流入结果逐周期流出。用这个模型跑不同尺寸的矩阵就能估算出阵列利用率和总周期数。5.3 关键参数的估算方法仿真跑起来之后需要关注几个关键指标阵列利用率实际计算量除以阵列峰值算力。如果矩阵维度是 768阵列是 256×256利用率大约是 (768/256)^2 分块后的填充率。实际算下来768 能被 256 整除利用率接近 100%但如果维度是 700就会有很多 PE 空转利用率可能掉到 70% 以下。访存带宽需求每周期需要从外部内存读多少数据。这个指标决定了内存子系统的设计。如果阵列每周期需要 256 字节的激活输入而外部内存带宽只有 128 字节/周期那阵列就会饿死。功耗估算用综合工具跑一下 PE 的面积和功耗乘以 PE 数量再加上缓冲和互连的功耗就能得到芯片的总功耗估算。这个数字要和目标场景的功耗预算对比如果超了就要回头调整阵列尺寸。6. 常见问题与排查技巧实录6.1 阵列利用率上不去怎么办这是最常见的性能问题。排查思路是先看矩阵维度再看分块策略最后看调度。如果矩阵维度不是阵列尺寸的整数倍利用率天然就上不去。解决办法有两个一是调整阵列尺寸匹配主力维度二是用编译器做 padding把维度补齐到整数倍。padding 会浪费一些计算但比 PE 空转好。如果分块策略有问题比如块切得太小导致部分和频繁写回就要优化分块大小。经验值是块越大越好直到片上缓冲放不下为止。如果调度有问题比如脉动阵列和 VPU 串行执行而不是并行就要重新安排指令顺序让两类单元尽量同时工作。6.2 精度损失怎么定位AI 芯片的精度问题往往出在几个地方矩阵乘法的累加位宽不够、逐元素操作的逼近误差、数据量化的截断误差。定位方法是逐层对比用芯片跑一遍用浮点参考模型跑一遍逐层对比输出差异。哪一层的误差突然变大问题就在那一层。然后检查那一层用到的硬件单元和精度配置。我遇到过一个典型案例某层的 LayerNorm 输出误差特别大查下来是开方运算的逼近精度不够。把开方单元从 8 位精度提到 12 位误差就回到了可接受范围。6.3 常见问题速查表问题现象可能原因排查方法解决思路阵列利用率低于 50%矩阵维度不匹配统计各层维度分布调整阵列尺寸或做 padding性能远低于峰值访存瓶颈分析访存次数和带宽优化分块和算子融合精度掉点严重累加位宽或逼近精度不足逐层对比浮点参考提高关键单元精度功耗超标阵列过大或访存过多功耗分解分析缩小阵列或优化数据复用编译器报错算子不支持查看不支持的算子列表用组合指令模拟或扩展指令集避坑技巧做架构评估时一定要用真实的模型和真实的输入数据跑仿真不要用随机数据。我见过用随机数据评估时利用率 90%换成真实模型后掉到 40% 的情况因为真实模型的维度分布和随机数据完全不同。7. 软硬件协同设计的工程实践体会做 AI 芯片这些年我越来越觉得这个领域的核心竞争力不在单点技术而在软硬件团队能不能真正协同。硬件团队懂时序、懂面积、懂功耗但不一定懂模型算法团队懂模型、懂精度但不一定懂硬件约束编译器团队夹在中间既要理解硬件能力又要理解模型需求。一个健康的协作模式是算法团队提前给出目标模型的算子分布和维度特征硬件团队据此设计数据流和指令集编译器团队在架构定义阶段就介入确保硬件能力能被编译器充分利用。三方定期对齐而不是各做各的。另一个体会是不要追求通用。AI 芯片和 CPU 不一样CPU 要支持所有程序AI 芯片只需要支持目标模型。把目标模型吃透针对性地优化比做一个什么都能跑但什么都跑不快的通用加速器有价值得多。TPUv1 的成功就是最好的例子——它不支持卷积不支持训练只支持推理时的矩阵乘法但它在自己的目标场景里做到了极致。最后分享一个我在实际项目中总结的判断标准如果一个架构设计需要编译器做大量聪明的优化才能跑出好性能那这个架构大概率有问题。好的架构应该让编译器的工作尽量简单直接硬件本身就把数据流和复用设计好编译器只需要做映射和调度。软硬件协同的最高境界是硬件把复杂性吃掉让软件层看到的是一个干净、规则、易用的接口。
