1. 从TPUv1的脉动阵列说起为什么AI芯片的硬件设计绕不开数据复用聊AI芯片的软硬件设计如果只盯着算力数字看很容易掉进一个误区——觉得堆的乘加单元越多芯片就越强。实际上真正决定一颗AI芯片能不能跑出理论峰值的关键在于数据在片上的流动方式。TPUv1作为早期专门为神经网络推理设计的芯片它最核心的贡献不是把频率拉多高而是用了一个非常聪明的结构脉动阵列Systolic Array。这个词听起来有点玄但拆开看其实很朴素——数据像心跳一样有节奏地在计算单元之间流动每个周期都有新的数据进来同时有结果流出去计算单元本身几乎不需要复杂的地址索引逻辑。我第一次接触脉动阵列的时候脑子里冒出来的问题是为什么不让每个乘法器自己去内存里取数答案很直接——取数太贵了。芯片里做一次乘加运算的能耗可能只有从DRAM搬一个字节数据的几百分之一。所以AI芯片设计的核心矛盾从来不是“算不动”而是“喂不饱”。脉动阵列解决的正是这个喂数据的问题它让权重在阵列里预先排好激活值从左侧流入部分和从上往下累积每个计算单元只跟相邻的单元打交道数据复用率被拉到极高。TPUv1的256×256阵列一个周期能完成65536次乘加而权重只需要在阵列里驻留一次就能服务整批输入。这里有个容易被忽略的细节脉动阵列的“脉动”是有方向的。通常权重是预先加载并保持静止的激活值横向流动部分和纵向流动。这种设计让每个PEProcessing Element只需要一个乘法器、一个加法器和几个寄存器面积和功耗都压得很低。但代价是灵活性——阵列一旦固定就很难高效处理稀疏、不规则的计算模式。这也是为什么后来的AI芯片在脉动阵列之外还要搭配向量单元、标量单元和专门的激活函数硬件。从软硬件协同的角度看脉动阵列的存在直接决定了编译器要做什么。你不能指望程序员手写阵列调度必须有一套映射工具把高层算子拆成适合阵列尺寸的tile再安排数据搬运的节奏。TPUv1的软件栈里XLA编译器就承担了这个角色它把计算图切分成适合256×256阵列的块生成指令序列控制DMA把数据从片外搬进来再按节拍送进阵列。这套流程里任何一个环节的节奏对不上阵列就会出现气泡算力利用率立刻掉下来。所以理解AI芯片的软硬件设计第一步不是去看某个具体型号的参数表而是先建立“数据复用优先”的思维。脉动阵列只是其中一种实现但它把问题本质暴露得很清楚硬件结构决定了数据流动的边界软件必须在这个边界内找到最优的调度方式。后面要聊的Transformer、MoE、量化、编译优化全都建立在这个基础之上。2. 脉动阵列的工程边界它擅长什么又在什么地方会卡住2.1 权重静止与激活流动的代价脉动阵列最舒服的场景是卷积和全连接层因为这两类计算的权重可以在推理开始前就确定下来整批数据共享同一组权重。TPUv1的设计就是围绕这个假设展开的权重从阵列上方或侧边加载一次然后保持不动激活值像流水一样从左边灌进去每个周期推进一格。这种模式下数据复用率极高功耗也控制得很好。但问题在于Transformer这类模型把计算模式彻底改了。自注意力机制里Q、K、V三个矩阵都是动态生成的权重不再静止而是随输入变化。更麻烦的是注意力矩阵的尺寸跟序列长度平方相关序列一长矩阵就变得又大又稀疏。脉动阵列处理这种动态、不规则的矩阵乘时效率会明显下降。你当然可以把QK^T拆成tile硬塞进阵列但每个tile的权重都不一样预加载的优势就没了数据搬运的开销反而成了瓶颈。我在实际项目里见过一个典型的误判团队看到某款芯片标称算力很高就直接把Transformer模型往上搬结果实测吞吐只有理论值的百分之十几。排查下来问题不在算力而在数据路径——注意力计算需要频繁地在不同存储层级之间搬Q、K、V而脉动阵列的输入缓冲区根本喂不上这个节奏。后来他们改成先做算子融合把QK^T和softmax尽量压在片上完成才把利用率拉回来。2.2 稀疏性与不规则计算的适配难题脉动阵列的另一个边界是稀疏性。神经网络剪枝之后权重矩阵里会出现大量零值理论上可以跳过这些零减少计算量。但脉动阵列的每个PE是按固定节拍工作的你很难让某个PE单独停下来因为数据是在阵列里同步流动的。跳过零值意味着要打乱流动节奏这会破坏整个阵列的同步性。业界常见的做法是在阵列之外加一个稀疏处理单元先把稀疏矩阵压缩成稠密块再送进阵列。但压缩和解压本身也有开销而且压缩后的块尺寸往往不规整阵列利用率还是会打折。更激进的做法是设计专门支持稀疏计算的阵列结构比如让PE之间可以动态路由但这会显著增加控制逻辑的复杂度面积和功耗都会上去。从软件角度看编译器需要知道硬件的稀疏支持能力才能决定剪枝策略。如果硬件只能处理结构化稀疏比如每四个元素里跳过一个那剪枝算法就必须按这个约束来设计不能随便乱剪。这种软硬件之间的约束传递是AI芯片设计里最容易被低估的环节。很多团队硬件做完了才让软件团队去适配结果发现硬件的一些“优化”在软件层面根本用不上或者用起来代价极大。2.3 阵列尺寸与tile划分的权衡TPUv1选了256×256的阵列这个尺寸不是拍脑袋定的。阵列越大单个周期能做的乘加越多但权重加载的时间和片上存储的需求也越大。如果模型某一层的权重矩阵只有128×128那256×256的阵列就有一半是闲置的。反过来如果阵列太小又需要频繁地分块和累加部分和的搬运开销会吃掉收益。实际设计里阵列尺寸往往跟目标模型的工作负载分布有关。如果主要跑的是BERT-base这类模型中间层维度大多是768或3072那阵列尺寸就要围绕这些数字来选尽量让tile划分整齐减少边缘浪费。编译器在这里的作用非常关键它需要根据阵列尺寸和片上缓存大小自动选择最优的tile大小和循环顺序。同一个矩阵乘不同的tile策略性能可能差好几倍。我个人的经验是评估一款AI芯片时不要只看它的峰值算力一定要问清楚它的阵列尺寸、片上缓存容量和编译器支持的tile策略。这三者决定了它在真实模型上的有效算力。很多标称几百TOPS的芯片跑起Transformer来还不如一些标称低得多的芯片原因就在这儿。3. Transformer把AI芯片设计带进了动态调度的深水区3.1 自注意力机制对硬件的新要求Transformer和CNN最大的区别在于它的计算图是动态的。CNN的卷积核权重在推理时是固定的硬件可以提前把权重排好数据流非常规整。但Transformer的自注意力里Q、K、V都是输入的函数每个token的权重都不一样。这意味着硬件不能依赖“权重静止”这个假设必须支持动态加载和动态调度。更具体地说自注意力包含几个关键步骤QK^T矩阵乘、softmax归一化、再跟V做矩阵乘。这三个步骤的数据依赖关系很强中间结果需要暂存而且softmax涉及指数运算和除法对硬件来说比单纯的乘加复杂得多。TPUv1那个时代这些操作主要靠向量单元和标量单元配合完成效率并不高。后来的芯片开始专门为softmax设计硬件单元比如用查找表实现指数函数用近似除法减少延迟。另一个挑战是序列长度。自注意力的计算量跟序列长度的平方成正比序列一长注意力矩阵就变得非常大。但很多注意力权重其实很小接近零理论上可以稀疏化。问题还是回到硬件脉动阵列处理不了这种动态稀疏需要专门的稀疏注意力硬件或者算法层面的近似。现在业界比较流行的做法是分块注意力或者线性注意力把复杂度从平方降到线性但代价是精度会有些损失需要根据任务来权衡。3.2 KV Cache带来的存储压力与带宽瓶颈推理阶段Transformer有一个绕不开的问题KV Cache。自回归生成时每生成一个新token都需要跟之前所有token的K和V做注意力计算。为了避免重复计算K和V会被缓存下来。序列越长KV Cache越大对片上存储和内存带宽的压力就越大。这个问题的本质是存储带宽跟不上计算需求。生成一个token需要读取整个KV Cache而KV Cache的大小跟序列长度和模型层数成正比。对于长序列场景KV Cache可能达到几十甚至上百MB远超片上SRAM的容量必须放到DRAM里。但DRAM的带宽有限读取KV Cache的时间可能比计算本身还长导致芯片算力大量闲置。硬件层面的应对策略主要有几个方向一是增大片上缓存把KV Cache尽量留在片内二是用更激进的量化把KV Cache压到8位甚至4位三是用PagedAttention这类软件技术把KV Cache分页管理减少内存碎片和无效读取。但每种方案都有代价增大缓存会增加面积和成本量化会损失精度分页管理会增加软件复杂度。软硬件团队必须一起权衡找到适合目标场景的平衡点。3.3 MoE与思维链对调度器的考验MoE混合专家模型把FFN层拆成多个专家每个token只激活其中一小部分。这种设计能在不显著增加计算量的前提下扩大模型容量但对硬件调度提出了很高要求。因为每个token走的专家路径可能不同计算变得非常不规则。脉动阵列这种规整结构处理起来很吃力需要配合动态路由和负载均衡机制。思维链Chain-of-Thought则是另一个维度的挑战。它让模型在推理时生成中间步骤序列长度动态变化而且每一步的计算量都不确定。这对硬件的动态调度能力要求极高编译器需要在运行时根据实际序列长度调整tile划分和资源分配。传统的静态编译流程很难应对这种动态性需要引入运行时调度或者即时编译。从软硬件协同的角度看MoE和思维链代表了一个趋势AI芯片不能只做规整的矩阵乘还必须具备处理动态、不规则计算的能力。这意味着硬件架构要更灵活编译器要更智能运行时系统要能快速响应负载变化。TPUv1那种纯脉动阵列的设计在这个趋势下已经不够用了后来的芯片都在往“阵列向量标量专用单元”的异构方向走。4. 软硬件协同的落地路径从算子映射到推理部署4.1 算子到硬件的映射逻辑把Transformer模型部署到AI芯片上第一步是算子映射。框架里的算子比如MatMul、Softmax、LayerNorm需要被拆解成硬件能执行的指令序列。这个过程通常由编译器完成但编译器能做的优化受限于硬件暴露的接口。以矩阵乘为例编译器需要决定把矩阵切成多大的tiletile按什么顺序循环数据什么时候从DRAM搬到片上部分和什么时候写回这些决策直接影响性能。如果硬件支持双缓冲编译器就可以让数据搬运和计算重叠起来减少等待时间。如果硬件不支持那搬运和计算就只能串行利用率会掉很多。我在实际项目里总结出一条经验评估编译器好不好用不要看它支持多少算子而要看它对关键算子的调度质量。拿一个标准的Transformer层让编译器生成指令然后看它怎么安排QK^T、softmax和PV这三个矩阵乘的流水线。好的编译器会把softmax尽量压在片上避免中间结果反复进出DRAM差的编译器可能每个算子都单独处理数据搬来搬去性能差好几倍。4.2 量化与精度取舍的实操细节量化是AI芯片部署里绕不开的一环。把FP32压到INT8模型大小和带宽需求都能降四倍算力利用率也能提升。但量化不是简单地把浮点数截断成整数需要处理数值范围、零点偏移和舍入误差。实际操作里我通常分几步走先做训练后量化PTQ用一小批校准数据统计每层的激活值范围确定缩放因子如果精度掉得太多再考虑量化感知训练QAT在训练时模拟量化误差让模型自己适应。对于Transformer注意力层的softmax输出范围很敏感通常需要保留更高精度或者用专门的近似算法。LayerNorm的均值和方差计算也容易受量化影响需要特别处理。KV Cache的量化更棘手因为它是动态增长的数值范围会随序列长度变化。常见的做法是按通道或者按页做量化每页单独统计范围减少误差。但这样会增加元数据的存储开销需要权衡。我见过一些团队为了省事直接对整个KV Cache用一组缩放因子结果长序列下精度崩得很厉害。后来改成分页量化精度才回来。4.3 推理部署中的流水线编排推理部署不是把模型跑起来就完事了还要考虑流水线编排。尤其是服务端场景多个请求并发进来硬件资源需要动态分配。如果每个请求都单独跑一遍模型KV Cache和权重会反复加载效率很低。更好的做法是把多个请求的输入拼成一批共享权重加载提高阵列利用率。但批处理也有代价批越大延迟越高。对于在线服务延迟是硬指标不能无限增大批次。所以需要在吞吐和延迟之间找平衡点。常见的策略是动态批处理请求进来后先攒一小会儿攒到一定数量或者等到超时再一起送进模型。这样既能提高利用率又不会让单个请求等太久。流水线编排还涉及算子融合。把多个小算子合并成一个大算子可以减少中间结果的搬运和kernel启动开销。比如把LayerNorm和后面的矩阵乘融合在一起中间结果不用写回DRAM直接在片上传递。但融合会增加编译器的复杂度而且不是所有硬件都支持。如果硬件没有可编程的片上缓存融合就很难做。5. 踩过的坑与实测经验那些文档里不会写的细节5.1 阵列利用率低先查数据路径而不是算力有一次调一个Transformer推理任务实测吞吐只有理论峰值的15%。第一反应是算力不够但算了一下模型的计算量明明远低于芯片标称值。后来用性能分析工具抓了一下发现阵列大部分时间在等数据。问题出在数据路径上Q、K、V三个矩阵的加载没有重叠好每次矩阵乘之前都要等数据搬完阵列空转。解决办法是调整编译器的调度策略把Q、K、V的加载提前跟上一层的计算重叠起来。同时把softmax的输出直接留在片上不要写回DRAM再读出来。改完之后利用率提到了60%以上。这个经历让我明白AI芯片的性能瓶颈往往不在计算单元而在数据搬运。排查性能问题先看数据路径再看算力。5.2 KV Cache量化不是越激进越好KV Cache量化能省带宽但精度损失不是线性的。我试过把KV Cache压到4位短序列下精度还行序列一长就崩了。原因是长序列下注意力权重的分布变得更尖锐量化误差被放大。后来改成8位并且按页做缩放精度就稳住了。还有一个细节KV Cache的量化缩放因子需要动态更新。如果整个推理过程用一组固定的缩放因子遇到数值范围变化大的输入就会出问题。比较好的做法是每隔一段时间重新统计一次范围或者用滑动窗口的方式更新。但这会增加运行时开销需要根据场景权衡。5.3 编译器tile策略对性能的影响远超预期同一个矩阵乘不同的tile策略性能可能差三到五倍。我做过一个对比实验对于768×768的矩阵乘tile大小从64×64到256×256都试了一遍。结果发现tile太小会导致部分和搬运次数增加tile太大又会导致片上缓存不够出现溢出。最优的tile大小跟阵列尺寸、缓存容量和数据类型都有关没有万能答案。更麻烦的是不同层的矩阵乘维度不一样最优tile策略也不同。如果编译器只用一套固定的tile策略某些层就会吃亏。好的编译器会根据每层的维度动态选择tile大小甚至在同一层内根据循环位置调整。这个优化对性能的影响比换一款芯片还大。5.4 动态形状是部署阶段最大的不确定性训练时通常用固定序列长度但推理时序列长度是变化的。如果编译器只针对固定长度优化遇到变长输入就会频繁重新编译延迟飙升。我见过一个服务因为没处理好动态形状每次请求都要重新编译一遍响应时间从几十毫秒涨到几秒。解决办法是在编译时预留多个形状档位比如128、256、512、1024运行时根据实际长度选择最接近的档位。对于超出最大档位的输入再走动态编译路径。这样大部分请求都能命中预编译的档位延迟可控。但档位设置也有讲究档位太多编译时间和存储开销大档位太少命中率低。需要根据实际流量分布来定。6. 从TPUv1到今天的AI芯片软硬件设计的变与不变回头看TPUv1它的设计哲学非常清晰用脉动阵列把矩阵乘的能效做到极致其他事情交给配套的向量和标量单元。这个思路在CNN时代非常成功因为CNN的计算模式规整权重静止数据复用率高。但到了Transformer时代计算模式变了动态性、稀疏性和不规则性成了常态纯脉动阵列的局限性就暴露出来了。今天的AI芯片设计核心挑战已经从“如何把矩阵乘做得更快”变成了“如何让数据在正确的时间出现在正确的位置”。脉动阵列依然重要但它只是整个数据路径中的一环。芯片里还需要有灵活的片上缓存、可编程的向量单元、专用的softmax和归一化硬件以及一套能感知硬件约束的编译器。软硬件之间的边界越来越模糊做硬件的人要懂模型做软件的人要懂架构否则做出来的东西很难在真实场景里跑出好性能。我个人在实际项目里的体会是AI芯片的评估不能只看纸面参数一定要拿目标模型跑一遍端到端的推理看有效吞吐、延迟和精度。很多问题只有在真实负载下才会暴露出来比如KV Cache的带宽瓶颈、动态形状的编译开销、量化误差的累积。这些细节文档里通常不会写但恰恰是决定项目成败的关键。最后分享一个小技巧如果你在调AI芯片的性能先别急着改模型或换硬件用性能分析工具把每个算子的耗时和数据搬运量抓出来画一张时间线图。十有八九瓶颈不在计算而在数据搬运或者同步等待。把数据路径理顺了性能自然就上来了。这个思路在TPUv1时代适用在今天依然适用。