十多年来，业界见证了谷歌如何凭借其自主研发的芯片打造出一个庞大的商业帝国。搜索、广告、YouTube 以及运行在 TPU 之上的每一代 Gemini 。很少有加速器像谷歌一样，在架构设计上引起如此多的关注，也很少有加速器像谷歌一样，引发如此多关于其性能和经济效益在设计公司之外会如何发展的讨论。Anthropico 是 TPU 的最大用户，预计到 2029 年，其 TPU 使用量将超过 DeepMind。

谷歌内部的成功从来都不是问题。问题在于，其他公司究竟能从中获得多少优势。你能否利用一个开放权重模型，通过熟悉的推理引擎进行部署，并在对你的业务至关重要的经济效益方面击败英伟达？

今天，Semianalysis发布了 TPUv7 Ironwood 在InferenceX官方预览版上的首个第三方推理测试结果。在与 B200/B300 的直接对比中，Ironwood 的性价比最高可提升 50%。其优势几乎覆盖了整个帕累托曲线，我们从两个方面分析了其经济效益：谷歌内部的总拥有成本 (TCO) 和实际客户支付的外部 TCO。

Ironwood (TPUv7) 是谷歌首次推出可直接购买或通过其云平台租赁的芯片，与其他公司竞争推理工作负载的 TPU 产品。早在 2025 年 11 月，我们就曾提到Anthropic 对 TPU 情有独钟，并承诺采购超过一百万个 TPU（其中约 40 万个为直接购买，60 万个为通过 GCP 租赁），主要用于训练，但也用于推理。我们的加速器模型包含了 Anthropic 和谷歌各季度 TPU 出货量的最新数据，以及 TPUv8i、v8t 和各种 v9/v10 的预估数据等等。

我们对TorchTPU（TPU的外部协议栈）的快速发展感到非常兴奋。本文稍后将讨论TPU软件外部化所需的后续工作，包括优化推测性解码、解耦预填充、键值缓存卸载、多轮代理工作负载等等。尽管如此，SemiAnalysis坚信TPU外部化正朝着正确的方向发展，并全速前进。此外，与仍在学习如何构建测试驱动软件文化的AMD不同，谷歌拥有数十年的软件工程经验和极其成熟的质量驱动文化，因此我们预计外部TPU软件将迅速成熟。

本文将全面介绍 TPU 内核和服务协议栈在开放权重模型方面所做的所有优化，包括动态规划注意力机制优化、MoE 路由和内核优化，以及减少 GDN 内核中的填充。我们还将深入探讨 TPU 系统，并讨论 TPU 性能工程师们为使该协议栈广泛可用而正在推进的下一步工作。

谷歌花了十多年时间展示其利用TPU构建的应用。现在，我们可以衡量一下其他厂商利用TPU能做出什么。

TPUv7 Ironwood对垒Blackwell 和 Blackwell Ultra

我们已经看到，即将推出的原生 TorchTPU vLLM 堆栈在与 Nvidia GPU 的直接对比测试中表现出色。谷歌使用 FP8 中的 Qwen3.5 397B 作为初始启动模型。后续章节将深入探讨为什么新的 TorchTPU 方案相比之前的 TorchAX 方案在外部 TPU vLLM/SGLang 服务方面有了显著改进。

一旦基础架构搭建完成，谷歌计划将支持扩展到其他开源权重模型，包括 Kimi K3 和 GLM5.3。我们相信，一旦少数模型得到充分优化，在正式发布初期为各种流行的开源模型添加优化支持就会容易得多。目前，vLLM 和 SGLang 的正式发布初期支持主要集中在 Nvidia 平台上，对 AMD 平台的支持也尚可。我们预计 TorchTPU 在不久的将来会足够稳定，届时 vLLM 和 SGLang 的维护者可能会将 TPU 添加到正式发布初期支持列表中。该技术栈预计将于 10 月左右结束内部测试并开源。

在聚合服务（FP8）和单token预测的同等条件下，我们发现 TPU 的每美元性能比运行 FP8 的 B200 和 B300 高出 50%。当使用 NVIDIA GPU 上的 FP4 模型提供服务时，性能相比 FP8 有所下降。TPUv7 本身不支持 FP4 计算，因此在 FP4 方面，NVIDIA GPU 仍然保持领先。但随着 TPUv8i 的出现，情况将有所改变，因为 TPUv8i 具有原生 FP4 支持，我们坚信 TPUv8i Boardfly 将能够与 Rubin NVL72 相媲美。

我们相信，一旦谷歌启用解耦服务和推测性解码，TPU 的性价比优势就能得以保持。届时，我们将在对比测试的双方都启用 MTP 和解耦服务。本文稍后将讨论谷歌计划在 TorchTPU 基础架构搭建完成后添加的优化措施。为了控制测试范围，谷歌最初专注于 8k1k 基准测试。然而，我们相信 TorchTPU 技术栈在代理工作负载方面也将表现出色，并将在今年晚些时候发布相关测试结果。本文稍后还将讨论 TPU 外部推理的后续步骤和路线图。

TPU 在性价比方面堪称王者

在每用户每秒生成 100 个token的交互速度下，Ironwood 每百万个token的成本约为 0.181 美元，而 B200 为 0.222 美元，B300 为 0.276 美元。这比 B200 的成本低约 19%，比 B300 的成本低约 34%，同时提供了相同的每用户生成速度。

为了进行比较，我们使用外部 TPU 的总拥有成本 (TCO) 与超大规模实验室购买 TPU 时使用的 B200/B300 TCO 进行比较。去年，SemiAnalysis 加速器模型率先披露了谷歌正在销售 TPU，而不仅仅是通过 Google Cloud 出租。TPU 较低的每小时成本弥补了其在大部分交互范围内较低的吞吐量。SemiAnalysis TCO 模型提供了 TPU 物料清单 (BoM) 估算和 TCO 估算的完整细分。

在大多数原始性能指标上，TPU 的性能都无法超越英伟达 GPU。然而，这种比较并没有考虑到 TPU 更低的总体拥有成本 (TCO)。TPU的买家主要关注的是他们能获得多少收益，也就是每美元的性能和每瓦的性能。

也就是说，在每用户每秒 20 个token的交互速率下，Ironwood 在原始吞吐量方面也遥遥领先，达到每芯片每秒 9,364 个token，而 B200 和 B300 分别为 8,903 个和 8,925 个。这比上述两种 GPU 在这些测试中的吞吐量高出约 5%。结合更低的每小时成本模型，Ironwood 每美元可获得的token数量比 B200 多 50.4%，比 B300 多 96.0%。

如果我们以每芯片小时 1.03 美元的价格来考虑谷歌内部工作负载的 TPU 总拥有成本 (TCO)，那么在并发数为 256 时，其每美元性能优势比 B200 高出 76.7%，比 B300 高出 130.2%。然而，这些高并发性能的提升是以延迟为代价的。例如，在并发数为 256 时，TPU 的平均 TTFT 为 5.41 秒，而 B200 为 3.75 秒，B300 为 2.40 秒。之前讨论的 50% 到 96% 的性能优势仅适用于此数据点，而非所有延迟目标。

从端到端延迟来看，Ironwood 在最高吞吐量点之后仍然保持竞争力。在 20 秒的中位响应时间下，根据帕累托曲线，TPU 的成本约为每百万token 0.098 美元，而 B200 为 0.106 美元，B300 为 0.132 美元。因此，TPU 的成本比 B200 低约 8%，比 B300 低约 25%。

尽管如此，B200 在部分性能对比曲线上仍然略胜一筹，例如在 30 秒左右的中位响应时间附近。然而，在更长的响应时间下，TPU 的每百万token成本优势得以恢复，同时在图中所示的重叠范围内，TPU 的成本优势也优于 B300。

GB300 NVL72与TPUv7的PK

谷歌内部多年来一直在生产环境中运行解耦服务，并且该路径已经过深度优化，但外部 TPU 服务堆栈的解耦路径尚未完全优化。因此，目前在解耦对解耦的对比中，GB200/300 NVL72 在每美元性能方面更具竞争力。但我们预计，随着谷歌、Inferact、RadixArk 和 SemiAnalysis 等公司不断进行优化，这种差距将在几个月内缩小，TPUv7 将能够与 GB200/300 NVL72 相媲美。我们将在后续文章中发布 TPUv7 解耦与 GB200/300 NVL72 解耦的对比评测。TPUv7 pod 可以通过其低延迟的 ICI 架构扩展到 1000 多个芯片，从而实现 NVIDIA NVL72 无法实现的针对大型模型解耦和超宽 EP 的优化。

即使将 TPUv7 聚合服务与 GB300 NVL72 解耦服务进行同等条件下的比较，使用内部总拥有成本 (TCO) 指标，TPUv7 在高延迟和中低延迟场景下也具有相当的竞争力。但当比较中端到端延迟时，GB300 NVL72 解耦服务相比 TPUv7 聚合服务则具有一项优势。

从GB300 NVL72分离式服务与TPUv7聚合服务的性能对比来看，在低端和高端端到端延迟方面，两者具有竞争力；但在中端延迟方面，GB300的性能性价比比TPUv7聚合服务高出约30%。我们坚信，一旦TPUv7上的分离式服务得到优化，它将在所有性能指标上与GB300 NVL72相媲美。

TorchTPU 后端推理,有何不同?

谷歌之前使用 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-inference在有 TPU 优化的 JAX 实现时会使用它。否则，TorchAX 会将原始的 PyTorch 模型转换为 JAX 可以执行的操作。其目标是在无需重写每个 PyTorch 模型的情况下，重用成熟的 TPU 原语和并行支持。

即将推出的第三代设计 TorchTPU 让 vLLM 可以将 TPU 视为原生 PyTorch 设备。它将成为未来 SGLang 和 vLLM 在 TPU 上的首选后端。

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

最初选择 JAX 是出于务实考虑。在上一次公告中，谷歌描述了底层优化、分页注意力机制以及将 vLLM 的工作模型适配到 TPU 执行等方面的困难。他们认为 JAX 能提供更成熟的 TPU 原语和并行支持。TorchAX 让团队无需重写每个 PyTorch 模型就能使用这些功能，而面向 TPU 的 JAX 实现可以共享相同的内核和编译器工作。

使用 TorchAX，开发者仍然用 PyTorch 编写模型，但 JAX 会在 TPU 上执行该模型。TorchAX 充当两个框架之间的转换层。正如TorchAX 文档所述，`torchax.tensor.Tensor` 的行为类似于普通的 PyTorch 张量，但它由 `jax.Array` 提供支持。当模型调用张量操作（在 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 当前的公共 TPU 堆栈使用SGLang-JAX。SGLang -JAX 与 vLLM 的 TorchAX 后端在性能上存在类似的权衡，但采用了不同的方法。它是一个独立的 JAX 原生服务引擎，而非通过 TorchAX 转换的 PyTorch SGLang 运行时。它将 SGLang 风格的调度和前缀缓存与 JAX 模型和 TPU 特有的内核相结合。Google 和 RadixArk 已宣布推出 SGL-torchtpu，作为额外的 PyTorch 原生后端，以扩展可用的服务路径。

TorchTPU将框架边界移到了 PyTorch 本身。它使用 PyTorch 的PrivateUse1 后端扩展点，在 device=”tpu” 上暴露一个普通的 torch.Tensor，而不是一个由 JAX 数组支持的包装器。PyTorch 调度器将 ATen 操作路由到 TPU 后端。开发者可以立即运行代码进行启动和调试，或者调用 torch.compile 来捕获图。在编译路径中，TorchDynamo 和 AOTAutograd 生成 FX 图，TorchTPU 将其降维为 StableHLO，然后 XLA 生成 TPU 可执行文件。谷歌的会议演示文稿明确指出，此路径使用 XLA 编译器，而不是 Inductor 和 Triton。

因此，“原生”一词具有特定的范围。它涵盖 PyTorch 张量、调度、即时执行、编译入口点和分布式 API。XLA 仍然是编译器，内核层继续使用 TPU 特有的实现。TorchTPU 可以调用 Pallas 和基于 JAX 的自定义内核。Helion的 TPU 后端也使用 Pallas。面向用户的框架变为 PyTorch 原生，而编译器和性能关键型内核仍然支持 TPU。

开发者可以使用 `.to("tpu")`，同时保留熟悉的分布式接口和服务引擎代码。谷歌文档支持 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 服务栈，但这并不意味着每个现有内核都能原封不动地迁移。许多方面，例如封装器、张量布局和运行时集成等，可能仍需要进行一些调整和验证。

vLLM 和 SGLang 是两大主流的开源生产级推理服务栈。Inferact、RadixArk 和 Red Hat 正与 Google 密切合作，投入巨资，力求使 TorchTPU 后端在这两大栈中都能提供一流的体验。这种社区支持对于 TPU 推理的外部化至关重要。

截至9月初，TorchTPU仍处于内部测试阶段，预计将于10月中旬的PyTorch大会期间开源。我们很高兴看到它将推动TPU在PyTorch、vLLM和SGLang中实现一流的体验。一旦TorchTPU vLLM和TorchTPU SGLang开源，我们将把InferenceX/AgentX的TPU基准测试从我们的分支迁移到我们的公共代码库。

TPU推理优化深度解析

为了优化原生 TorchTPU 服务栈首个启动模型的 TPU 推理性能，我们投入了数百工时，并提交了大量 PR。正如本文后续部分所述，TPU 架构与 GPU 架构有所不同，因此需要一套专属的优化方案才能高效运行模型。下文将介绍针对 Ironwood 平台上 Qwen3.5-397B 及其他模型所采用的优化措施，这些措施显著提升了（或在某些情况下实现了）性能。大多数 Pallas 内核优化方案都可以移植到新的原生 TorchTPU 栈中。

并非TPU特有的一个问题是如何有效地将并行性从模型架构映射到硬件。例如，Qwen3.5的GQA注意力层有32个查询头，但只有2个共享的键值头。在TP8（八个逻辑设备）中，查询计算可以均匀地分布在8个设备上，每个设备4个查询头。然而，键值头在这些设备上的分配并不均匀。每个键值头被16个查询头共享，这意味着四个设备需要相同的键值数据来计算其本地注意力。当TP超过键值头数量时，标准做法是在TP进程间复制键值头。vLLM已经支持这种做法。TPU后端需要进行兼容性修复才能激活这种现有行为，并避免不必要的全对全通信。

为了实现更高并发的服务，TPU 后端还增加了对八路注意力数据并行（DP8）与八路专家并行（EP8）相结合的支持，简称为“DP 注意力”或 DEP8。与将每个请求的注意力分散到所有八个设备不同，每个设备处理不同的请求子集，并将这些请求的两个键值对（KV）头保存在本地。注意力权重会被复制，但键值缓存保存的是不同的请求历史记录，而不是同一历史记录的重复副本。同时，512 个路由专家仍然分布在各个设备上，避免了对规模更大的专家权重池进行复制。为了实现这一点，需要在服务引擎中协调请求分配、循环状态槽和块表。

借助 DP 注意力机制和 EP，TPU GroupedGEMM 实现会在专家运行之前收集所有token激活信息和路由元数据，然后使用 reduce-scatter 算法对加权输出求和，并将每个token的结果返回到其注意力排名。后端最初分别收集选定的专家 ID 和路由权重。谷歌将这些信息合并为一个 all-gather 操作，从而避免了额外的收集操作，因为对于如此小的数组，额外的收集操作的延迟可能会主导传输时间。PR 报告称，在其 DeepSeek -V3 测试中，每层大约节省了 80 微秒。虽然 80 微秒听起来不多，但 DeepSeek -V3/R1 有 58 层需要进行此 AllGather 操作。在这些层中，每次前向传播单独节省的时间约为 4.64 毫秒（80 微秒 × 58），从而降低了 TPOT 并提高了交互性。

Google 还在SparseCore 上实现了 ReduceScatter 集体计算（SparseCore 更适合数据移动等不规则操作），并利用 Ironwood 更快的芯片间链路，在芯片间交换部分和之前，先将每个芯片内的贡献合并。这种每芯片双设备的布局以及芯片间的 ICI 网络将在下一节中介绍。双缓冲允许内核在累积另一个数据块的同时传输一个数据块，从而将本地归约和芯片间传输与速度较慢的芯片间流量重叠。在 SparseCore 上运行集体计算还可以释放 TensorCore 的执行资源，用于其他操作。

资料来源：谷歌和SemiAnalysis

在上图中，通信阶段相互重叠，从而缩短了端到端执行时间。例如，考虑标记为 t_1 的点。此时，片内 DMA ScatterReduce (P1) 已完成 MB 1 的处理。现在可以对 MB0 执行跨 ICI dim 0 的 ScatterReduce (P2.0)，同时对后续 MB 执行 P1。与基线相比，此优化在 8k1k 架构上，并发数从 64 到 512 时吞吐量提高了 4.1% 到 14.2%，其中并发数为 256 时提高了 8.5%，在 1k8k 架构上，并发数为 512 时提高了 26.1%。

将此集合从 TensorCore 迁移到 SparseCore 也意味着在集合运行期间，可以在 TensorCore 上流水线处理其他工作。然而，并非所有集合都值得迁移。在某些情况下，将集合卸载到 SparseCore 实际上会降低性能。例如，对于 Qwen3.5，现在有一个阈值来决定何时将所有 reduce 和 所有 gather 操作卸载到 SparseCore 。如果将小型集合保留在 TensorCore 上，即使它们可以放入 VMEM 并且 SparseCore 的卸载开销会增加，也会降低它们的速度。默认阈值取决于 VMEM 容量。PR 报告称，在并发数为 64 时，8k1k 吞吐量提高了 2.7%；在并发数为 128 时，吞吐量提高了 5.7%。

优化 MoE 路由和专家内核

混合专家层会为每个专家生成不规则的token组，这些杂乱的token组必须被重塑成 TPU 矩阵单元能够高效处理的形式。以下部分更改是对路由和分组矩阵乘法后端的一般性改进，适用于任何混合专家模型。其他更改则是在 Qwen3.5 上直接测量的。

第二版分组矩阵乘法算法改进了专家输入传递给 MXU 的方式。它移除了冗余的瓦片计算，将大小转移为有效行数而不是填充后的最大值，对专家权重进行三重缓冲，以便在计算当前组权重的同时，下一组的权重已经开始计算，并将组元数据生成融合到内核中。

随后，专家输入的不规则重组被转移到 SparseCore 上。SparseCore 负责数据移动，将每个专家的 token 收集到连续的组中，而 TensorCore 则负责运行专家矩阵乘法。后续的重写改进了内存读取流水线，并将合并工作分配到各个 token 和隐藏维度上。PR 报告显示，与原始 SparseCore 内核相比，8k1k 服务吞吐量提高了 12%，同时 TTFT 和 TPOT 也更低。

进一步的优化将不规则的 gather-reduce 路径中的 top-k 权重收集操作移到了 SparseCore 内核。这里的 gather 指的是从单个设备的内存中读取选定的条目，而不是跨设备进行全设备收集。此前，TensorCore 在预处理期间收集路由权重和源索引，这限制了 TensorCore 和 SparseCore 的重叠。将这些收集操作移到 SparseCore 内核中，在 DeepSeek -V3 微基准测试中，使用 16 路专家并行处理，批处理大小为 2k，将 TensorCore 的开销从 29 µs 降低到 14 µs，并将整体操作延迟从 146 µs 降低到 137 µs。

对于小批量数据，通用的不规则路径成本高于其所处理的工作量。因此，专门的小批量置换算法会构建独热矩阵，并使用普通矩阵乘法将标记置换为其对应的专家标记，然后再对结果进行逆置换。在 8k1k 工作负载下，当并发数为 64 时，吞吐量提高了 7.3%；当并发数为 128 时，吞吐量提高了 5.1%。

一项正在进行中的更改将专家 ID 和token索引打包到一个排序键中，因此 XLA 可以执行更简单的排序，同时保持所需的顺序。排序延迟从 106.6 微秒降至 21.7 微秒。PR 报告称，8k1k 服务性能提升了 0.6% 到 8.5%，但同样的更改也支持 FP8 全集收集。

优化门控 DeltaNet Pallas 内核

门控 DeltaNet (GDN) 是一种循环计算。也就是说，每一步运行状态都会衰减，并更新一个秩为 1 的项，然后进行投影以产生输出。请参阅以下文章，深入了解 GDN 等线性注意力机制：

以下变更展示了其矩阵运算、向量更新和状态转移如何在 MXU、VPU、VMEM 和 HBM 之间进行调度。Qwen3.5的初始支持添加了纯 JAX 实现的因果 Conv1D 和 GDN，并将它们连接到 TPU 操作符调度，同时启用了循环状态缓存。后续的 PR 对这些实现进行了优化。

另一个易于实现的优化方法是重新排列输出投影计算中的代数运算，使 MXU 和 VPU 的工作重叠。此前，VPU 首先将排名为 1 的更新应用于衰减后的状态，然后 MXU 将更新后的状态乘以查询以生成输出。这种依赖关系迫使 MXU 等待。

将衰变状态记为 S，校正向量记为 Δ，密钥和查询分别记为 k 和 q，则输出可以展开为：

现在，MXU 可以直接从衰减状态计算 Sq，而 VPU 则为下一个 token 构建更新状态。当前更新对输出的贡献是单独计算的，使用 kᵀq（每个头进行标量点积运算），然后进行一个小的缩放向量加法。这使得状态更新从 MXU 的依赖路径中移除。据报道，在并发数为 64 时，8k1k 吞吐量提升了 2.79%；在并发数为 512 时，吞吐量提升了 4.48%。

下一个改动通过在解码循环内部对 Q 和 K 进行切片来减少向量寄存器溢出，从而减少同时保持有效的值数量。解码 64 位内核的速度提升了约 20%，但端到端性能提升较小，在并发数为 512 时，8k1k 架构的性能提升为 0.8%，1k8k 架构的性能提升为 3.8%，因为内核只是解码步骤的一部分。

异步状态传输利用双缓冲技术将DMA与计算重叠。第二组缓冲区消耗的额外VMEM最初导致数据并行注意力下降。重用现有暂存缓冲区并缩短临时缓冲区的生命周期恢复了该容量并消除了性能下降。据报道，在并发数为512时，8k1k吞吐量提升了11.3%。

GDN v3将 Conv1D 和 GDN 融合到一个内核中，从而减少了 HBM 往返次数，改进了预填充布局，并将混合预填充/解码执行统一到一条路径中。据报道，内核级加速比分别为：解码 1.41 倍，预填充 1.60 倍，混合批次 2.14 倍。这些测量结果仅涵盖内核层面，并未体现端到端服务的性能提升。

只需提高内核内各操作之间的重叠度，即可显著提升性能！

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

Qwen3.5 是一种混合模型，包含两种状态：GQA 层会累积一个随每个 token 增长的键值历史，而 GDN 层则为每个请求维护一个固定大小的循环状态。状态分配、存储精度、注意力页面大小和物理数据布局共同决定了 HBM 的实际可用量以及注意力机制的运行效率。

批量分页注意力机制将序列批量处理，预先计算页面元数据，并使用三重缓冲来改进流水线并减少填充。此更改为共享注意力后端奠定了基础。PR 中的工作负载示例使用 Qwen3-32B，因此其测量结果应与 Qwen3.5-397B 的结果分开。

现在循环状态采用紧凑分配方式，每个活动请求大约占用一个槽位，而不是像以前那样num_blocks为每个层组分配槽位。在报告的配置中，这回收了约 76 GiB 的 HBM，并将注意力块池扩大了 71%，在并发数为 64 时，1k8k 输出吞吐量提高了 18%。

将循环状态存储在 BF16 中，可使其 HBM 占用空间减半，同时将 FP32 运算保留在 VMEM 内。这既节省了内存容量和传输带宽，又保留了 FP32 运算。据报道，在并发数为 512 时，1k8k 吞吐量提升了 15%。

移除旧的混合页面大小对齐约束后，批量注意力机制可以使用合适的 2 的幂次方页面大小，例如 256 个 token。在并发数为 512 时，1k8k 吞吐量提升了约 7%。这适用于非前缀缓存路径，并且与之后引入的对齐检查点模式不同。

通过基于键值对数量选择重塑路径，而不是施加仅其他形状才需要的布局约束，可以避免不必要的键值布局复制。在并发数为 512 时，Qwen3.5 8k1k 的吞吐量提高了约 4.1%。

减少填充和低并发开销

当只有四到八个请求同时进行时，针对数百个并发请求优化的配置会造成资源浪费。编译后的形状桶过大，元数据的大小被限制在配置的最大值范围内，填充标记还会触发不必要的专家计算。InferenceX 的最佳运行点恰好位于这种低并发情况，这也促使我们进行了以下更改。

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

针对 InferenceX 的显式 Qwen3.5 调优，缩小了旋转表，调整了每个 DP 等级的序列限制，并为并发数为 4 的情况添加了一个专用的注意力桶。在并发数为 4 的情况下，8k1k 的性能提升了 13.3%，1k1k 的性能提升了 15.5%。

进一步的低并发调优切换到TP8注意力机制并启用专家并行，减小最小token桶大小，并将填充token路由到专家零，以避免触发额外的专家权重负载。在并发数为4时，1k1k的综合性能提升了22.9%，在并发数为8时提升了18.1%。在8k1k上，并发数为4时吞吐量提升了9.2%，但在并发数为8时下降了5.3%。

为混合模型启用前缀缓存

上述性能提升是在随机输入基准测试中测得的，其中请求之间没有任何共享信息。前缀缓存对于代理和多轮对话工作负载至关重要，因为这些工作负载会重用较长的系统提示和对话历史记录。对于混合模型，缓存的前缀必须在前缀末尾同时保留其键值块和 GDN 循环状态。该循环状态通常会在请求继续进行时被覆盖，这使得前缀缓存比“常规”注意力模型更难。

支持 DP 的混合前缀缓存通过为 GDN 提供单独的槽位来读取检查点和写入实时状态，从而解决了这个问题。因此，继续执行请求不再会覆盖与其缓存前缀关联的检查点。状态地址源自块表，该结构也用于定位 KV 块。检查点以对齐的缓存粒度获取，因此保存的状态始终与 KV 块边界对齐。此模式需要一个完整的检查点池，而不是前面描述的紧凑的按请求分配，它牺牲了一些 HBM 以换取跨请求的重用。当实际共享前缀时，这种优势尤为明显。

分页注意力中的通道布局和管道深度

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

在批量注意力内核中，KV缓存沿头部维度打包键值对，对于FP8，打包因子为4。每个设备只有一个KV头部的模型只有两个需要打包的内容，因此每个tile的一半空间都被浪费在了填充上。序列通道布局通过将页面标记放在128通道轴上，并将头部维度放在子通道轴上来解决这个问题。可用KV页面数量翻倍（在报告的配置中从5141增加到10283），并且头部维度只需要是32的倍数而不是128，这使得头部维度为64的模型也能使用该内核。在低并发情况下，这种布局会使每个标记的延迟增加约3%，但在8k1k架构上并发128时，额外的容量使吞吐量提高了16.5%，并将中值TTFT降低了95%，因为请求不再需要等待KV空间。

RPA v3 的块大小启发式算法也存在类似的盲点。在 v7x 解码过程中，它将 KV 计算块的大小设置为与 KV 获取块的大小相同，约为 16k 个 token，这导致预取缓冲区几乎没有 VMEM 可用。将这两个块大小分开后，KV 获取块的大小保持在 16k 个 token，而计算块的大小减少到 4k 个 token。这使得流水线能够在 MXU 之前运行，并将解码吞吐量从每秒 64.9k 个 token 提高到每秒 96.3k 个 token，在 Qwen3-0.6B 上进行了四次运行，结果均显示吞吐量提升了 49%。吞吐量在整个块大小范围内呈现出清晰的倒 U 型曲线。由于调整后的参数表每个形状只存储一个块大小，因此无法体现这一改进，这也是为什么该更改最初是以环境覆盖的形式实现的。后续版本扩展了该表，使其能够分别存储获取块大小和计算块大小。

TPU系统级协同设计与网络化详解

正如我们在此前文章中所说，谷歌在推理方面的每token成本优势源于协同设计。谷歌并没有专注于最大化单芯片性能，而是将计算芯片、芯片间互连架构和编译器协同设计，从而能够将计算和通信作为一个整体系统进行优化。这种联合优化有助于解释Ironwood的性价比优势。以下章节将探讨芯片和网络如何协同工作。Anthropic是TPU的忠实拥趸，并大量使用TPU进行训练。

Ironwood芯片

TPUv7 Ironwood 打破了 TPU v4 和 TPU v5p 定义的“MegaCore”设计理念，后者将两个物理核心融合到一个共享同一内存空间的逻辑加速器中。Ironwood 则采用两个独立的计算芯片，每个芯片运行其自身的独立逻辑设备。这些芯片通过高带宽的芯片间链路连接，而非统一的内存结构。JAX 和其他框架现在将它们作为每个芯片上的两个独立设备公开。每个 Ironwood 芯片包含 2 个 TensorCore 和 4 个第三代 SparseCore。SparseCore 可以加速嵌入查找和其他稀疏运算，否则这些运算会使稠密矩阵引擎不堪重负。

在内存方面，每颗芯片的HBM容量约为Trillium的6倍，这一提升直接关系到键值缓存的余量和批处理大小。值得一提的是，Ironwood是首款原生支持FP8硬件的TPU产品，而之前的几代产品则需要通过软件模拟FP8。

MXU 以及形状为何重要

矩阵乘法单元（MXU）是实际执行乘法运算的引擎，它是一个脉动阵列：一个二维的乘加单元网格。权重被加载到阵列中并保持静止。激活信号从边缘流入，部分和逐个单元地在网格中传播，并不断累积。最终结果从另一端流出，计算过程中无需访问内存。从 TPU v5 到 v5 的每一代都使用 128x128 的 MXU，每个周期可执行 16,384 次 MAC 运算。从 TPU v6e 开始，一直到 Ironwood，阵列尺寸翻倍至 256x256，每个周期可执行 65,536 次 MAC 运算，每个周期的 FLOP 运算量是前代设计的 4 倍。

关键在于，更大的脉动数组只有在能够保持满载的情况下才是免费的。矩阵维度需要在两个方向上至少填充到 MXU 的边长，旧版本为 128，v6e 和 v7 版本为 256，XLA 编译器会忠实地填充任何较小的轴以填满图块。每个填充的单元格在该周期内仍然占用一个 MAC 单元，乘以一个对结果没有任何贡献的零。Llama 3 8B 以 128 的注意力头维度为例对此进行了说明。在 Ironwood 的 256x256 MXU 上，该注意力头维度正好是数组原始宽度的一半，这使得两个注意力矩阵乘法运算的 MXU 利用率最高只能达到 50%（而不是 75%，因为只有注意力头维度是 128）。

其影响远不止于单个模型的注意力层。过去可以随意设置的架构超参数（例如头部尺寸、张量并行分片后保留的键值对头部数量以及 MoE 专家宽度）现在越来越需要考虑 TPU 瓦片的几何形状，因为形状不匹配会直接影响吞吐量，无论堆栈的其他部分性能如何。内核作者也需要明确地考虑这一点。像Ragged Paged Attention这样的生产级注意力内核会使用显式的打包尺寸来减少填充，因为 XLA 的默认瓦片布局效率低下。

TPU很挑剔

GPU 矩阵核心消耗的图块较小，因此各种不同的头部尺寸、专家宽度和分片后键值对头部数量的影响都接近峰值。在 H100 或 B200 上，64 维（而非 128 维）的性能提升几乎可以忽略不计，并且可以获得更便宜的注意力层。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 使用二维环面，每个芯片连接到 4 个相邻芯片；从 TPU v4 和 v5p 开始，谷歌转向三维环面，每个芯片沿 +/-X、+/-Y 和 +/-Z 轴连接到 6 个相邻芯片。Ironwood（TPUv7）保留了这种三维环面结构。其基本构建模块是一个 4x4x4 的立方体，包含 64 个芯片，尺寸设计为可以完美地映射到一个物理服务器机架上。

环面与普通网格的区别在于其环绕连接。连接线段的两端形成一个环，将最坏情况下的跳数从 N 减少到 N/2。这与吃豆人游戏中迷宫看起来比实际更小的原理相同。谷歌更进一步，开发出“扭曲环面”，一种类似莫比乌斯环的环绕结构，可以进一步减少平均跳数。为了扩展到单个 4x4x4 立方体之外，谷歌使用光路交换机 (OCS) 将多个立方体拼接在一起，从而在更大的可重构拓扑结构中保持环绕特性，并扩展到 Ironwood 数据中心完整的 9216 芯片超级节点及其 42.5 FP8 exaflops 的聚合计算能力。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采用更扁平、分层的高基数交换机结构。Boardfly拓扑结构会增加每个芯片的网络接入成本。

与同等规模的 3D 环面相比，TPU 8i 的优势在于网络直径减少了 50% 以上，在 1024 到 1152 个芯片的规模范围内，跳数从大约 16 跳减少到大约 7 跳。更少的跳数意味着集体操作的尾延迟显著降低，这在跨 MoE 层路由token或运行多轮代理工作负载时至关重要，因为每增加一跳都会累积成用户可见的延迟。TPU 8i 还提供了 19.2 Tb/s 的 ICI 带宽（是上一代的两倍）和 384 MB 的片上 SRAM（是上一代的三倍），其容量专门用于存储片上推理和代理模型的 KV 缓存，而无需往返 HBM。

然而，这些硬件上的提升空间本身并不能带来实际效益。要将更小的网络直径和更大的片上缓存转化为token成本的降低，软件栈仍然需要跟上步伐。下一节将介绍实现这一目标的路线图，包括推测性解码、预填充解码分离以及在生产环境中充分利用 TPUv8 所需的更广泛的模型支持。

构建TPU基础的下一步工作

这些预览结果为在同一硬件上进行进一步优化奠定了基础。我们预计随着 TPU 外部化的持续推进，性能将得到提升，但软件栈的各个层面仍需进行大量工作。

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

谷歌团队首先需要关注的是优化推测性解码。一个低成本的小型预测器会预测几个标记，主模型会在一次前向传播中验证所有这些预测，任何不匹配的预测都会被丢弃。所有保留下来的标记都与模型自行生成的标记完全相同。推测性解码是无损的，不会造成任何质量损失。

其原理在于解码过程受限于带宽而非计算能力。要为单个用户生成一个token，需要从 HBM 中读取整个模型的权重。读取权重是主要的开销，无论验证一个token还是五个token，开销几乎不变。在内存系统完成所有工作时，矩阵单元大多处于空闲状态。推测性解码利用这些空闲的计算能力，将一次非常昂贵的权重读取操作分摊到多个token上，这就是为什么像 MTP 或 DSpark 这样一次性生成多个token的方法如此高效且成本低廉的原因。

预填充解码分解和键值缓存卸载

预填充解码 (PD) 分离是 TPU 团队正在努力实现的另一项优化。它将预填充和解码操作分离到不同的 TPU 池中，从而允许每个池独立进行调优和扩展，以匹配工作负载。

谷歌内部早已为 Gemini 服务运行 PD 解耦技术，但将其外部化的工作直到几个月前才开始。这项工作包括在 llm-d 中添加 TPU 支持，以及开源 TPU-Sync（原名 TPU-raiden），即谷歌的解耦键值缓存传输库。TPU-Sync 可与 JAX 和原生 TorchTPU 协议栈原生协作，并通过提取原生 PJRTBuffer 硬件描述符来实现零拷贝传输。我们相信，借助解耦 PD 技术，TPUv7 将成为 GB200/GB300 的强劲对手，甚至在性价比方面超越它们。

TPU-Sync 还支持原生 TPU KV 缓存 DRAM 卸载，这对于大型模型和中大型批次大小的数据处理至关重要，因为此时 HBM 内存已无法满足所有用户的 KV 缓存需求。很高兴看到谷歌也将其 DRAM 卸载优化技术公开。

Google 正在将其原生 TPU 卸载栈外部化，并支持行业标准的 Mooncake Store 卸载库以及 Mooncake Store 的 DRAM P2P 池化支持。我们相信 Mooncake Store 的支持将使用 tpu-sync 中的原语来实现。

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

此外，Mooncake Store 将来自多个服务器的 NVMe 存储汇集到一个逻辑池中，并且还支持像 WEKA/Vast 这样的传统分布式文件系统后端。

AgentX TPU

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

从宏观层面来看，智能体的工作负荷具有四个要素：

多轮交互：一次会话包含用户与助手之间数十次甚至数百次交互，而聊天机器人场景中通常只有几次交互。这类工作负载结合了长时间上下文处理、高预填充重用、子代理突发以及大量的工具调用。

上下文信息丰富：系统提示、工具定义以及大量的操作回合使得上下文信息积累迅速。

前缀重用率高：由于对话以线性方式进行，第 n-1 轮的输出通常会连接到第 n 轮，因此大部分上下文信息可以从键值缓存中获取，而无需重新计算（这取决于可用于存储键值张量的存储空间大小）。随着 n 的增大，缓存输入与未缓存输入的比率通常趋近于 1。

子代理突发：一个会话启动多个具有全新上下文的短生命周期子代理，从而创建突发的 KV 缓存模式。

我们听说很多 Google TPU 用户都在询问 AgentX 的性能结果，我们很高兴地宣布，我们正在努力将 TPU 也引入 AgentX！

Ironwood 本身不支持 FP4，因此目前最公平的比较对象是 Blackwell 的 FP8：FP8 对 FP8 之间没有质量损失，而 FP4 对 FP8 则会在 FP4 端引入质量差异。Google 的 TPUv8i 具有原生 FP4 加速功能，因此当我们在 InferenceX/AgentX 上启用 TPUv8i 时，我们将进行 FP4 对 FP4 的比较。

我们坚信，在原生 TorchTPU 稳定的基础上构建原生 vLLM 和 SGLang 支持，是 TPU 外部化的正确方向。TorchTPU 旨在取代即将弃用的 TorchAX 堆栈。接下来，Google 将致力于在 AgentX 等单轮迭代和/或代理工作负载上启用 Kimi K3 和 GLM5.3，以及 Google 自有的开放权重模型，例如 Gemma4。一旦一些模型在 TorchTPU 堆栈上运行，添加新模型将变得更加容易，TPU 支持也将更接近正式发布之日。

罗马不是从零开始建成的，所以我们也不应指望TPU的外部化会立即发生。但我们坚信，它正在以极快的速度发生。

本文来自微信公众号“半导体行业观察”（ID：icbank），作者：Semianalysis