AI 芯片这个话题聊架构、聊算力、聊制程的文章已经够多了但真正落到软硬件怎么配合设计这个层面能讲清楚的人并不多。我做了几年加速器相关的软硬件协同工作从 RTL 仿真到编译器算子映射都趟过一遍深知这里面的坑有多深。这篇就围绕 AI 芯片软硬件设计这个主题把脉动阵列的基本原理、Transformer 这类模型在硬件上的映射逻辑、以及 TPUv1 当年为什么那么设计掰开揉碎讲一遍。不管你是刚入行的数字 IC 设计工程师还是做模型部署想理解底层硬件的算法同学又或者只是对AI 芯片到底难在哪好奇的技术爱好者这篇都能给你一些能直接用的东西。我会尽量少堆公式多用类比和实际设计中的取舍来讲因为软硬件协同这件事本质上就是一连串为什么这么选的决策链。1. 为什么 AI 芯片的软硬件必须一起设计1.1 通用处理器跑神经网络到底卡在哪先想一个问题为什么不用现成的 CPU 或 GPU 直接跑神经网络非要专门设计 AI 芯片很多人第一反应是CPU 算力不够这个答案对但太粗糙。真正的问题在于数据搬运的开销远大于计算本身。我拿一个具体的例子算给你看。假设一个全连接层输入维度 1024输出维度 1024权重矩阵就是 1024×1024用 FP16 存储大约 2MB。每次推理要读这 2MB 权重做大约 200 万次乘加运算。在 CPU 上这 2MB 数据要从主存搬到缓存、再搬到寄存器而 CPU 的算力单元其实大部分时间在等数据。这就是所谓的内存墙——算力增长快但内存带宽增长慢两者差距越拉越大。GPU 缓解了一部分问题因为它有几千个核心可以并行吞吐上去了。但 GPU 的功耗和面积效率对专用场景来说仍然浪费它的调度逻辑、缓存层级是为通用图形计算设计的跑矩阵乘的时候有大量控制开销。AI 芯片的核心思路就是把数据流设计成让计算单元尽量不空等把控制逻辑砍到最简把面积和功耗全砸在乘加阵列和片上存储上。这就引出了软硬件协同设计的根本矛盾硬件是固定的、并行的、有严格时序的而模型是变化的、算子多样的、精度需求各异的。硬件设计者如果不懂模型的计算模式就会设计出看起来很猛但实际跑不满的架构算法工程师如果不懂硬件的数据流约束就会写出理论 FLOPs 很高但硬件利用率只有 20%的模型。所以这两拨人必须坐在一起从架构定义阶段就对齐。1.2 软硬件协同的三个层次我把协同设计分成三个层次来理解这样你在实际工作中能快速定位自己处在哪一层。第一层是架构层协同。这一层决定芯片的宏观结构用什么样的计算阵列、片上缓存多大、片外带宽多少、支持哪些数据精度。这一层的决策一旦流片就无法更改所以必须在设计前就把目标模型的计算特征摸清楚。比如你要跑 Transformer那注意力机制里的矩阵乘维度变化很大阵列就不能设计得太死板。第二层是编译层协同。硬件定了之后怎么把模型算子映射到硬件上就是编译器的事。这一层要解决算子融合、数据切分、流水线调度、精度量化等问题。编译器做得好不好直接决定硬件利用率。我见过同一个硬件换个编译器性能差三倍的情况。第三层是运行时协同。实际部署时batch size 会变、输入长度会变、多模型可能混跑运行时调度要动态分配硬件资源。这一层最容易被忽视但对实际体验影响巨大。提示很多团队把这三层割裂开架构师画完框图就交给编译器团队编译器团队做完就交给部署团队结果每一层都在为上一层的不合理设计打补丁。真正高效的团队会让三个层次的人定期同步甚至让编译器工程师参与架构评审。1.3 一个反直觉的结论算力不是最重要的指标新手看 AI 芯片参数第一眼盯的是 TOPS每秒万亿次运算。但实际部署中有效算力才是关键也就是在真实模型上能跑出来的算力。一块标称 256 TOPS 的芯片如果跑某个模型只能达到 30 TOPS 的有效算力那它还不如一块标称 64 TOPS 但利用率能到 80% 的芯片。有效算力 标称算力 × 硬件利用率。而硬件利用率取决于数据复用做得够不够、访存瓶颈有没有卡住、算子映射是否匹配阵列形状。这些全是软硬件协同要解决的问题。所以我在评估一块 AI 芯片时从来不只看纸面参数而是会问跑 ResNet-50 的利用率是多少跑 BERT-base 呢跑一个自定义的注意力层呢这些数字才说明问题。2. 脉动阵列AI 芯片最经典的计算骨架2.1 脉动阵列到底脉动在哪里脉动阵列Systolic Array这个名字听起来玄乎其实核心思想特别朴素让数据像心跳一样有节奏地在计算单元之间流动每个计算单元只做最简单的乘加数据流过就出结果。我打个比方。想象一条流水线上面站着一排工人每个工人手里拿着一个权重值。数据从左边一个个传进来每经过一个工人就做一次乘法并累加然后传给下一个工人。同时权重从上往下传数据从左往右传两者在交叉点相遇就做一次乘加。这样一个矩阵乘法的所有运算就能在数据流动的过程中完成不需要反复从内存里取数。这种设计的精妙之处在于每个计算单元PE只需要一个乘法器、一个加法器、几个寄存器没有复杂的地址计算没有指令译码。控制逻辑极简面积和功耗几乎全用在计算上。这就是为什么 TPUv1 能在那个年代做到那么高的能效比。脉动阵列的脉动体现在数据是有节奏地、周期性地流动的。每个时钟周期数据在阵列里前进一格就像心脏跳动推动血液流动一样。这种规律性让时序设计变得简单也让数据复用达到极致——一个权重值被加载进 PE 后可以被多个输入数据复用不需要重复访存。2.2 数据流设计权重固定还是输出固定脉动阵列不是只有一种形态根据数据在阵列里的流动方式主要分几种数据流。这个选择直接决定硬件的访存模式和编译器的工作方式是架构设计里最关键的决策之一。数据流类型权重流动输入流动输出流动适用场景权重固定WS静止横向流动纵向累加权重复用高的卷积输出固定OS横向流动纵向流动静止大矩阵乘行固定RS静止对角流动横向输出平衡型设计权重固定Weight Stationary是最常见的方案。权重预先加载到 PE 里就不动了输入数据横向流过部分和纵向累加。这种方案适合卷积神经网络因为卷积核权重在整个特征图上反复使用权重固定能最大化复用。输出固定Output Stationary则是让输出部分和留在 PE 里不动权重和输入都流动进来。这种方案适合大矩阵乘因为输出累加不需要反复搬运省了带宽。实际芯片里往往是混合的。TPUv1 用的是权重固定为主的脉动阵列256×256 的乘加阵列专门为卷积和全连接层优化。但到了 Transformer 时代矩阵乘的形态变了纯权重固定的阵列就不一定最优了这也是为什么后来的 AI 芯片开始支持更灵活的数据流配置。2.3 阵列尺寸怎么定一个真实的权衡过程假设你要设计一个脉动阵列尺寸定多大256×256 还是 128×128这不是拍脑袋决定的要算账。阵列面积大致和尺寸的平方成正比。256×256 的阵列有 65536 个 PE每个 PE 假设占 1000 平方微米含乘法器、加法器、寄存器光阵列就占 65 平方毫米这在很多工艺下已经是大芯片了。而 128×128 只有 16384 个 PE面积降到四分之一。但尺寸小意味着单次能处理的矩阵块小需要更多次的数据加载和部分和累加访存开销上升。这里有个经验公式阵列尺寸应该匹配目标模型里最常见的矩阵维度。如果目标模型的主力矩阵乘是 512×512那 256×256 的阵列刚好可以分四块处理复用效率高如果主力是 128×128那 256 的阵列就有一半 PE 在空转。我在实际项目中遇到过这个问题一开始定了 256×256结果发现目标模型里大量小矩阵乘利用率只有 40%。后来改成 128×128 加多阵列并行利用率提到 75%。所以阵列尺寸不是越大越好匹配才是王道。注意阵列尺寸还受限于时钟频率。阵列越大信号在阵列里传播的路径越长时序收敛越难频率就上不去。256×256 的阵列想跑到 1GHz 以上在先进工艺下都很有挑战。很多时候降频换面积是划算的。2.4 脉动阵列的软肋不规则计算怎么办脉动阵列最大的问题是它只擅长规则的矩阵乘和卷积。一旦遇到不规则的计算——比如注意力机制里的 softmax、LayerNorm、各种激活函数、稀疏操作——阵列就使不上劲了。这些操作怎么办通常是在阵列旁边挂一个向量处理单元VPU专门处理这些逐元素或归约操作。阵列负责矩阵乘这种大块头VPU 负责零碎活。软硬件协同在这里的体现就是编译器要能识别哪些算子该丢给阵列哪些该丢给 VPU并且安排好两者的流水线别让阵列等 VPU也别让 VPU 等阵列。我见过一个设计阵列很强但 VPU 太弱结果跑 Transformer 的时候注意力里的 softmax 成了瓶颈阵列大部分时间在等 VPU 算完。后来把 VPU 的并行度提高了一倍整体性能提升了 40%。这个教训说明木桶效应在 AI 芯片里特别明显软硬件设计要均衡不能有明显的短板。3. Transformer 上硬件比 CNN 难在哪3.1 注意力机制的计算特征拆解Transformer 的核心是自注意力Self-Attention它的计算可以拆成几步输入分别乘以三个权重矩阵得到 Q、K、V然后 Q 和 K 的转置做矩阵乘得到注意力分数分数经过 softmax 归一化最后和 V 做矩阵乘得到输出。从硬件角度看这里面有几种截然不同的计算模式。Q、K、V 的生成是标准的矩阵乘脉动阵列很擅长。QK^T 也是矩阵乘但维度是序列长度×序列长度序列长度可变从几十到几千都有可能。softmax 是归约操作需要找最大值、算指数、求和、再归一化这是 VPU 的活。最后和 V 相乘又是矩阵乘。问题在于序列长度是动态的。CNN 的输入尺寸相对固定硬件可以针对性地优化。但 Transformer 的序列长度在推理时可能变化比如对话场景输入长度每次都不一样。这就要求硬件和编译器能动态适应不能写死。3.2 序列长度动态变化带来的调度难题假设你的脉动阵列是 128×128 的处理 QK^T 时如果序列长度是 128刚好一次搞定如果是 512就要分四次如果是 100那就有 28 行 PE 在空转。这种维度不匹配导致的利用率下降是 Transformer 硬件加速的核心痛点。解决办法有几种。一种是padding把不足的维度补齐简单但浪费算力。一种是动态分块编译器根据实际序列长度选择最优的分块策略这需要硬件支持灵活的数据加载。还有一种是多序列并行把多个短序列拼在一起凑满阵列但这要求注意力计算能正确处理序列边界实现复杂度高。我在实际项目里用过动态分块加多序列并行的组合方案。对于短序列把多个 batch 的序列拼起来填满阵列对于长序列切成块做分块注意力。编译器里要维护一套复杂的调度逻辑但换来的是平均利用率从 50% 提到 80% 以上。这套逻辑写起来很痛苦但效果确实值。3.3 KV Cache推理场景的访存大头Transformer 推理有个特殊之处自回归生成时每生成一个新 token都要和之前所有 token 的 K、V 做注意力计算。为了避免重复计算会把之前算好的 K、V 缓存起来这就是KV Cache。KV Cache 的大小 2 × 层数 × 序列长度 × 隐藏维度 × 精度字节数。以一个 12 层、隐藏维度 768、序列长度 2048、FP16 的模型为例KV Cache 大约是 2×12×2048×768×2 字节约 72MB。这个数字看着不大但它是每生成一个 token 都要完整读一遍的。生成 1000 个 token就要读 72GB 的数据。访存带宽成了绝对瓶颈。所以现代 AI 芯片设计里KV Cache 的管理和搬运是重中之重。硬件上要有大容量的片上缓存或者高带宽的片外存储接口软件上要做 KV Cache 的分页管理、量化压缩、甚至选择性丢弃。这些全是软硬件协同的活。我见过一个优化案例把 KV Cache 做 INT8 量化精度损失很小但访存带宽直接减半推理吞吐翻倍。3.4 从 CNN 到 Transformer硬件设计的范式转移CNN 时代的硬件设计思路是权重复用最大化因为卷积核小且反复使用。Transformer 时代权重矩阵大得多而且注意力计算里 Q、K、V 都是动态生成的没有固定的权重可以长期驻留。这意味着硬件设计要从权重固定转向更灵活的数据流。有些新架构开始支持可重构的互连让 PE 之间的连接方式可以根据算子动态配置。还有些架构增加了专门的注意力加速单元把 QK^T、softmax、AV 这条链路做成流水线。TPUv1 那个年代还没有 Transformer它的设计完全是为 CNN 优化的。后来 TPU 的迭代版本不断增加对 Transformer 的支持比如更大的片上缓存、更灵活的向量单元、对 KV Cache 的专门优化。这个演进过程本身就是软硬件协同的活教材模型在变硬件必须跟着变而硬件的变化又反过来影响模型的设计选择。4. TPUv1 的设计哲学与当代启示4.1 TPUv1 的核心参数与设计取舍TPUv1 是 2015 年左右的设计现在看参数不算亮眼但它的设计思路至今仍有参考价值。它有一颗 256×256 的脉动阵列峰值算力 92 TOPSINT8片上 24MB 的统一缓存用 DDR3 做片外存储。注意几个关键取舍。第一它只做推理不做训练所以不需要支持反向传播不需要存梯度硬件可以大幅简化。第二它主打 INT8 量化牺牲一点精度换能效和面积。第三它用 DDR3 而不是更贵的 HBM因为推理场景对带宽的需求相对可控用便宜的内存降低整体成本。这些取舍背后是清晰的目标定位在数据中心里以低成本、高能效地跑推理。它不追求训练能力不追求最高精度不追求最先进的内存。这种有所不为的设计哲学比堆参数更难也更有价值。4.2 指令集与编译器的配合TPUv1 的指令集非常精简主要就是几条读权重、读输入、矩阵乘、激活、写输出。这种精简指令集的好处是硬件译码简单但代价是编译器要做大量工作。编译器要把高层模型拆成这些基本指令的序列还要安排数据在片上缓存和阵列之间的搬运。TPUv1 的编译器有个关键设计它把整个模型编译成一个指令序列而不是运行时动态调度。这样运行时开销极小但灵活性差模型一变就要重新编译。这种静态编译的思路在推理场景很合适因为推理的模型和输入形状相对固定。但到了 Transformer 这种动态性强的场景静态编译就不够用了需要更动态的调度。这也是为什么后来的 AI 芯片编译器越来越复杂要在静态优化和动态适应之间找平衡。4.3 从 TPUv1 看软硬件协同的边界TPUv1 的成功很大程度上是因为它明确了自己的边界只做推理、只做矩阵乘为主的计算、只支持有限的精度。在这个边界内软硬件可以深度协同把效率做到极致。一旦超出边界性能就断崖式下跌。这对我们的启示是设计 AI 芯片时明确不做什么比做什么更重要。你想支持所有模型、所有精度、所有场景结果就是什么都不精。反过来如果你能清晰定义目标场景软硬件就能围绕这个场景做深度优化。当然现实中的芯片往往要兼顾多个场景那就需要在架构上做模块化设计让不同场景用不同的硬件单元编译器负责调度。但无论如何边界清晰是高效协同的前提。5. 软硬件协同中的实操经验与避坑指南5.1 精度量化软硬件必须一起定量化是 AI 芯片绕不开的话题。把 FP32 降到 INT8算力翻几倍功耗降一半但精度会掉。掉多少这取决于模型、取决于量化策略、也取决于硬件的量化单元怎么设计。我踩过的一个坑硬件团队自己定了量化方案用对称量化、per-tensor 缩放。结果算法团队拿过去一跑某些层的精度掉得厉害因为那些层的权重分布不对称per-tensor 缩放根本不够。后来改成 per-channel 非对称量化精度回来了但硬件要支持更复杂的缩放计算面积增加了一些。这个教训是量化方案必须软硬件一起定。算法团队要告诉硬件团队哪些层对精度敏感需要更细粒度的量化硬件团队要告诉算法团队支持哪些量化模式代价是多少。双方在架构定义阶段就要对齐不能各干各的。5.2 算子融合的边界在哪里算子融合是编译器优化的常用手段把多个算子合并成一个减少中间结果的访存。比如 Conv BN ReLU 融合成一个中间结果不用写回内存再读出来。但融合不是越多越好。融合后的算子可能超出硬件的单次处理能力需要拆分融合可能改变数值精度需要重新验证融合可能让调度变得复杂反而降低利用率。我见过一个过度融合的案例把十几个算子融成一个巨大的 kernel结果寄存器不够用频繁 spill性能反而下降。合理的做法是优先融合访存密集的相邻算子对计算密集的算子保持独立。具体融合到什么程度要用 profiler 实测看融合前后的访存量和利用率变化。这个边界因硬件而异没有通用答案。5.3 编译器与硬件的契约编译器和硬件之间需要一份清晰的契约硬件提供哪些指令、哪些寄存器、哪些内存空间、时序约束是什么编译器承诺生成符合这些约束的代码。这份契约如果模糊就会出现编译器生成的代码跑不对、或者跑得慢的情况。我在项目里坚持做的一件事是维护一份硬件抽象层的精确文档并且用自动化测试保证编译器和硬件的一致性。每次硬件改动都要跑一遍回归测试确保编译器生成的代码仍然正确。这个投入看起来大但比后期调试省事得多。提示如果你的团队还没有硬件抽象层的文档建议现在就补。哪怕只是一份简单的指令列表和寄存器映射表也能省下大量沟通成本。5.4 性能分析别只看端到端数字优化性能时很多人只看端到端的延迟或吞吐。这个数字当然重要但它不告诉你瓶颈在哪。要定位瓶颈必须做细粒度的性能分析每个算子的耗时、每次访存的带宽利用率、阵列的活跃周期比例。我常用的方法是分层 profiling先看端到端找出最慢的几个算子再看这些算子的内部是计算受限还是访存受限如果是访存受限再看是哪一级存储的带宽不够。这样一层层往下挖才能找到真正的瓶颈。有个案例端到端看某个模型推理慢怀疑是阵列不够强。但 profiling 发现阵列利用率其实有 70%真正的问题是片外带宽跑满了数据供不上。解决办法不是换更大的阵列而是优化数据复用把更多数据留在片上。如果只看端到端数字很容易做出错误的优化决策。6. 写在最后的一些个人体会做 AI 芯片软硬件协同这些年我最大的体会是这不是一个纯技术问题而是一个沟通和权衡问题。硬件工程师、编译器工程师、算法工程师三拨人的语言体系不一样关注点不一样很容易各说各话。能把这三拨人拉到一张桌子上用共同的语言讨论问题比任何单点技术突破都重要。另一个体会是不要迷信纸面参数要相信实测数据。TOPS 再高利用率上不去也是白搭。我在评估任何方案时都会先跑一个真实模型看实际性能再决定要不要深入。这个习惯帮我避开了很多看起来很美的坑。最后AI 芯片这个领域变化太快今天的最优解明天可能就过时了。保持学习保持对模型演进的敏感保持和一线工程师的交流比死守某个架构方案更重要。Transformer 之后还会有新的架构硬件设计要留出足够的灵活性去适应而不是把宝全押在某一种计算模式上。