我直接开讲。十多年前大家聊起硬件加速脑子里蹦出来的还是GPU跑深度学习那套。直到2016年Google把TPU的论文发出来很多人才意识到原来为了跑神经网络真有人敢从零开始造一颗专用芯片。而我拿到这篇论文读了三遍之后的感受是这哪里是一颗芯片这就是一个把“极致专一”四个字写进硅片里的计算怪兽。这篇博文我不打算照着论文给你翻译一遍那太枯燥。我会从一个从业者的视角把Google TPU的架构原理拆开揉碎聊清楚它为什么长成那样它到底解决了什么问题以及有哪些设计思路到今天依然值得我们做AI Infra、做高性能计算甚至做应用性能优化的时候借鉴。不管你是做AI后端开发、算法工程化、还是对芯片设计有兴趣这篇文章都应该能给你一些不一样的启发。1. 内容整体设计与思路拆解1.1 为什么会有TPU它到底在解决谁的痛点要理解TPU先得理解2015年前后的尴尬局面。那会儿深度学习已经火起来了ImageNet分类、语音识别、机器翻译一个接一个被神经网络拿下。但算法跑得越快算力缺口就越大。大家手里的主流武器还是NVIDIA的GPUCUDA生态也变得成熟。可问题在于大规模部署推理任务的时候GPU架构的“通用性”反而变成了累赘GPU要照顾图形渲染的变形、光照、纹理还要兼顾科学计算的各种精度轮到纯跑神经网络推理的时候里面有大量晶体管和带宽其实在“陪跑”。Google的做法很直接我不需要一块能玩3D游戏的显卡我也不需要一块能跑分子动力学的加速卡我只需要一块能把深度神经网络推理过程里的矩阵乘法和激活函数算到飞起的芯片。于是TPU诞生了。它的定位从第一天起就非常清楚ASIC专用集成电路服务推理阶段。相比GPU的“什么都行但什么都不算极致”TPU选择了“我不管别的我把你推理时每秒要做的矩阵运算次数堆到最高”。1.2 从“拆掉没用的功能”开始做设计TPU的设计思路如果用一句话概括就是“找到你计算里唯一的核心然后不计代价地把这条路做到最短”。神经网络推理的计算核心是什么很多人第一反应是卷积但落实到硬件层面无论你是什么卷积、全连接、池化最终都逃不过大量的矩阵乘法和向量乘法。而矩阵乘法的本质就是“乘累加”操作。TPU的每一个计算周期都在围绕乘累加优化。它不关心double精度、不关心复杂分支预测、不关心乱序执行。因为推理阶段的模型一旦训练完它的网络结构就固定了整个计算路径是可控且可预测的。这种“因为确定所以极致”的哲学可以说把ASIC的优势发挥得淋漓尽致。1.3 为什么不上GPU而是另起炉灶很多人问过我Google这种体量的公司多买点GPU把集群堆起来不就行了答案是成本和能效的问题。以TPU v1的论文数据为例它把芯片放在硬盘的2.5寸槽位上一个PCIe插槽功耗只有40W但推理性能却超过了同代旗舰GPU一个数量级。这是什么概念你原来机房里跑一组生产级神经网络推理需要8块GPU用TPU可能只需要1块而且总功率还低得多。对数据中心来说功耗和散热不只是电费问题还牵扯到机柜承重、冷量分配、机房面积这些全是真金白银。再一个就是确定性。GPU生态好、灵活性高但恰恰是这种灵活导致你在生产环境里很难榨干它的每一分性能。而TPU呢它的编程模型就两个核心指令一个负责矩阵运算一个负责主机通信。你没法用它干别的但它干这一件事的效率GPU确实比不了。2. 核心细节解析与实操要点2.1 TPU的顶层架构一块只为“计算流”服务的芯片TPU的整体架构可以从它的数据流路径来看。数据和指令从主机过来之后会先在片上做一个调度和分配。核心的计算引擎是一个极大的二维乘累加阵列这个阵列就是TPU的“心脏”。左右两边各有一个SRAM缓冲区分别负责喂数据和收集计算结果。底部还有一块激活函数单元专门处理非线性变换。整条链路像一条流水线一样输入一批数据输出一批结果全过程中没有中断、没有分支跳转带来的pipeline清空因为所有计算路径都是固定的硬件可以在一开始就把后面的指令预取好。我当初第一次看到TPU的die照片时最震撼的就是那个大矩阵几乎占满整个芯片面积。这不是为了铺满而铺满因为所有其他模块都在为它让路。做芯片的人都清楚面积就是成本Google愿意把这么大的面积投入单个功能块说明他们对“矩阵运算路径最短化”的决心有多大。2.2 乘累加阵列脉动阵列的数学原理与物理实现TPU最核心的技术点就是那个乘累加阵列业界也常叫它脉动阵列Systolic Array。这里不拽术语我用大白话解释它为什么高效。常规的CPU一次乘法算完再算加法数据要到内存里取取回来算完再存回去循环往复。这种模式叫做“取指-译码-执行-访存”每一步都有开销。而TPU的脉动阵列不是这个玩法。它把几十上百个乘法器和加法器拼接成一个网状结构数据从一个方向流入权重从另一个方向流入每经过一个节点就做一次“乘加”结果直接流向下一层。打个好理解的比方把这个阵列想象成一个生产流水线。工人站在流水线旁边零件从传送带上传过来每个工人只做一件事把自己手边的一个零件数取出来和传送带上的数相乘再加到自己的小本子上然后把本子传给下一个人。整条产线转起来之后第一个工人手边零件的计算结果会一路滚动到最后形成批量完成的效果。这个过程完美的利用了每个乘累加单元把内存访问次数降到最低。关键参数方面TPU v1的脉动阵列是256×256的大小也就是说它一个周期能完成65536次乘累加运算。以700MHz的主频来算理论峰值就是92 TOPS每秒92万亿次整数运算。这个数字在当时是非常恐怖的。2.3 片上存储的容量分配为什么这里的SRAM这么重要TPU里有两块大的片上SRAM一块存权重一块存中间激活值。搞过CUDA优化的同学都知道全局内存访问的延迟和带宽是性能瓶颈。TPU的方案就是“把数据尽量留在芯片上不让它出门”。举个例子当你在跑一个典型的卷积模型时权重矩阵可能有好几MB。如果每次计算都要从DDR内存里取权重那总线带宽就会成为脖子上的绳子。TPU把权重预先存在片上的统一缓冲区里计算时可以连续不断地向脉动阵列喂数据。论文里给出的数据是TPU v1片上SRAM容量是24MB实际分成了两块一个存权重一个存中间结果。这个容量对当时的模型尺寸来说已经可以做到把绝大多数模型的核心权重都塞进芯片内部。再一个容易被忽视的是累加器。在脉动阵列的底部TPU专门放了一块很大的累加器存储而不是直接把结果写回统一缓冲区。这是因为卷积计算经常要累加多轮结果频繁存取中间结果会消耗大量功耗。在片上直接用一个34-bit甚至更宽的累加器做连续累加只有当整个计算窗口收敛之后才写回缓冲区。这个细节非常有意思它反映出的设计思路是能不上内存总线就不上能多算一会儿就多算一会儿。2.4 控制单元为什么TPU几乎不需要分支预测传统CPU里最复杂也最耗电的模块之一就是分支预测和乱序执行引擎。CPU为了在处理逻辑判断和分支跳转时不让流水线空转花掉了大量的晶体管做预测。但TPU根本不在乎这些。因为神经网络推理的计算流程是高度确定和循环化的。一个卷积层的计算就是套在循环里的乘累加无非是输入维度的长短整个控制逻辑几乎可以被看作是一个超大号的for循环。所以TPU的控制单元非常简洁它只需要按照预先生成的指令序列一条一条按顺序执行就好不需要分支预测不需要乱序执行不需要寄存器重命名。这些被省下来的晶体管和功耗通通转移给了脉动阵列。对我们做工程的人而言这个设计选择背后的思维特别值得学习如果确定自己的任务99%的时间都是一种固定模式为什么还要花大量资源给剩下那1%的极端case做优化呢现实中很多性能问题就是因为想“太通用”而变慢的。3. 实操过程与核心环节实现3.1 数据如何从主机送入TPU为什么用DMA引擎在使用TPU的实际流程中数据通路是这么走的主机Host CPU先把TPU要执行的指令和权重写入TPU的输入内存里然后TPU通过DMA直接内存访问引擎把数据取走。你别小看这一步DMA的价值在于它不需要主机CPU去逐条搬运数据效率会高很多。层面到芯片内部DMA引擎会把权重数据从外部DRAM搬进权重SRAM再把中间激活值从统一缓冲区搬到脉动阵列的输入端。如果你做过嵌入式开发或者FPGA开发这一套流程会非常有共鸣它就是“数据搬移和计算重叠”的经典做法。我在这里分享一个实操心得TPU虽然是一个专用芯片但你在管理它的输入序列时一定要保持数据的按序性。因为TPU为了极致效率指令流是顺序执行的一旦出现一个需要同步等待的突发请求整个流水线都会停顿。这跟异步批处理的设计哲学是相通的你在调度的时候就应该尽量把一个小批量内的输入组织得整齐划一不要让单个慢样本拖累整个批次。3.2 TPU的指令集只有5条指令的超精简世界观同学们注意了这是TPU设计中最有意思的部分之一。它的指令集总共只有CISC风格但数目极少的核心指令。论文里概括起来大概是5大类Read从主存读数据到TPU片上缓冲区Write把TPU计算结果写回主存MatrixMultiply执行矩阵乘法这是最核心的指令周期最长也决定了整个TPU的利用率Activation执行激活函数和池化操作Jump辅助类的跳转指令其中MatrixMultiply指令非常夸张它自带了很多参数包括矩阵的行列数、步长、可选的激活函数类型、以及累加目标地址。一个周期执行这条指令时整个脉动阵列都在并行工作算力怪兽真正启动。这里要注意的一个坑是虽然指令少但指令的编码和参数组织方式非常关键。因为MatrixMultiply执行一次可能持续几百个周期你在写调度代码的时候必须清晰地预测到指令之间的依赖关系。TPU不会帮你自动规避依赖它的思路就是一条路走到黑出了问题只能靠编译器或者上层编程框架去避免。3.3 权重驻留优化片上命中率的黄金法则在跑模型推理的时候TPU有一个运行时的核心优化目标把权重尽可能长时间地驻留在片上的SRAM里。因为从外部存储搬运权重的功耗和延迟远大于在片上重复读取。实际操作中很多模型例如一个7层的CNN其权重大小在几MB量级。TPU v1的权重SRAM是16MB左右有些资料说统一缓冲区共24MB。这意味着你完全可以在启动推理之前先把所有层的权重全部加载到片上之后整个推理过程都不需要再从外部取权重。这种情况下的性能会非常漂亮。如果你的模型大到权重超过片上SRAM那就需要以层为单位进行分块调度这时你需要仔细规划每一层计算和下一层权重预加载的重叠否则就会因为等待数据而空转。这也是很多其他AI芯片设计中的经典话题叫“计算和通信重叠”你可以在任何一本讲GPU优化的书里看到类似思想TPU则把这一点芯片化了。3.4 数值精度的选择8位整数为什么够用TPU v1时代它的脉动阵列做的是8位整数乘法累加用32位。这在很多搞传统的HPC的人看来可能不太可思议因为CPU和GPU讲究的是FP32甚至FP64的浮点精度。但Google研究团队当时发现神经网络推理对权重和激活值精度并不敏感。推理阶段的参数已经是通过训练收敛出来的噪声容忍度相对高。8位整数就能提供足够好的结果这带来的直接收益是相比FP32的运算单元8位乘法器的面积和功耗都大幅降低同等的芯片面积可以塞更多计算单元。你如果查资料会发现后来的TPU各代产品开始支持BF16等格式那是为了兼顾训练和推理但在TPU v1时代“8位整数足够”就是Google交出的答案。这个决策背后也有教训可以总结数值精度不是越高越好而是够用最好。每提高一档精度吞吐、功耗、内存带宽的代价都是指数级上升的。你在优化自己的推理服务时如果框架允许完全可以试试INT8量化性能提升往往立竿见影。4. 常见问题与排查技巧实录4.1 为什么TPU那么快但我的模型没有提速这是一个非常常见的问题我在帮朋友分析TPU落地场景时也反复遇到过。很多人以为拿到TPU就能无脑升高性能结果是模型结构差异巨大最终的运行效率天差地别。在TPU上效率最高的模型是卷积层非常规整、通道数全是固定倍数的网络。如果模型里有大量不规则的分支、动态shape、全连接层占比特别大TPU的实际利用率会直线下降。因为全连接层本质上可以转化为矩阵乘法但动态shape会导致权重的排布和数据的组织不够规整脉冲阵列有时候会空转。排查思路你自己模拟TPU的内部运行模式把模型的每一层抽象成矩阵运算看看你的运算在一个256×256的网格上能否完整填充。如果很多层只有很小的矩阵格子里大量单元闲着那TPU的优势就发挥不出来。4.2 内存带宽是不是TPU的瓶颈很多做惯了GPU优化的人会直觉地认为TPU的瓶颈肯定也是带宽。其实在TPU v1的负载下片上SRAM命中率做到较高水平时外部DRAM带宽不是主要矛盾。TPU架构的主要瓶颈点有两个一个是同步操作等待另一个是批处理大小。如果主机每次只丢过来一个很小的批次TPU的长流水线效应会被频繁打断芯片能效就会下降。所以我在实际部署推理任务时宁可在主机端多攒一批请求再发给TPU也不要实时逐条送。TPU本质上是个“吃大饭”的架构你一次给它一顿满汉全席它当然爽你一次给它一粒米它就懒得动。4.3 如何判断TPU是否“吃饱”了一个很有效的指标是TPU的忙闲比也就是脉动阵列实际处于计算状态的周期占总周期的比例。简单点说就是看芯片利用率。你可以通过TPU性能计数器拿到这些数据这有点像看CPU的IPC每周期指令数或者GPU的SM占用率。一般来讲如果你的模型属于典型的卷积网络并且输入批次组织得当忙闲比能到80%以上。如果这个值低于50%你应该回头检查数据搬运、指令调度、权重驻留三个环节里有没有哪个卡住了。我把排查流程总结成了一张速查表遇到TPU性能问题时按这个顺序检查通常很快能定位问题检查项参考操作低效时的信号批次大小增大单次推理的batch size忙闲比低整体吞吐上不去权重驻留统计各层权重总计大小对比片上SRAM权重反复从外部加载功耗偏高动态shape将输入固定为最接近的长方形阵列利用率出现大量空洞指令依赖在调度时避免同一缓冲区的读写冲突周期性停顿导致延迟升高4.4 热功耗问题的另一种解法有工程经验的都会关注散热。TPU v1只有40W功耗放进2.5寸硬盘槽位都没问题。到了后面几代TPU比如v3、v4性能和功耗都水涨船高就得依赖专门的液冷机柜了。这里面有一个通用经验芯片的性能约束通常会在散热条件卡死之后显得尤为明显。你在自己的环境里做性能调优时也一样把频率和电压当作可调节参数如果你的机器散热条件好可以适当拉高运行频率来换取性能反过来如果机房散热已经到极限别硬撑调低频率可能会让整体稳定性更好。5. 软件栈与编译器的协同设计5.1 TPU的出现也逼出了专用编译器硬件再强没有合适的软件栈就是一块砖头。Google很早就明白如果要让开发者愿意用TPU就必须提供一套好用的编译工具链。当时主推的方式就是通过TensorFlow的计算图把网络结构自动转换为TPU可执行的指令序列。这套编译器的角色其实非常像传统编译器中的指令调度器但它更直观把计算图里的卷积层、全连接层、池化层识别出来然后一一映射到TPU的MatrixMultiply指令和Activation指令上。我在实际使用过程中感触比较深的一点是这种“图映射”模式看起来简单但中间的策略其实极其讲究。因为同一层卷积可能有N种切片方式每种方式对应的脉动阵列利用率和内存访问量都不同。编译器需要选择最优的layout这对编译器的优化深度提出了很高的要求。5.2 为什么说TPU和TensorFlow是命运共同体TPU有一个特点那就是它与TensorFlow深度绑定。这个特性在早期是一把双刃剑。好处是整条链路从框架到芯片Google可以全栈优化没有中间商赚差价坏处是如果你不用TensorFlow想驾驭TPU就非常困难。后来Google也意识到开放的重要性开始通过XLAAccelerated Linear Algebra编译器提供更通用一些的接口把任意前端模型编译到TPU可执行文件。这本质上就是编译器后端扩展的思路。我个人的体会是当你在评估一个专用AI加速芯片时千万别只看硬件指标一定要把它的编译器成熟度放在同等位置来评估。硬件再强编译不出高效的指令序列最终结果一定让你失望。5.3 从TPU编译器里可复用的软件设计思路TPU编译器的很多设计思想其实可以反向借鉴到普通的后端服务性能优化里。比如“算子融合”这个概念就是把多个连续的小算子合并成一整个大算子减少中间结果的写入和读取。神经网络推理里最常见的融合操作是ConvBNReLU。TPU编译器会把它们融合成一次计算这样中间结果根本不用写回内存直接留在片上寄存器/累加器里持续计算。你在做服务端性能优化时也能用同样的思路把多个小请求合并成一个大请求或者把多个步骤的数据导入过程合并成一次I/O往往能获得非常显著的延迟下降。6. 横向对比与未来走向6.1 TPU与GPU、FPGA、其他NPU的差异在哪很多刚接触AI硬件的人会混淆TPU、GPU、FPGA之间的区别。简单画个对比关系就清楚了GPU通用性较强的并行加速器适合训练也适合推理生态无敌但是功耗和专用效率稍弱。FPGA可重构逻辑灵活性高开发周期长适合特殊算子但量产成本较高。TPU/NPU专用的AI加速器性能功耗比极高灵活度最低适合大规模生产环境部署。CPU通用性最强但并行算力远逊于以上三类适合做调度而非密集计算。在实际数据中心的规划中我看到越来越多的架构是“CPUGPU/NPU”的混合方案CPU负责任务调度和预处理加速卡则专注张量计算。这跟TPU最开始的定位是完全一致的即便到了今天这个思路依然有效。6.2 TPU各代产品演进的几大趋势从TPU v1到v4几个非常明显的演进方向值得关注。计算精度从INT8扩展到BF16可以支持训练任务不再只是推理专用。互联能力不断加强从单机单卡走向多卡互联和超级计算机级别的组网。TPU v3起引入高速互联v4更是直接在光交换机上做数据交换。更大的片上SRAM、更高的主频、更复杂的内存层次结构。这些都在力图减少数据搬移的代价。编译器变得更加开放不只是TensorFlow专属逐步支持PyTorch等主流框架。如果说TPU v1只是一个“验证专用计算思路”的实验品那到后来的几代它已经成长为一套完整的数据中心级计算系统。你不能只把它看作一个芯片它其实是一整套包含芯片、编译、系统网络、散热运维协同运作的大型系统。6.3 AI芯片竞争的终局大家都殊途同归这几年市场上涌现了大量NPU/ASIC AI芯片你会发现大家最后都往同一个方向靠拢大矩阵乘累加阵列、片上大缓存、编译器与框架的深度绑定、多芯片互联。这说明TPU当年的架构选择经受住了时间检验。专用化、数据本地化、计算通信重叠这三个原则是任何AI计算芯片绕不开的核心话题。你把这三点想透了再回头去看市面上任何一款AI加速芯片的datasheet基本都能快速判断出它的优劣和适用场景。7. 写在最后一点很个人的心得做了这么多年性能优化和AI基础设施我觉得TPU带给我最大的启发不是“Google有多牛”而是一种思维方式的转变不要总想着一个东西什么都能干当你足够了解你的业务时你会知道最核心的那1%工作是什么然后为这一件事做到极致的专一优化其收益远大于在通用方案上修修补补。很多人在优化线上服务的性能时也会犯同样的毛病。今天调调缓存明天加加并发后天换换数据库却从没认真想过你这个系统99%的请求到底在执行什么操作。如果你能像Google设计TPU那样先把核心路径分析透然后集中资源把这条路径优化成一条流水线效果往往比零敲碎打强得多。最后再分享一个小技巧也是我在这几次AI芯片项目落地过程中反复用到的经验拿到任何一款加速硬件第一件事不要先跑基准测试看FPS先去看它的数据流路径。你能不能用你最常用的模型把整条数据通路走一遍哪些数据该放在哪里哪些步骤可以重叠等你把这些问题想通性能上的问题大概率都能自行化解。芯片是这样系统设计也是这样。