在 Agent 架构下,模型不再是单次调用的函数,而是一个持续运行的推理循环 —— 接收用户请求、调用工具、反思中间结果、决定下一步行动。这种 “推理即循环” 的模式对底层硬件提出了截然不同的要求:不仅是高吞吐,更是低延迟与可预测性。Google 第八代 TPU(代号 Ironwood 后续架构)给出了自己的答案,通过 TPU 8i 与 TPU 8t 的双芯片设计,分别针对推理与服务场景进行深度优化。本文聚焦 TPU 8i,解析其面向 Agent 推理的硬件架构创新。
打破内存墙:384MB 片上 SRAM 的工程逻辑
传统推理系统的核心瓶颈在于内存墙。当模型处理长上下文时,KV Cache(Key-Value 缓存)规模急剧膨胀,如果无法全部驻留在片上 SRAM,就必须频繁从高带宽内存(HBM)读写数据,导致计算单元大量空闲等待。Google 在 TPU 8i 中直接将片上 SRAM 扩容至 384MB—— 这已经是前代产品的 3 倍。这一数字并非随意选择,而是基于生产环境中推理模型 KV Cache 的实际 footprint 进行倒推:只有让活跃工作集完整驻留在硅片上,才能消除预填充阶段与解码阶段之间的内存访问阻塞。
更大的片上 SRAM 直接改变了推理流水线的物理实现方式。在传统架构中,预填充(Prefill)阶段处理长输入时,解码(Decode)阶段往往需要等待 KV Cache 从 HBM 加载完成,形成所谓的 “候诊室效应”。TPU 8i 通过将推理模型的完整活跃状态保留在片上,使预填充与解码之间的状态传递延迟从毫秒级压缩到微秒级。对于需要多轮交互的 Agent 场景,这意味着每一次工具调用后的上下文恢复都足够快,不会成为用户体验的瓶颈。
Boardfly 拓扑:面向全互联通信的网络重构
网络拓扑是 TPU 8i 最容易被忽视却最具战略价值的创新。Google 此前在训练集群中广泛采用 3D Torus 拓扑,这种拓扑在节点间的邻居通信场景下效率极高,但在 Agent 推理场景中暴露了致命缺陷:任意节点之间的通信可能需要跨越 16 跳(以 1024 芯片配置为例),而推理任务(尤其是 MoE 模型)中的全互联通信模式要求任意芯片都能以最低延迟与其他芯片交换数据。
TPU 8i 引入了名为 Boardfly 的层次化拓扑结构,其设计灵感源自 Dragonfly 网络。单个构建块(Building Block)由四个芯片组成环形,内部通过 ICI 链路互联;八个这样的构建块通过铜缆组成一个全互联的组(Group),组间通过光电路交换机(OCS)连接。该结构将 1024 芯片 Pod 的最大网络直径从 16 跳压缩至 7 跳,降低了 56%。对于依赖跨芯片路由 token 的 MoE 推理以及 Chain-of-Thought 等需要跨核心同步的推理模式,这一改进直接转化为更可预测的尾延迟。
Collectives Acceleration Engine:采样瓶颈的硬件卸载
Agent 推理中的另一个关键瓶颈是采样与集体操作。在自回归解码过程中,每个 token 的生成都需要跨核心进行注意力聚合与归约操作,这些全局同步步骤在软件层面实现时往往成为延迟的主要来源。TPU 8i 专门设计了 Collectives Acceleration Engine(CAE),将原本由通用计算单元执行的集合操作卸载到专用硬件,使集体通信的片上延迟降低了 5 倍。
CAE 的设计哲学与 Agent 工作负载高度契合。当多个 Agent 并行运行,每个 Agent 都在进行自回归采样时,集体操作的效率直接决定了系统的整体吞吐量。传统架构中,这些操作需要占用大量计算周期进行跨芯片同步,导致推理核心在等待中空闲。CAE 通过硬件加速将同步开销降至几乎可忽略的水平,使得 TPU 8i 能够在维持低延迟的前提下同时运行数百万个 Agent 实例。
面向部署的实际考量
从系统部署角度,TPU 8i 的创新为生产环境带来了几个可直接参考的工程参数。首先,384MB 片上 SRAM 意味着单芯片可承载的 KV Cache 规模有明确上限,部署长上下文 Agent 时应以该数值为基准进行批处理分组,避免因缓存溢出导致性能退化。其次,Boardfly 拓扑的 7 跳上限为跨 Pod 调度提供了网络直径的硬约束,当 Agent 需要跨多个 Pod 协作时,应优先在单 Pod 内部完成完整推理循环,仅在必要时才进行跨 Pod 通信以规避尾延迟风险。第三,CAE 对集体操作的加速效果在 MoE 模型上最为显著,如果是 Dense 模型则收益相对有限,选型时应根据模型架构进行差异化部署。
Google 官方数据显示,TPU 8i 相比前代在推理场景下实现了 80% 的性价比提升与 2 倍的能效提升。这些收益并非来自单一硬件指标,而是内存、网络、集体操作三个维度协同优化的结果。对于构建大规模 Agent 系统的团队而言,理解这些硬件特性的底层逻辑,是做出高效部署决策的前提。
资料来源:本文核心数据来自 Google Cloud 官方博客对第八代 TPU 的技术解读,涵盖 TPU 8i 的片上 SRAM 规格、Boardfly 拓扑设计及 CAE 性能数据。