TPU推理外部化全面加速：InferenceX首份第三方性能数据发布

十多年来，业界一直见证谷歌在自有硅片上建立帝国。搜索、广告、YouTube以及每一代Gemini都运行在TPU上。很少有加速器能像它一样受到架构上的审视，也引发了如此多关于其性能和经济性在设计公司之外的讨论。

Anthropic是TPU的最大用户，到2029年超过Deepmind自身的使用量。

谷歌内部的成功从来不是问题。问题是行业其他部门能获得多少优势。你能否利用一个开放权重模型，通过熟悉的推理引擎来提供，并在经济学上击败NVIDIA？

今天，我们将发布TPUv7 Ironwood的首批第三方推断结果InferenceX官方预览。在我们与B200/B300的公平比较中，Ironwood每美元性能提升了高达50%。

其优势涵盖了帕累托曲线的大部分区域，我们从两方面分析了经济学：谷歌内部的总拥有成本和实际客户支付的外部总拥有成本。

Ironwood（TPUv7）是谷歌第一代用芯片竞争他人推理工作负载的版本，芯片可直接购买或通过自家云租赁。Anthropic非常喜欢TPU，承诺开发超过一百万个TPU（直接购买约40万+，GCP租赁约60万+），主要用于培训，也用于推理。

我们的加速器模型包含了Anthropic最新的TPU出货量和谷歌按季度的整体TPU出货量，以及TPUv8i、v8t和各种v9/v10等的估算

我们对新的 TorchTPU 栈——TPU 外部协议栈——开发速度感到兴奋。文章后面，我们将讨论 TPU 软件外部化所需的未来工作，包括优化推测性解码、拆分预填充、KV 缓存卸载、多回合代理工作负载等。

即便如此，我们 SemiAnalysis坚信TPU外部化正朝着正确方向发展，并全力推进。此外，与仍在学习如何构建以测试为先的软件文化不同，谷歌拥有数十年的软件工程经验和极为成熟的质量驱动文化，因此我们预计外部TPU软件将迅速成熟。

本文将介绍TPU内核和开权模型服务栈的所有优化，包括DP注意力优化、MoE路由和内核优化，以及减少GDN内核填充。我们还将深入探讨TPU系统，讨论出色的TPU性能工程师们正在采取的下一步步骤，以使该堆栈更广泛可用。

谷歌已经花了十多年时间展示其用TPU构建的能力。现在我们可以衡量行业其他部门对它们的能力。

特别感谢谷歌团队（Chris Chan、Jahangir Hasan、张王元、Anne Stern、Puneith Kaul、Ruizi Dong、Sangam Jindal、Qi Zhou、Madhan Jaganathan、Gang Ji、Jun Wan、Devanshu Jain、Jiaxin Cao、Srinath Mandalapu、Haowen Ning）和Inferact团队，感谢他们为这场精彩的TPU基础和表现贡献良多！

此外，还要感谢RadixArk团队，他们也在参与TorchTPU SGLang的开发。

InferenceX官方预览：TPUv7 Ironwood对阵Blackwell及Blackwell Ultra

我们已经看到即将推出的原生 TorchTPU vLLM 协议栈在与 Nvidia GPU 的全面比较中取得了强劲成果。谷歌在 FP8 中使用 Qwen3.5 397B 作为初始启动模型。

后续章节将深入探讨为何新的 TorchTPU 方法相比之前的 TorchAX 路径在外部 TPU vLLM/SGLang 服务方面有了显著改进。

一旦奠定了基础，谷歌计划将支持扩展到其他开放权重模型，包括Kimi K3和GLM5.3。一旦少数模型优化良好，我们相信在0日附近为各种热门开放模型添加优化支持将变得更加容易。

目前，vLLM和SGLang将对0日支持集中在英伟达，AMD的0日支持还算过得去。我们预计TorchTPU在不久的将来将足够稳定，vLLM和SGLang维护者可能会将TPU加入该0日名单。

该技术栈预计将在10月左右结束私测并开源。

在与FP8及单代币预测的聚合服务进行苹果对比中，TPU每美元性能比运行FP8的B200和B300高达50%。在使用FP4的模型上运行NVIDIA GPU时，存在质量损失。

TPUv7没有FP4原生计算，因此在FP4上，NVIDIA GPU仍保持领先。TPUv8i支持FP4，情况将改变，我们坚信TPUv8i Boardfly将能与Rubin NVL72竞争。

我们相信，一旦谷歌启用拆分服务和推测解码，且在比较双方都启用MTP和拆分服务后，TPU每美元的性能优势将持续存在。文章后面，我们将讨论谷歌计划在TorchTPU基础建立后添加的优化。

为了控制启动范围，谷歌最初专注于8k1k的基准测试。不过，我们相信TorchTPU协议栈在代理工作负载上也会表现出色，并将在今年晚些时候发布相关结果。

文章后面，我们将讨论TPU外部推断的下一步步骤和路线图。

TPU在每美元性能方面称霸

在每用户100个代币的交互性中，Ironwood每百万个代币的成本约为0.181美元，而B200为0.222美元，B300为0.276美元。这比B200低约19%，比B300低34%，同时实现了相同的用户生成速度。

在这次比较中，我们用外部TPU的TCO对比B200/B300的TCO，作为一个超大规模实验室的TPU购买。去年，SemiAnalysis Accelerator Model 率先揭露了谷歌正在销售 TPU 而非通过 Google Cloud 出租的案例.TPU较低的小时成本抵消了其在大部分交互性范围内吞吐量较低的缺陷。

半分析TCO模型提供了TPU的BoM估计和TCO估计的完整分解。

在大部分原始性能曲线中，TPU并未优于英伟达GPU。然而，这一比较并未考虑到TPU较低的TCO。TPU的买家主要关注的是他们能获得多少收入，即每美元的性能和每瓦的性能。

话虽如此，在20 tok/s/用户的交互性上，Ironwood在原始吞吐量方面也领先，达到9,364个总token/s/芯片，而B200为8,903个，B300为8,925个。

这比这两批GPU的吞吐量高出约5%。结合较低的建模小时成本，Ironwood每美元提供的代币数量比B200多50.4%，比B300多96.0%。

如果以 Google 内部工作负载的 TPU 总有成本为每芯片小时 1.03 美元，那么并发 256 时的每美元性能优势提升至 B200 的 76.7%，比 B300 的 130.2%。

然而，这些高并发结果伴随着延迟权衡。例如，在并发 256 时，TPU 平均 TTFT 为 5.41 秒，而 B200 为 3.75 秒，B300 为 2.40 秒。

前述 50% 到 96% 的优势针对该数据点，而非所有延迟目标。

从端到端延迟来看，Ironwood 在最高吞吐量点之后依然具有竞争力。在 20 秒的中位响应时间中，帕累托曲线显示 TPU 约为每百万个总代币 0.098 美元，而 B200 为 0.106 美元，B300 为 0.132 美元。

因此，TPU 的成本比 B200 低约 8%，比 B300 低 25%。

尽管如此，B200在苹果对苹果曲线的一小部分仍可能领先，包括大约30秒的中位响应时间。然而，TPU在更长的响应时间下，恢复了每百万代币成本较低的地位，同时在重叠区间内仍保持对B300的成本优势。

SemiAnalysis 是一份由读者支持的出版物。如需接收新文章并支持我们的工作，请考虑成为订阅者。

从苹果到香蕉：GB300 NVL72 Disagg verus TPUv7 Agg

谷歌多年来一直在生产环境中运行拆分服务，这一路径经过高度优化，但外部TPU服务栈尚未有完全优化的分区路径。因此，目前GB200/300 NVL72在分区对分区的表现上更具竞争力。

但我们预计这个差距将在几个月内缩小，TPUv7将与Google、Inferact、RadixArk和SemiAnalysis的土地优化竞争，我们将在后续文章中发布TPUv7与GB200/300 NVL72的分区对比。

TPUv7 Pods通过低延迟ICI结构可扩展至超过1k+芯片，从而实现NVIDIA NVL72无法实现的大型分组和超宽EP优化。

即使是在TPUv7聚合服务verus GB300 NVL72分组的苹果对香蕉比较中，使用内部TCO，TPUv7在高延迟和低到中等延迟方面也有相当的竞争力。

但在右侧中等端对端延迟上，GB300 NVL72分组相比TPuv7聚合服务有一个优势。

从中来看，它在低端和高端对端端延迟下都具有竞争力，但在中间阶段，GB300在每美元上相较TPUv7聚合约有30%的优势。我们坚信，一旦TPUv7上的细分服务得到优化，它将在全帕累托上与GB300 NVL72竞争。

推理服务于新原生的一流 TorchTPU 后端与之前的 TorchAX 协议栈

之前的 TorchAX vLLM 后端

谷歌此前曾使用 TorchAX 后端将基于 PyTorch 的 vLLM 和 SGLang 实现翻译到 JAX。TorchAX 旨在让开发者通过 JAX 运行 PyTorch 模型定义。

然而，这种方法存在的问题促使谷歌与 PyTorch、vLLM 和 SGLang 社区合作开发一种名为 TorchTPU 的新方法。它允许开发者将 TPU 作为原生 PyTorch 设备使用，而 StableHLO、XLA 和 TPU 优化的 Pallas 内核则在幕后处理底层执行。

在文章后续章节中，我们将展示这款新 TorchTPU 后端与 NVIDIA Blackwell 及 Blackwell Ultra GPU 的性能比较。

vLLM TPU 支持经历了三个不同阶段的发展：

其首个TPU原型机使用PyTorch/XLA.其懒惰的执行模型将操作收集成计算图供 XLA 编译，而不是立即执行每个 PyTorch 操作。
目前的公共后端，TPU 推断当有 TPU 优化的 JAX 实现时，使用 TPU 优化的实现。否则，TorchAX 会将原始 PyTorch 模型转换为 JAX 可执行的操作。目标是重用成熟的 TPU 原语和并行支持，而无需重写每个 PyTorch 模型。
即将推出的第三个设计 TorchTPU，使 vLLM 能够将 TPU 视为原生 PyTorch 设备。它将成为 SGLang 和 vLLM 在 TPU 上的首选后端。

下图放大了第二阶段，展示了vLLM的tpu推理后端如何结合了直接JAX模型和通过TorchAX运行的PyTorch模型。

最初选择JAX是务实的。在最后公告谷歌描述了在低层优化、分页注意力以及将vLLM的工作者模型与TPU执行匹配方面遇到的困难。他们认为JAX会提供更成熟的TPU原语和并行支持。

TorchAX让团队无需重写每个PyTorch模型即可使用这些能力，而面向TPU的JAX实现则可以共享相同的内核和编译器工作。

在 TorchAX 中，开发者仍然在 PyTorch 中编写模型，但 JAX 在 TPU 上执行。TorchAX 作为两个框架之间的转换层。作为TorchAX 文档解释，Torchax.tensor。

张量的行为类似于普通的PyTorch张量，但背后有jax。数组。当模型调用张量操作（在PyTorch中称为ATen操作）时，TorchAX通过__torch_dispatch__钩截取并将其转换为一个或多个JAX操作。

因此，即使JAX执行计算，模型仍可保留在PyTorch中编写。

vLLM 还必须管理生成过程中变化的状态，尤其是 KV 缓存，它存储了早期令牌的注意力信息以便重用。其模型包装器将模型权重和KV缓存作为显式状态呈现给JAX。

这一过程称为函数化，将原本在Python对象内发生的变化转化为编译器可追踪的输入和输出。jax.jit随后可以将每个推理步骤捕获为计算图，并通过JAX的编译流水线进行重复执行。

从那里开始，每个组件都有不同的角色。Pallas 为性能关键操作提供了优化的 TPU 内核。StableHLO 以编译器可处理的标准化格式表示程序，XLA 将其转换为 TPU 的可执行代码。

在 vLLM 的标准 PyTorch 路径中，torch.compile 负责图形捕获和编译的协调。在 TorchAX 路径中，JAX/XLA 流水线已经执行这些任务，因此 vLLM 绕过了其通常的 torch.compile 路径。

这些转换层导致了许多问题，TorchTPU vLLM 栈旨在解决这些问题。

之前的 SGLang JAX 后端

SGLang 当前的公共 TPU 协议栈用途SGLang-JAX.SGLang-JAX 与 vLLM 的 TorchAX 后端做出类似权衡，但采用不同的方法。它是一个独立的 JAX 原生服务引擎，而非通过 TorchAX 转换的 PyTorch SGLang 运行时。

它结合了 SGLang 风格的调度和前缀缓存，配合 JAX 模型和 TPU 专用内核。Google 和 RadixArk 宣布 SGL-torchtpu 作为额外的 PyTorch 原生后端，以扩展可用的服务路径。

转向新的原生、一流的 TorchTPU 后端

火炬TPU将框架边界移动到PyTorch本身。它使用 PyTorch 的PrivateUse1 后端扩展点以暴露普通火炬。设备上的张量=“tpu”，而非由JAX数组支持的包装器。

PyTorch调度器将ATen操作路由到TPU后端。开发者可以急切地运行代码进行调用和调试，或调用torch.compile进行图捕获。在编译路径中，TorchDynamo和AOTAutograd生成FX图，TorchTPU将其降为StableHLO，XLA生成TPU可执行文件。

谷歌的会议甲板使编译器选择明确该路径使用XLA，而非电感和特里同。

因此，“Native”具有特定的范围。它涵盖了 PyTorch 张量、调度、急切执行、编译入口和分布式 API。XLA 仍然是编译器，内核层继续使用 TPU 专用的实现。

TorchTPU 可以调用 Pallas 和 JAX 支持的自定义内核。Helion 的 TPU 后端也输出 Pallas。面向用户的框架成为 PyTorch 原生框架，而编译器和性能关键的内核则保持 TPU 支持。

开发者可以使用 .to（“tpu”），同时保留熟悉的分布式接口和服务引擎代码。Google 文档支持 DDP、FSDP2、DTensor、每个设备一个进程，以及 MPMD 和 SPMD 执行。

这些接口符合 PyTorch 服务引擎已假设的进程和分布式模型。

对于 vLLM 和 SGLang，机会是重用更多上游模型代码、调度器、连续批处理、API 和功能逻辑，而不是跨越 PyTorch 到 JAX 边界重建它们。

这应该能降低调动模型和引擎特性的成本，即使最终 TPU 实现仍需专用内核。

原生 PyTorch 支持并不消除对 TPU 特定优化的需求。Pallas 和 TorchAX 有不同的用途，TorchAX 允许 PyTorch 模型代码通过 JAX 执行，而 Pallas 则用于直接为 TPU 硬件实现性能关键操作。

移除 TorchAX 并不意味着移除其底层的 Pallas 内核。由于 TorchTPU 允许原生 Pytorch 代码调用 Pallas 和 JAX 内核，早期栈开发的内核能够迁移到新栈。

类似于GPU PyTorch vLLM，高性能推理仍需定制内核。张量形状和布局仍需调优，以高效利用TPU的矩阵乘法单元。我们认为大多数现有Pallas内核可以轻松迁移到新的原生TorchTPU服务栈，但这并不意味着所有现有内核都会保持不变。

许多内容如包装器、张量布局和运行时集成等仍需稍作调整和验证。

与谷歌TPU团队的强力生态系统合作

vLLM 和 SGLang 是两大主要的开源生产推理服务栈。Inferact、RadixArk 和 Red Hat 正在与 Google 大力合作，使 TorchTPU 后端在两种栈中都成为一流的体验。

这种社区支持对于外部化 TPU 推理极为重要。

截至九月初，TorchTPU 仍处于私测阶段，预计将在十月中旬 PyTorch 大会期间开源。我们很高兴看到它推动 TPU 在 PyTorch、vLLM 和 SGLang 中获得一流的体验。

一旦 TorchTPU vLLM 和 TorchTPU SGLang 开源，我们将将 InferenceX/AgentX 的 TPU 基准测试从我们的分支迁移到我们的公开仓库。

TPU推理优化深度解析

为本地TorchTPU服务栈的首个启动模型优化TPU推理的冲刺耗费数百小时的工程时间和多次PR。如文中后文所述，TPU架构不同于GPU架构，因此需要自身的优化以高效运行模型。

下面，我们将介绍那些提升（或在某些情况下启用）Ironwood上Qwen3.5-397B及其他模型性能的优化。大多数Pallas内核优化可迁移到新的原生TorchTPU栈中。

DP 注意力优化

一个不一定是TPU特有的问题是如何有效地将模型架构映射到硬件的并行性。例如，Qwen3.5的GQA注意力层有32个查询头，但只有2个共享的KV头。

使用TP8（八个逻辑设备），查询计算可以均匀分布在8个设备之间，每个设备有4个查询头。然而，KV头在这些设备间分配不均匀。每个KV头由16个查询头共享，这意味着四个设备需要相同的KV数据来计算它们的本地注意力。

当TP超过KV头数时，标准做法是跨TP等级复制KV头。vLLM已经支持这一点。TPU后端需要兼容性修复激活这种现有行为，避免不必要的全对全沟通。

为了实现更高并发的服务，TPU后端还增加了合并支持八路注意力数据并行处理（DP8）采用八路专家并行（EP8），简称为“DP注意”或DEP8。

每个请求的注意力不分布在全部八个设备之间，而是每个设备处理不同的请求子集，并将这些请求的两个KV头保持在本地。注意力权重会被复制，但KV缓存保留的是不同的请求历史，而不是同一历史的重复副本。

与此同时，512个路由专家分布在各设备中，避免了对更大专家权重池的复制。实现这一工作是必要的协调请求分配、递归状态槽和块表在发球机上。

优化沟通

借助DP关注和EP，TPU的GroupedGEMM实现在专家运行前收集令牌激活和路由元数据，然后利用减少散布法将加权输出加总，并将每个令牌的结果返回其注意力排名。

后端最初分别收集选定的专家ID和路由权重。谷歌将这些组合成一个全集，去除了一个额外的集合，其延迟可能主导此类小阵列的传输时间。

PR报告每层大约80微秒保存在 DeepSeek-V3 测量中。虽然 80 微秒听起来很小，但 DeepSeek-V3/R1 有 58 个层需要 AllGather。在这些层之间，单独节省的每前向传递约 4.64 毫秒（80 微秒× 58 毫秒），减少了 TPOT 并提高了交互性。

谷歌也一样在 SparseCore 上实现了 ReduceScatter 集体（更适合数据移动等不规则操作）并利用 Ironwood 更快的芯片对芯片链路，在各芯片间交换部分求和前合并各芯片内的贡献。

这种每芯片两设备布局和芯片间 ICI 网络将在下一节介绍。双缓冲允许内核传输一个块，同时累积另一个块，从而与较慢的芯片间流量重叠局部减少和芯片间传输。

在 SparseCore 上运行集体还释放了 TensorCore 执行时间用于其他操作。

在上图中，通信阶段重叠，缩短了端到端的执行时间。例如，考虑标记为t_1的点。此时，芯片内DMA散点缩减（P1）已完成，针对MB（微批次）1。

ICI的散点缩减现在可以在IMI dim 0（P2.0）上进行同时进行的P1 用于后续的 MB。在基线下，这种优化结果为8k1k时吞吐量提升4.1%到14.2%在64至512并发中，其中256并发处占8.5%，并发512时1k8k占26.1%。

感谢阅读，SemiAnalysis！这篇帖子是公开的，欢迎大家分享。

该集合从 TensorCore 迁移到 SparseCore 也意味着在集体运行时可以对 TensorCore 进行额外的工作流水线。然而，并非每个集合都值得迁移。

在某些情况下，将一个集合卸载到 SparseCores 实际上加重性能。例如，Qwen3.5 a阈值现在决定何时将所有约简和全采集操作卸载到稀疏核,当小集体能在 VMEM 和 SparseCore 中时，将它们保留在 TensorCore 上，会让它们变慢.默认阈值基于 VMEM 容量。

PR 报告称，在并发 64 时吞吐量提升了 2.7%，并发 128 时提升了 5.7%。

优化 MoE 路由与专家内核

专家混合层为每个专家生成不规则的令牌组，这些杂乱的组必须被重塑为TPU矩阵单元能够高效消耗的组。以下的一些变更是对路由和分组matmul后端的总体改进，惠及任何MoE模型。

其他改动则直接在Qwen3.5上测量。

SemiAnalysis 是一份由读者支持的出版物。如需接收新文章并支持我们的工作，请考虑成为免费或付费订阅者。

该分组Matmul的第二种版本重新设计专家输入输入 MXU 的方式。它消除了冗余的瓦片计算，将 size 转移为有效行数而非填充的最大值，三重缓冲区将专家权重处理，使下一组的权重在当前组计算时已在运行中，并将组元数据生成整合进内核。

专家输入的不规则重组随后进入了SparseCore.SparseCore 负责数据流动，将每位专家的标记收集成连续的组，而 TensorCore 则负责执行专家矩阵乘法。

A后续重写改进了内存读取流水线，并将合并工作分成令牌和隐藏维度。PR报告显示，与原始SparseCore内核相比，8k1k的服务吞吐量提升了12%，同时TTFT和TPOT也降低了。

进一步优化将顶K权重收集器从粗糙采集-减少路径移动到SparseCore.这里的 gather 指的是从一个设备上读取选定的条目，而不是设备间的全收集。

此前，TensorCore 在预处理时收集路由权重和源索引，限制了 TensorCore 与 SparseCore 的重叠。将这些采集迁移到 SparseCore 内核后，TensorCore 开销从 29 微秒降至 14 微秒，在 DeepSeek-V3 微基准测试中将 146 微秒降至 137 微秒，采用 16 路专家并行。

对于小批量，一般的崎岖路径成本高于其安排的工作量。A专用小批量排列因此，构建单热矩阵，并使用普通矩阵乘法将符号置换给其专家，并对结果进行逆置换。

在8k1k工作负载下，这在并发64时提高了7.3%，在并发128时提高了5.1%。

一次进行中的变更将专家ID和令牌索引打包到一个排序键中，因此XLA在保持所需排序的同时执行更简单的排序。排序延迟从106.6微秒降至21.7微秒。

PR报告8k1k的服务提升为0.6%至8.5%，但同样的变更也支持FP8全收集。

优化门控DeltaNet Pallas内核

门控DeltaNet（GDN）是一种循环计算。也就是说，每一步运行状态都会衰减，更新一个秩一项，并投影输出。请查看以下文章，深入了解线性注意力机制如GDN：

以下变化展示了其矩阵操作、向量更新和状态传输在MXU、VPU、VMEM和HBM上的调度方式。Qwen3.5 的初始支持增加了纯JAX因果Conv1D和GDN实现，将其连接到TPU运算符调度，并支持递归状态缓存。

后续PR对这些实现进行了优化。

另一个容易实现的优化是：在输出投影计算中重新排列代数以使MXU和VPU工作重叠.此前，VPU首先对衰减状态应用第一秩更新，然后MXU将更新状态乘以查询以生成输出。

这种依赖性迫使MXU等待。

将衰减状态写为S，校正向量写作Δ，键和查询写作k和q，输出可以展开：

MXU现在可以直接从衰减状态计算Sq，而VPU则构建下一标记的更新状态。当前更新对输出的贡献分别通过kTq计算，kTq是每个点的标量点积，随后进行小规模向量加法。

这使状态更新从MXU的依赖路径中移除。报告的8k1k吞吐量提升在并发64时为2.79%，在并发512时为4.48%。

下一次变动减少向量寄存器溢出通过在解码循环中切片 Q 和 K，使同时保持活跃的值减少。decode-64 内核速度约提升 20%，但端到端收益较小，8k1k 时为 0.8%，并发 512 下为 3.8%，因为内核只是解码步骤的一部分。

异步状态传输DMA与双缓冲计算重叠。第二组缓冲区消耗的额外VMEM最初使数据并行注意力回归。重复利用现有暂存缓冲区并缩短临时缓冲区寿命可恢复该容量并消除回归。

报告的8k1k吞吐量提升为11.3%，并发512。

GDN v3将Conv1D和GDN融合到一个内核中，以减少HBM往返，改进预填充布局，并将混合预填充/解码执行统一为一条路径。报告的内核级加速为：解码时1.41×预填充时1.60×，混合批次时提升2.14×。

这些测量仅涵盖内核，无法确定端到端的服务提升。

仅仅通过改善内核内部操作之间的重叠，就能实现显著的性能提升！

管理混合状态和分页注意力

Qwen3.5 是一个混合模型，包含两种状态类型：GQA 层积累的 KV 历史，随着每个令牌的增加而增长，而 GDN 层则每个请求保持固定大小的重复状态。

状态分配、存储精度、注意力页大小和物理数据布局都决定了 HBM 的实际可用量以及注意力运行的效率。

SemiAnalysis 是一份由读者支持的出版物。如需接收新文章并支持我们的工作，请考虑成为订阅者。

批量、杂乱、分页、关注批量处理序列，预计算页面元数据，并设置三重缓冲区以改善流水线并减少填充。这一变化为共享注意力后端工作提供了基础。

PR的工作负载示例使用Qwen3-32B，因此其测量应与Qwen3.5-397B结果分开。

复发状态现在是紧致分配，每个激活请求大约只有一个槽位，而不是num_blocks每个层组的槽位。在报告的配置中，这回收了约76 GiB的HBM，并将注意力块池扩展了71%，在并发64时提升了18k8k的输出吞吐量。

在BF16中存储递归状态将HBM占用减半，同时将算术在VMEM中的FP32中保留。这节省了内存容量和传输带宽，同时保留了FP32的运算。在并发512时，报告的1k8k吞吐量提升为15%。

移除旧的混合页面大小对齐约束让批量关注使用合适的二次幂页面大小，包括256个令牌。1k8k吞吐量在并发512时提升约7%。这适用于非前缀缓存路径，与后来引入的对齐检查点模式不同。

不必要的KV布局复制品包括通过根据KV磁头数选择重塑路径来避免，而不是施加只有其他形状才需要的布局约束。Qwen3.5 8k1k吞吐量在并发512时提升约4.1%。

减少填充和低并发开销

当只有四到八个请求在运行时，针对数百个并发请求进行配置的配置会浪费工作。编译后的形状桶过大，元数据被限制在配置的最大值，填充令牌会触发虚拟专家工作。

InferenceX的操作点正好处于这些低并发度，这促使了以下变化。

请求元数据现在是按活跃请求数量划分而不是总是使用配置的最大值。在8k1k并发-64测试中，GDN调度开销从283微秒降至97微秒，吞吐量从2,328个令牌/芯片/秒提升至2,516个。

InferenceX 的显式 Qwen3.5 调优减少旋转表，限制每个DP等级的序列大小，并为并发四增加专属注意力桶。并发四的综合收益为8千千人时13.3%，1千千人时15.5%。

A进一步的低并发调优切换为TP8注意力加专家并行，减少最小令牌桶，并将填充令牌路由至专家零，避免触发额外的专家权重负载。1k1k的综合收益为四次并发时22.9%，八次时18.1%。

在8k1k下，四次并发时吞吐量增长9.2%，并发八时下降5.3%。

支持混合模型的前缀缓存

上述收益是在随机输入基准测试中测量的，即请求之间没有共享。前缀缓存对于代理型和多回合工作负载至关重要，这些工作负载重用了长的系统提示和对话历史。

对于混合模型，缓存前缀必须保留其KV块和前缀末尾的GDN重复状态。该重复状态通常在请求继续时立即被覆盖，使得前缀缓存比“常规”注意力模型更困难。

带有DP支持的混合前缀缓存通过为 GDN 分别设置读取检查点和写入活状态的独立槽位来解决这个问题。因此，继续请求不再覆盖与其缓存前缀相关的检查点。

状态地址由块表派生，该结构与已定位 KV 块相同。检查点以对齐的缓存粒度获取，因此保存的状态始终与 KV 块边界对齐。该模式需要完整的检查点池，而非之前描述的紧凑的每个请求分配，后者会以部分 HBM 交换请求间的重用。

当前缀实际共享时，好处显现出来。

分页注意力中的通道布局和流水线深度

两个硬件事实决定了 KV 缓存在 TPU 上的布局方式。我们将在下一节中更详细地讨论 TPU 硬件。首先，向量单元适用于最后一个维度为 128 个通道宽的图块，因此任何尾随维度小于该维度的数组都会被填充。其次，Pallas 内核通过双缓冲隐藏 HBM 延迟。它在计算当前块的同时获取下一个块。两个块都必须适合 VMEM，因此计算块的大小决定了预取的深度。这两个限制都会影响容量和吞吐量。

在批量注意力内核中，KV 缓存沿着头维度打包键和值，对于 FP8，打包因子为 4。每个设备只有一个 KV 头的模型只需要打包两件东西，因此每块瓷砖的一半都浪费在填充上。 A 车道上的顺序布局通过将页面的标记放在 128 通道轴上并将头部尺寸放在子通道轴上来解决此问题。可用的 KV 页增加了一倍（在报告的配置中从 5,141 增加到 10,283），并且头部尺寸只需要是 32 的倍数而不是 128，这使得头部尺寸为 64 的模型也可以使用内核。在低并发情况下，该布局的每个令牌延迟成本约为 3%，但在 8k1k 上并发 128 时，额外的容量将吞吐量提高了 16.5%，并将中值 TTFT 降低了 95%，因为请求不再等待 KV 空间。

RPA v3 块大小启发式也有类似的盲点。在 v7x 解码期间，它将 KV 计算块设置为等于 KV 获取块，大约 16k 个令牌，这几乎没有为预取缓冲区留下 VMEM。将两者分开块大小将 KV 获取块保持在 16k 令牌，同时将计算块减少到 4k。这为管道提供了领先于 MXU 运行的空间，并将解码吞吐量从每秒 64.9k 令牌增加到 96.3k 令牌，在 Qwen3-0.6B 上的四次运行中再现了 49% 的增益。吞吐量在块大小扫描中遵循清晰的倒 U 形曲线。调整后的参数表无法代表改进，因为它为每个形状存储了单个块大小，这就是更改首先作为环境覆盖进行的原因。后续操作扩展了表以分别保存获取和计算大小。

TPU 系统级协同设计和网络解释器

。谷歌没有专注于最大单芯片性能，而是将计算芯片、芯片到芯片结构和编译器设计在一起，使其能够作为单个系统优化计算和通信。这种联合优化有助于解释 Ironwood 的性价比表现。以下部分研究芯片和网络如何协同工作。 Anthropic 是 TPU 的忠实粉丝，并大量使用它们进行训练。

铁木片

TPUv7 Ironwood 打破了定义 TPU v4 和 TPU v5p 的“MegaCore”惯例，其中两个物理核心被融合到共享一个内存空间的单个逻辑加速器中。相反，Ironwood 有两个独立的计算芯片，每个芯片都运行自己独立的逻辑设备。这些芯片通过高带宽芯片到芯片链接而不是统一的内存结构连接起来。 JAX 和其他框架现在将它们公开为每个芯片两个不同的设备。每个 Ironwood 芯片都带有 2 个 TensorCore 和 4 个第三代 SparseCore。 SparseCores 加速了嵌入查找和其他稀疏操作，否则这些操作会阻碍密集矩阵引擎。

在内存方面，每个芯片的 HBM 容量约为 Trillium 的 6 倍，这一跳跃对 KV 缓存空间和批量大小直接重要。值得注意的是，Ironwood 是第一代具有原生 FP8 硬件支持的 TPU，而前几代必须在软件中模拟 FP8。

在我们的 TPUv7 文章中阅读有关 Ironwood 的更多信息

MXU 以及为什么形状很重要

矩阵乘法单元是实际执行乘法的引擎，它是一个脉动数组：乘法累加单元的二维网格。权重被加载到阵列中并保持静止。激活从边缘涌入，部分总和在网格中逐个单元地波动，并不断累积。最终的结果从另一端流出，而不会在计算过程中触及内存。 v5 之前的每一代 TPU 都使用 128x128 MXU，每个周期适合 16,384 个 MAC。从 TPU v6e 开始一直到 Ironwood，该阵列增加了一倍，达到 256x256，即每个周期 65,536 个 MAC，每个周期的 FLOP 数比之前的设计多 4 倍。

请阅读缩放书（又名机器学习系统圣经），获取有关 TPU 和一般扩展的优秀顶级信息。

SemiAnalysis 是一份读者支持的出版物。要接收新帖子并支持我们的工作，请考虑成为订阅者。

问题是，只有当你能保持它充满时，更大的脉动阵列才是免费的。矩阵尺寸需要在两个方向上至少填充到 MXU 的边长，老一代为 128，v6e 和 v7 为 256，XLA 编译器会尽职尽责地填充任何较小的轴以填充图块。每个填充单元在该周期中仍然占用一个 MAC 单元，乘以零对答案没有任何贡献。 Llama 3 8B 用 128 的注意力头维度说明了这一点。在 Ironwood 的 256x256 MXU 上，该头维度恰好是阵列本机宽度的一半，这将两个注意力矩阵相乘限制在最大 50% MXU 利用率（不是 75%，因为只有头维度为 128）。

这一含义比一个模型的注意力层更深。过去采用任意超参数（例如头尺寸、张量并行分片后保留的 KV 头数和 MoE 专家宽度）的架构选择越来越需要考虑 TPU 块几何形状，因为无论堆栈的其余部分有多好，形状不匹配都会对交付的吞吐量产生直接影响。内核作者还需要明确地考虑这一点。生产注意内核如参差不齐的分页注意力当 XLA 的默认平铺产生低效布局时，使用显式打包尺寸来减少填充。

TPU 很挑剔

GPU 矩阵核心消耗小块，因此各种磁头尺寸、专家宽度和分片后 KV 磁头计数都接近峰值。尺寸为 64 而不是 128 在 H100 或 B200 上几乎没有任何作用，并且获得了更便宜的关注层。 GPU 架构允许研究人员优先考虑评估性能，而不会损失推理性能。但对于 TPU 来说，这个超参数成为一种权衡

如前所述，这些选择在 256 宽脉动阵列上并不是免费的。通过将 Llama 3 8B 的注意力矩阵乘法上限限制为 50% 的算法，在写入一行内核代码之前，头部维度 64 将注意力矩阵乘法上限限制为 25%。 gpt-oss 的发布值为 64。DeepSeek 的 MLA 将其查询/键维度分割为 128 加 64，总共 192，这对于任何 2 的幂数组来说都是笨拙的，对于宽数组来说更糟糕。

其结果是，启动成本差异很大，而且与模型的受欢迎程度关系不大。一个形状完全消失的模型需要调度、服务和调整工作，这些工作以周为单位来衡量。一个对抗瓷砖几何形状的模型需要新的内核才能达到同等水平，更不用说在性价比上获胜了。内核工程师是有限的，因此对硬件未开始运行的模型进行外部化排序是合理的举动，我们希望任何处于此位置的加速器供应商都会采取相同的行为。

环面拓扑

大规模推理和训练运行在成千上万的芯片上，这些芯片需要不断地相互通信。TPU通过一个名为ICI（芯片间互连）的专用网络进行点对点连接。

ICI完全绕过主机CPU，使芯片能够直接交换激活值和梯度，而不是通过PCIe和通用网络接口卡路由。ICI背后的拓扑结构随着每一代的发展而演变：TPU v2和v3使用二维环路（2D torus），每个芯片与4个邻居相连，而从TPU v4和v5p开始，谷歌改用三维环路（3D torus），每个芯片沿 /-X、/-Y 和 /-Z 轴与6个邻居相连。

Ironwood（TPUv7）保持这种三维环路。基本构建模块是由64个芯片组成的4x4x4立方体，其尺寸可以完美映射到一个物理服务器机架上。

环面和普通网格之间的区别在于它的环绕连接。将一条线的两端连接起来就形成了一个环，将最坏情况下的跳数从 N 减少到 N/2。这也是让《吃豆人》的迷宫感觉比实际小的同样技巧。

谷歌在此基础上进一步推出了“扭曲环面”，一种类似莫比乌斯环的环绕方式，更进一步减少了平均跳数。为了超越单个 4x4x4 立方体的规模，谷歌使用光学电路交换机（OCS）将多个立方体连接在一起，从而在更大可重构拓扑中保持环绕特性，可扩展到 Ironwood 的完整 9,216 芯片超级集群及其 42.5 FP8 亿亿次浮点运算总计算能力。

OCS 的操作优势在于，谷歌可以在几秒钟内物理重连故障的链路或损坏的芯片，使用镜子而不是派技术人员去重新连接运行中的数据中心的铜线。

这种横向扩展设计在当时是革命性的。在 NVL72 机架出现之前，那些无法适配单个 8-GPU 节点的模型必须使用流水线并行，因为节点间的 InfiniBand 速度很慢。

在 TPU pod 上，ICI 环状网络为整个 pod 提供了类似 NVLink 的带宽，在 v5p 上可支持多达 8,960 个芯片。有了这种带宽，DSV3 级别的模型可以通过张量并行、专家并行和数据并行（FSDP 风格的权重分片）在整个 pod 上进行分片，而无需在流水线阶段拆分层。

尽管 TPU 环形网络的跳数比单跳 NVLink 交换设计多，但我们即将发布的 CollectiveX/NetworkingX 结果显示，在许多情况下，对于小型 EP 消息，TPU 环形网络的延迟低于单跳 NVSwitch。

展望未来：TPUv8i 的 Boardfly 网络

谷歌新宣布的第八代 TPU 产品线有两款专用芯片：用于训练的 TPU 8t 和用于推理的 TPU 8i。这是谷歌首次将训练和推理分离为不同的架构设计，而不是提供一种同时为两者优化的架构。

TPU 8t 保持了用于规模扩展的 3D 圈拓扑血统，而 TPU 8i 则用一种名为“Boardfly”的新拓扑取代了圈拓扑。这个名字暗示了长期在超级计算中使用的蜻蜓式高基数网络设计。

Boardfly 使用了由高基数交换机组成的更平坦、分层的结构，而不是最近邻网格。Boardfly 拓扑影响每颗芯片的网络连接资本支出。

其收益是网络直径相比类似规模的三维环面减少了50%以上，在1,024到1,152芯片范围内，跳数大约从16跳降至7跳。跳数减少意味着集体操作的尾延迟显著降低，这在跨MoE层路由令牌或运行多轮代理任务时尤其重要，因为每多一次跳转都会累积成用户可感知的延迟。

TPU 8i 以19.2 Tb/s的ICI带宽支持这一点，是上一代的两倍，并配备384 MB片上SRAM，是上一代的三倍，专门用于将推理和代理模型的KV缓存保存在芯片上，而不是来回访问HBM。

不过，这些硬件余量本身并不会带来收益。将直径更小的网络和更大的片上缓存转化为每令牌成本的收益，仍然需要软件堆栈跟上。下一部分将介绍实现这一目标的路线图，包括推测解码、预填充解码的分离，以及在生产中充分利用 TPUv8 所需的更广泛的模型支持。

铺设 TPU 基础的下一步

这些预览结果为在同一硬件上的进一步优化奠定了起点。我们预计随着TPU外部化的推进，性能将会有所提升，但整个软件栈仍有大量工作需要完成。

启用与优化规范解码（MTP）

谷歌团队首先需要关注的是优化推测性解码。一个廉价的小型绘图员预测多个代币，主模型在一次前向传递中验证所有代币，任何不匹配的猜测都会被丢弃。

所有幸存的部分都与模型单独生成的完全相同。推测解码是无损的，质量完全不下降。

它之所以有效，是因为译码受带宽限制，而非计算受限。要为单个用户输出单个令牌，你会将整个模型的权重从HBM流式传输出来。这个权重读取是主导成本，无论你验证一个令牌还是五个令牌，成本几乎没有变化。

矩阵单元大多处于空闲状态，而内存系统则完成所有工作。投机性解码将空闲计算的空闲计算用来摊销一次非常昂贵的权重读段到多个令牌之间，这也是为什么像MTP或DSpark这样一次性草拟多个令牌的方法既便宜又高效。

预填充-解码拆分与KV缓存卸载

预填充-解码（PD）拆分是TPU团队正在外部化的另一项优化。它将预填充和解码分隔在不同的TPU池中，允许每个池独立调整和扩展以匹配工作负载。

谷歌长期以来一直在为Gemini服务运行PD拆分，但外部化工作仅在几个月前才开始。这项工作包括在LLM-d中支持TPU，以及开源TPU-Sync（前称TPU-raiden），这是谷歌的拆分KV缓存传输库。

TPU-Sync原生支持JAX和原生TorchTPU栈，通过提取本地PJRTBuffer硬件描述符实现零拷贝传输。通过拆分PD，我们相信TPUv7可以强劲竞争GB200/GB300，甚至在每美元性能上胜出。

SemiAnalysis 是一份由读者支持的出版物。如需接收新文章并支持我们的工作，请考虑成为订阅者。

TPU-Sync 还支持原生 TPU KV 缓存 DRAM 卸载对于大型模型和中大型批次，HBM无法为所有用户保留KV缓存，这也非常需要。看到谷歌将其DRAM卸载优化外部化向公众展示，令人欣慰。

谷歌正在外部化其原生TPU卸载栈，支持行业标准的月饼店卸载库同时支持月饼商店的DRAM P2P池化。我们相信月饼商店的支持将通过tpu-sync中的原语实现。

通过P2P池化，每个TPU主机上的KV缓存存储被聚合到单一逻辑内存池中，因此任何服务器上的TPU都可以访问其他服务器的KV缓存。这将多个节点的内存贡献统一到一个共享逻辑池中。

此外，Mooncake Store 将多台服务器的 NVMe 存储池整合到单一逻辑池中除了支持传统的分布式文件系统后端，如 WEKA/Vast。

AgentX TPU

KV缓存卸载对于长上下文、多回合代理工作负载尤为重要，优化的KV缓存存储可以实现高KV缓存命中率。

在高层次上，代理性工作负载由四个要素来描述：

多回合：一个会话包含用户与助手之间的数十甚至数百次交互，而聊天机器人场景中只有少数几次。这些工作负载结合了长上下文和高预填充重复使用，以及子代理突发和大量工具调用。
长上下文：系统提示、工具定义以及大量回合使上下文迅速积累。
高前缀重用：由于对话线性进行，通常从第n-1回合输出连接到第n回合，大多数上下文可以从KV缓存中提供，而不必重新计算（这取决于可用存储KV张量的存储量）。随着n的增长，缓存输入与未缓存输入的比例通常趋近于1。
子代理突发：一个会话会启动多个短寿命的子代理，并获得新的上下文，从而创建突发式的KV缓存模式。

Ironwood 没有原生 FP4，所以目前公平的比较是 FP8 Blackwell：FP8 和 FP8 没有质量损失，而 FP4 和 FP8 在 FP4 方面会带来质量差异。

Google 的 TPUv8i 确实有原生 FP4 加速，所以当我们在 InferenceX/AgentX 上提到 TPUv8i 时，我们会比较 FP4 和 FP4。

SemiAnalysis 是一份由读者支持的出版物。如需接收新文章并支持我们的工作，请考虑成为订阅者。

我们坚信，在原生 TorchTPU 稳定基础上构建原生 vLLM 和 SGLang 支持，是实现 TPU 外部化的正确方向。TorchTPU 旨在取代即将被弃用的之前的 TorchAX 协议栈。

谷歌接下来将致力于在单回合和/或代理工作负载如 AgentX 上支持 Kimi K3 和 GLM5.3，同时支持谷歌自家的开放权重模型如 Gemma4。

一旦部分模型在 TorchTPU 栈上运行，添加新模型将变得容易得多，TPU 支持也将在更接近零日时实现。

罗马并非建于零日，因此不应期望TPU的外部化会立刻发生。但我们坚信这一进程正以极其迅速的速度发生。

TPU成本效益总拥有成本

我们对TPUv7的完整BOM和TCO估算已列举在半分析 AI TCO 模型.结构概述如下。

参与讨论