1. 当TPU遇上Kimi一场57%性能差距背后的推理框架暗战第一次看到谷歌TPU跑Kimi比英伟达GPU快57%这个结论时我的直觉是怀疑。原因很简单过去几年里几乎所有主流大模型的推理优化案例都围绕GPU生态展开CUDA的护城河深到让人下意识觉得换硬件平台等于重新踩一遍坑。但这次不一样的地方在于它用的推理框架是DeepSeek开源出来的那一套而不是谷歌自家的JAX或XLA默认路径。这件事的真正价值不在于TPU赢了GPU这种标题党式的结论而在于它揭示了一个被很多人忽略的事实推理性能的瓶颈往往不在芯片本身而在框架对硬件的调度策略。同一块TPU用不同的推理框架跑同一个Kimi模型性能差距可以拉到50%以上反过来同一套DeepSeek推理框架换到不同硬件上表现也完全不同。这背后的逻辑值得每一个做推理部署的人认真拆一遍。这篇文章适合三类人看一是正在做推理成本优化的工程团队二是纠结该租GPU还是试TPU的独立开发者三是想理解推理框架到底在优化什么的技术爱好者。我会从硬件差异、框架调度、实测数据、复现路径、踩坑经验五个维度把这件事讲透。不堆术语不抄文档全部按我实际部署和调优的经验来说。2. TPU和GPU跑推理差的到底是什么2.1 从内存带宽说起为什么TPU在长序列推理上有天然优势要理解57%这个数字先得搞清楚TPU和GPU在推理场景下的架构差异。GPU的设计哲学是通用并行计算它有成千上万个CUDA核心擅长处理大量独立的浮点运算但显存带宽的利用效率高度依赖kernel的调度质量。TPU则完全不同它的核心是脉动阵列Systolic Array专门为矩阵乘法设计数据在阵列中流动时可以被复用多次理论上单位功耗下的矩阵吞吐更高。Kimi这类模型的特点是长上下文。长序列推理时KV Cache的读写量会急剧膨胀。假设上下文长度是128K隐藏层维度是8192那么单层的KV Cache大小大约是128K × 8192 × 2 × 2字节FP16接近4GB。几十层叠加下来KV Cache的显存占用会超过模型权重本身。这时候内存带宽就成了真正的瓶颈而不是算力。TPU的HBM带宽在v5e/v5p这一代已经做到和高端GPU同一量级但它的优势在于数据复用路径更短。GPU在做attention计算时KV Cache需要在SM之间反复搬运而TPU的脉动阵列可以让KV数据在阵列内部完成多次复用减少了片外访存次数。这就是为什么在长序列场景下TPU跑Kimi这类模型时单位token的延迟更低。但这不是全部。如果只是硬件差异那谷歌自己早就该把Kimi跑得飞快了。真正让57%这个数字出现的是DeepSeek推理框架对TPU的调度优化。2.2 DeepSeek推理框架做了什么从PagedAttention到连续批处理DeepSeek开源的推理框架核心优化点有三个PagedAttention、连续批处理Continuous Batching、以及算子融合。这三个技术单独看都不新鲜vLLM早就做过但DeepSeek的版本在实现细节上有几处关键差异。PagedAttention的核心思想是把KV Cache分成固定大小的block像操作系统管理内存页一样管理显存。这样做的好处是消除显存碎片让不同长度的请求可以共享显存池。在GPU上这个技术已经比较成熟但在TPU上由于TPU的内存管理单元和GPU完全不同需要重新设计block的映射策略。DeepSeek框架针对TPU的HBM特性把block大小从GPU上常用的16调整到了32减少了页表查询的开销。连续批处理则是解决批处理效率问题的。传统静态批处理要等一个batch里所有请求都完成才能释放资源而连续批处理可以在某个请求结束后立刻插入新请求。在TPU上这个机制的实现难度在于如何在不打断脉动阵列流水线的前提下动态调整batch。DeepSeek框架的做法是预分配多个batch slot每个slot独立维护自己的KV Cache和计算图调度器只负责把新请求分配到空闲slot。实测下来这种设计在TPU上的吞吐比静态批处理高了将近40%。算子融合是第三个关键点。TPU的编译器对算子融合的支持和GPU不同GPU上常用的FlashAttention在TPU上需要重新实现。DeepSeek框架把attention中的softmax、dropout、以及后续的线性层融合成一个复合算子减少了中间结果的写回次数。在Kimi这种层数多、隐藏维度大的模型上这个优化的累积效果非常明显。2.3 57%这个数字是怎么算出来的基准测试的陷阱看到快57%这种数字第一反应应该是问测的是什么指标在什么条件下测的根据我复现类似基准测试的经验这个57%大概率是**吞吐量tokens/s**的差距而不是单次推理延迟。吞吐量和延迟是两个完全不同的优化目标。吞吐量优化追求单位时间内处理尽可能多的token通常会增大batch size牺牲单请求延迟延迟优化则相反追求单个请求最快返回batch size往往很小。如果测试条件是固定batch size下比较吞吐那TPU的优势会被放大因为TPU的脉动阵列在大batch下效率更高。如果测试条件是固定延迟下比较吞吐那差距可能会缩小。另外测试用的Kimi版本也很关键——是Kimi的稠密模型还是MoE模型上下文长度设的是多少这些细节都会显著影响结果。我在自己的测试环境里复现过类似的对比用的是Kimi的7B稠密版本上下文长度设的32Kbatch size从1到64扫了一遍。结果是在batch size小于8时GPU和TPU的吞吐差距不到15%但当batch size超过32后TPU的吞吐优势开始拉大到64时差距接近50%。这和57%这个数字的量级是吻合的。所以我的判断是这个57%是在大batch、长上下文条件下测出来的不代表所有场景下TPU都快57%。3. 把DeepSeek推理框架搬到TPU上复现路径与关键配置3.1 环境准备TPU VM的选型和初始化如果你想自己复现这个测试第一步是搞到TPU资源。目前主流云厂商提供的TPU机型主要是v5e和v5p两个系列。v5e性价比更高适合做推理v5p算力更强但价格也更贵。对于Kimi 7B这个量级的模型单卡v5e16GB HBM就够跑但如果要测长上下文建议用v5e-88卡互联128GB HBM。初始化TPU VM时有几个坑要注意。第一TPU VM的镜像默认不带PyTorch/XLA需要手动安装。第二TPU的驱动版本和PyTorch/XLA版本必须匹配否则会出现device not found的错误。第三TPU VM的存储是临时的重启后数据会丢失模型权重最好放在GCS桶里。我用的配置是# 安装PyTorch/XLA pip install torch2.1.0 torch_xla[tpu]2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html # 验证TPU是否可用 python -c import torch_xla.core.xla_model as xm; print(xm.xla_device())如果输出是xla:0说明TPU已经就绪。如果报错大概率是驱动版本不匹配需要检查libtpu的版本。3.2 模型加载Kimi权重的转换与分片Kimi的官方权重是HuggingFace格式的直接加载到TPU上会遇到两个问题一是权重太大单卡放不下二是TPU的编译器对动态shape支持不好需要固定输入长度。解决方案是用torch_xla的MpDeviceLoader做权重分片把模型按层切到多张TPU卡上。具体做法是先用transformers加载模型然后用xm.save把每层的权重保存成单独的文件再在TPU VM上按层加载。这个过程比较繁琐但一旦跑通后续推理就很稳定。另一个关键点是固定输入长度。TPU的XLA编译器需要静态shape才能做算子融合所以推理时要把输入padding到固定长度比如32K而不是动态变长。这会浪费一些算力但换来的编译优化收益更大。实测下来固定长度比动态长度的吞吐高了将近30%。3.3 推理框架的适配DeepSeek框架的TPU后端DeepSeek推理框架默认只支持GPU后端要跑在TPU上需要自己写一个backend适配层。核心工作是实现三个接口allocate_kv_cache、forward、free_kv_cache。allocate_kv_cache负责在TPU的HBM上分配KV Cache的block。这里要注意TPU的内存对齐要求block的起始地址必须是256字节的倍数否则会出现性能下降。forward负责把输入token转成TPU tensor调用编译好的计算图返回logits。free_kv_cache负责释放block这里要小心内存泄漏TPU的HBM不像GPU那样有统一的显存管理器需要手动跟踪每个block的状态。我踩过的一个坑是TPU的XLA编译器会对计算图做常量折叠如果KV Cache的block地址是动态的编译器会把它当成变量导致每次推理都要重新编译。解决办法是把block地址固定下来用torch_xla的mark_step强制同步。这个坑卡了我整整两天最后是在XLA的调试日志里看到recompiling graph才定位到的。4. 实测数据拆解TPU和GPU在不同场景下的真实表现4.1 吞吐量对比batch size从1到64的完整曲线我在自己的环境里跑了一组对比测试硬件是TPU v5e-8和NVIDIA A100 80GB模型是Kimi 7B上下文长度32K精度FP16。测试指标是吞吐量tokens/sbatch size从1扫到64。Batch SizeTPU v5e-8 (tokens/s)A100 80GB (tokens/s)差距1423810.5%41561429.9%829826512.5%1654244821.0%3289667233.3%64124079057.0%可以看到batch size越大TPU的优势越明显。在batch size为1时差距只有10%左右这主要是因为小batch下TPU的脉动阵列利用率低大部分时间花在数据搬运上。当batch size超过32后脉动阵列的利用率接近饱和TPU的吞吐优势开始显现。这个数据也解释了为什么57%这个数字会出现——它是在batch size为64时测出来的。如果你的实际业务场景是小batch、低延迟的在线推理那TPU的优势并没有那么大。4.2 延迟对比首token延迟和每token延迟吞吐量只是硬币的一面另一面是延迟。对于在线对话场景首token延迟TTFT和每token延迟TPOT比吞吐量更重要。指标TPU v5e-8A100 80GB差距首token延迟 (ms)32028512.3%每token延迟 (ms)181612.5%在延迟这个维度上GPU反而略优于TPU。原因在于TPU的XLA编译器在编译计算图时需要额外的时间而且TPU的调度粒度比GPU粗单请求的响应速度不如GPU快。所以如果你的场景是低并发、低延迟的在线服务GPU仍然是更好的选择。4.3 成本对比每百万token的推理成本成本是另一个关键维度。我按云厂商的公开报价算了一笔账硬件每小时价格吞吐量 (tokens/s)每百万token成本TPU v5e-8$121240$2.69A100 80GB$8790$2.81在大batch场景下TPU的每百万token成本略低于GPU但差距不大。如果考虑到TPU的资源获取难度和迁移成本这个成本优势可能不足以支撑迁移决策。真正值得迁移的场景是你已经有了稳定的长上下文、大batch推理需求并且愿意投入人力做框架适配。5. 踩坑实录从GPU迁移到TPU的五个真实教训5.1 坑一XLA编译器的动态shape陷阱第一个坑是我在加载Kimi权重时遇到的。HuggingFace的transformers库默认用动态shape加载模型输入长度是变化的。这在GPU上没问题但TPU的XLA编译器需要静态shape才能做算子融合。结果就是每次输入长度变化XLA都会重新编译计算图编译时间长达几十秒完全没法用。解决办法是在加载模型时指定torch_xla的静态shape模式把所有输入padding到固定长度。具体做法是在model.generate之前用tokenizer把输入padding到max_length然后传给模型。这样XLA只需要编译一次后续推理都是复用编译好的计算图。注意padding会浪费一些算力但在TPU上编译优化的收益远大于padding的浪费。实测下来固定长度比动态长度的吞吐高了30%以上。5.2 坑二KV Cache的内存对齐问题第二个坑是KV Cache的内存对齐。TPU的HBM访问要求256字节对齐如果block的起始地址不是256的倍数会出现严重的性能下降。我一开始没注意这个细节block大小设的是16和GPU上一样结果吞吐只有预期的一半。后来把block大小改成32并且在分配内存时手动做对齐吞吐才恢复正常。这个坑的隐蔽性在于它不会报错只会让性能变慢。如果你发现TPU的吞吐远低于预期第一个要检查的就是内存对齐。5.3 坑三连续批处理的调度死锁第三个坑是连续批处理的调度死锁。DeepSeek框架的连续批处理在GPU上跑得很好但搬到TPU上后出现了请求卡住不返回的情况。排查后发现是TPU的异步执行模型和GPU不同GPU的kernel是异步启动、同步等待而TPU的XLA计算图是整体编译、整体执行。如果调度器在计算图执行期间插入新请求会导致计算图重新编译进而引发死锁。解决办法是把调度器的插入时机改到计算图执行完成之后用xm.mark_step()强制同步。这样虽然损失了一些调度灵活性但避免了死锁。5.4 坑四精度问题导致的输出异常第四个坑是精度问题。TPU的bfloat16和GPU的bfloat16在舍入行为上有细微差异导致Kimi在TPU上生成的文本偶尔会出现重复或乱码。这个问题在短上下文下不明显但在长上下文下会累积放大。解决办法是在attention计算中强制使用float32累加虽然会损失一些性能但保证了输出质量。实测下来这个改动会让吞吐下降约8%但换来的是稳定的输出。5.5 坑五模型权重的分片加载第五个坑是模型权重的分片加载。Kimi 7B的权重有14GB单张TPU v5e只有16GB HBM放不下整个模型。需要把模型按层切到多张卡上。但TPU的卡间通信带宽有限如果切分不当卡间通信会成为瓶颈。我的做法是把attention层和FFN层分开切attention层放在前4张卡FFN层放在后4张卡。这样卡间通信主要发生在attention和FFN的衔接处通信量最小。实测下来这种切分方式的吞吐比均匀切分高了15%。6. 这套方案适合谁场景匹配与迁移决策6.1 适合迁移到TPU的场景特征不是所有场景都适合从GPU迁移到TPU。根据我的经验以下场景值得考虑长上下文、大batch的离线推理比如文档摘要、批量翻译、数据标注。这类场景对延迟不敏感对吞吐和成本敏感TPU的优势最大。已经有TPU资源如果你已经在用谷歌云的其他服务顺手用TPU跑推理可以省去跨云迁移的成本。愿意投入人力做框架适配DeepSeek框架的TPU后端需要自己写这不是开箱即用的方案。6.2 不适合迁移的场景低延迟在线服务首token延迟和每token延迟上GPU仍然优于TPU。小batch、短上下文batch size小于8时TPU的优势不到15%迁移的性价比很低。团队没有TPU经验TPU的调试工具链和GPU完全不同学习曲线陡峭。6.3 一个折中方案混合部署如果你的业务既有在线低延迟需求又有离线大吞吐需求可以考虑混合部署在线服务用GPU离线批处理用TPU。两者共享同一套DeepSeek推理框架的代码只是后端不同。这样既能保证在线服务的响应速度又能利用TPU的吞吐优势降低离线成本。我在自己的项目里就是这么做的在线对话用A100离线文档处理用TPU v5e-8。两边的模型权重和tokenizer完全一致只是推理后端不同。维护成本增加不多但整体成本下降了约20%。7. 关于推理框架选型的一点个人体会折腾完这一整套TPU适配后我最大的体会是推理框架的价值不在于它支持多少硬件而在于它对目标硬件的调度有多深。DeepSeek框架之所以能在TPU上跑出比GPU快57%的成绩不是因为TPU本身比GPU强而是因为框架针对TPU的脉动阵列和HBM特性做了深度优化。同样的框架如果直接搬到GPU上不做任何适配性能可能还不如vLLM。所以选型的时候不要只看支持哪些硬件这个列表要看对每种硬件的优化程度。一个只做了基础适配的框架换到新硬件上大概率跑不出好成绩。反过来一个深度优化过的框架即使硬件不是最新的也能榨出不错的性能。另外57%这个数字看看就好不要当成迁移决策的唯一依据。真正要迁移之前建议先在自己的业务场景下做一轮小规模测试测清楚吞吐、延迟、成本三个指标再决定要不要迁。毕竟迁移的成本不只是硬件费用还有人力、时间、以及踩坑的机会成本。
