确定 TPU 的经验性能，需要超越产品规格参数，转向细粒度的真实测量。理论吞吐量指标虽然表明其拥有强大的计算能力，但关键仍在于：这些加速器在特定架构环境和工作负载约束下实际表现如何。为此，我们编写了一个用于评估 TPU 示例工作负载的微基准测试套件。

微基准测试是评估 TPU 等加速器的关键工具，因为它能够对基础硬件能力提供细粒度、量化的测量。它将性能评估拆分为多个关键功能领域，使你能够准确判断设备是否达到了其理论性能规格，并识别特定的性能差距或架构相关瓶颈。

本文将深入介绍该微基准测试套件：它测量什么、它与 TPU 架构有何关系，以及如何对其进行调优。我们还将展示一个示例工作负载的案例研究。

核心类别与指标

该微基准测试套件围绕五个核心领域提供标准化性能评估：

• 网络（Network）：刻画对跨芯片互连（ICI）进行模型多芯片扩展至关重要的集合通信操作的互连性能。关键测试包括 all-gather、all-reduce、reduce-scatter 和 all-to-all。指标包括吞吐量（GB/s）和延迟（秒）。
• 计算（Compute）：量化原始计算吞吐量。主要操作是通用矩阵乘法（GEMM/Matmul），以 TFLOPs 和模型 FLOPs 利用率（MFU）衡量；MFU 表示工作负载对 TPU Matrix Unit（MXU）架构的利用效率。
• 高带宽内存（HBM）：测量片上内存的有效读写带宽。这用于评估大规模向量操作期间基础内存访问的速度；这类操作旨在使内存互连充分饱和。核心指标是以 GB/s 衡量的内存带宽。
• 主机传输（Host Transfer）：通过测量主机 CPU 与加速器 HBM 之间经 PCIe 互连的数据传输速率，刻画 I/O 效率。这包括主机到设备（H2D）和设备到主机（D2H）传输，单位为 GB/s。
• Ragged-Paged Attention（RPAv3）：聚焦于面向 Transformer 模型的基准测试，特别衡量 Attention（tokamax splash）和批量矩阵乘法（BMM）等原语的吞吐量。这些测试有助于预测端到端推理性能，并识别首 Token 时间（TTFT）延迟中的瓶颈。

对 TPU 模块进行微基准测试

我们套件中的 Network、Compute、HBM 和 Host Transfer 微基准测试针对相互隔离的组件进行测试，从而对理论性能提供细粒度评估。这种详细的组件分析对于从多个维度调试和优化系统性能至关重要。例如，它可以帮助判断设备是否达到了规格目标，或者是否存在 VMEM 局部性、互连效率等特定架构瓶颈，限制了工作负载的规模和速度。

微基准测试套件

进行此类经验性评估的核心工具，是我们的 accelerator-microbenchmarks GitHub 仓库。该套件通过多个关键目录，为硬件验证提供结构化方法：

• src/：包含微基准测试的基础实现和源代码。
• configs/：管理每项测试的调优参数，例如数据缓冲区大小、迭代次数以及支持的数据类型（如 bf16、fp8、fp32）。

shell

使用示例配置运行计算微基准测试

kubectl apply -f Ironwood/guides/collectives/tpu7x-2x2x1-ici-all-gather-microbenchmark.yaml

通过使用这些配置方案，你可以将前文讨论的高层指标（如 ICI 带宽或 MXU 利用率）与实际运行时数据关联起来。例如，configs 文件夹中的集合通信 YAML 文件可直接测量架构部分所述的 ICI 环面（torus）效率，有助于确保物理互连健康状况符合理论预期。

性能调优的影响

微基准测试将性能优化从试错过程转变为一门经验工程学科。通过建立“光速”（Speed-of-Light，SOL）基线，即硬件模块的理论极限，它为训练和推理效率提供了明确的优化目标。

Roofline 模型：诊断的北极星

微基准测试的核心价值在于其能够定义 Roofline 模型的边界。通过测量 MXU 峰值吞吐量和 HBM 带宽，开发者可以立即将工作负载的瓶颈归类为：

• 计算受限（Compute-bound）：模型已触及 MXU 平台上限。调优内存访问不会带来收益；优化必须聚焦于内核效率或减少 FLOPs。
• 内存受限（Memory-bound）：模型受到 Roofline 倾斜“屋檐”部分的限制。可通过优化数据局部性（VMEM）或减少 HBM 流量来提升性能。
• 网络受限（Network-bound）：ICI 或 DCN 中的停顿导致芯片无法发挥其计算潜力。这表明需要采用更好的分片策略或通信重叠技术。

硬件感知型模型架构

微基准测试揭示了 TPU 硬件独特的“个性”。例如，在 Ironwood（TPU7x）架构上，256x256 脉动阵列对操作数形状施加了物理约束。基准测试证实，使用 head_dim 为 128 的模型（在较早版本的 Llama 中较为常见）会获得较低的 MXU 利用率。这一洞察促使研究人员协同设计前沿模型，使其维度与 256 字节边界对齐，从而最大化硬件效率。

软件调节手段与优化策略

一旦识别出瓶颈，微基准测试可以帮助你选择合适的软件调节手段：

• 内核选择：你可以比较标准 JAX 操作与专用 Pallas 内核或 Tokamax Splash Attention。微基准测试可为哪种实现最能降低反向传播开销提供可靠依据。
• 分片与网格调优：通过测量不同拓扑结构下的集合通信性能（例如 2x2x2 与 4x2x1），你可以优化 Fully Sharded Data Parallel（FSDP）参数，确保通信时间被计算完全掩盖。也就是说，任务完成得足够快，从而 TPU 无需停下来等待响应。
• 重物化（Rematerialization）：微基准测试有助于计算计算周期与内存压力之间的精确权衡，指导何时应将计算数据存储在内存中，以及何时应“重物化”（重新计算）数据，以避免 HBM 溢出。

预测性优化

最后，微基准测试数据支持对大规模部署进行预测建模。你无需在完整 TPU Slice 上运行成本高昂的测试，而是可以利用测试结果构建分析模型，在扩展之前准确预测训练吞吐量、TTFT、每个输出 Token 时间（TPOT）和 MFU。

案例研究：Ironwood（TPU 7x）性能分析

在一个示例中，我们分析了部署在 4x4x4 TPU 7x 配置上的 110B 混合专家（MoE）训练工作负载的性能调优。我们利用微基准测试建立硬件基线，并推动对复杂稀疏工作负载的优化。结果如下：

• Roofline 分析：对训练工作负载的诊断分析揭示了性能瓶颈的清晰分化。前向传播和稠密核心操作主要受计算限制，达到 1.85 PFLOPS。相反，MoE 架构中的路由和 Attention 原语仍受内存限制；由于 HBM 饱和，它们的性能仅达到 SOL 的 30% 至 60%。在专家分发期间，稀疏激活模式变得高度 I/O 受限，需要将输入请求发送到子网络。
• 与微基准测试的关联：微基准测试结果直接指导了优化措施，使训练步耗时缩短了 21.2%。ICI 集合通信基准测试识别出通信停顿问题，随后通过 SparseCore collective offloading 得以解决；HBM 带宽测试则识别出受内存限制的 Attention 原语，进而推动采用 Tokamax Splash Attention。

为 TPU 性能优化带来工程纪律

总之，微基准测试套件提供了必要的结构化验证，弥合理论规格与实际运行时性能之间的差距，使你能够在大规模部署中充分释放 TPU 架构的潜力。

如需了解更多信息，请参阅以下资源：

1. Ironwood：面向推理时代的首款 Google TPU
2. 在 TPU v7x（Ironwood）上优化前沿模型训练
3. AI 加速器性能与基准测试
4. 开发者使用 Ironwood TPU 训练指南

Determining the empirical performance of TPUs requires moving beyond product specifications to granular, real-world measurements. While theoretical throughput metrics indicate significant computational power, the critical measure remains how these accelerators behave within specific architectural environments and workload constraints. To help, we authored a microbenchmark suite to evaluate TPUs for sample workloads.

Microbenchmarks are critical tools in evaluating accelerators like TPUs because they offer granular, quantitative measures of fundamental hardware capabilities. They segment performance evaluation into key functional areas, allowing you to accurately assess whether a device is achieving its theoretical performance specifications, and to identify specific performance gaps or architecture-specific bottlenecks.

In this blog, we take a closer look at the microbenchmark suite: what it measures, how it relates to the TPU architecture, and how to tune it. We also present a case study on a sample workload.

Key categories and metrics

The microbenchmark suite provides standardized performance evaluations across five core areas:

• Network : Characterizes interconnect performance for collective communication operations essential for scaling models across multiple chips through inter-chip interconnect (ICI). Key tests include all-gather , all-reduce , reduce-scatter , and all-to-all. Metrics include throughput (GB/s), and latency (seconds).
• Compute : Quantifies raw computational throughput. The primary operation is General Matrix Multiply (GEMM/Matmul), measured in TFLOPs and Model FLOPs Utilization (MFU), which indicates how effectively the workload utilizes the TPU's Matrix Unit (MXU) architecture.
• High Bandwidth Memory (HBM) : Measures effective bandwidth to and from on-chip memory. This assesses the speed of basic memory access during large vector operations intended to saturate the memory fabric. The core metric is memory bandwidth measured in GB/s.
• Host Transfer : Characterizes I/O efficiency by measuring data transfer rates between the host CPU and accelerator HBM over the PCIe interconnect. This includes Host-to-Device (H2D) and Device-to-Host (D2H) transfers, measured in GB/s.
• Ragged-Paged Attention (RPAv3) : Focuses on benchmarks tailored for transformer-based models, specifically measuring throughput for primitives like Attention (tokamax splash) and Batch Matrix Multiplication (BMM). These help predict end-to-end inference performance and identify bottlenecks in Time-to-First-Token (TTFT) latency.

Microbenchmarking TPU modules

The Network, Compute, HBM, and Host Transfer microbenchmarks in our suite test isolated components, providing a granular assessment of theoretical performance. This detailed component analysis is crucial for debugging and optimizing system performance across various dimensions. For instance, it can help to identify whether a device is achieving its specifications or if specific architectural bottlenecks, such as VMEM locality or interconnect efficiency, are limiting the scale and speed of workloads.

Microbenchmark suite

The core tool for this kind of empirical evaluation is our accelerator-microbenchmarks GitHub repository. This suite provides a structured approach to hardware validation through several key directories:

• src/ : Contains the fundamental implementations and source code for the microbenchmarks

• configs/ : Manages the tuning parameters for each test, such as data buffer sizes, number of iterations, and supported data types (e.g., bf16, fp8, fp32)

  Run a compute microbenchmark with a sample configuration

  kubectl apply -f Ironwood/guides/collectives/tpu7x-2x2x1-ici-all-gather-microbenchmark.yaml

Shell

By utilizing these recipes, you can relate the high-level metrics discussed previously (like ICI bandwidth or MXU utilization) to actual runtime data. For example, the collective YAMLs in the configs folder allow for direct measurement of the ICI torus efficiency described in the architecture section, helping to ensure the physical interconnect health matches theoretical expectations.

Performance tuning impact

Microbenchmarks take performance optimization from a trial-and-error process into an empirical engineering discipline. By establishing a "Speed-of-Light" (SOL) baseline (the theoretical limit of the hardware modules), they provide a clear target for training and inference efficiency.

The Roofline model: A diagnostic North Star

The core value of microbenchmarks lies in their ability to define the boundaries of the Roofline model. By measuring peak MXU throughput and HBM bandwidth, developers can immediately categorize a workload's bottleneck as:

• Compute-bound : The model is hitting the MXU plateau. Tuning memory access will yield no benefit; optimizations must focus on kernel efficiency or reducing FLOPs.
• Memory-bound : The model is limited by the sloped "eaves" of the roofline. Performance can be improved by optimizing data locality (VMEM) or reducing HBM traffic.
• Network-bound : Stalls in ICI or DCN prevent the chips from reaching their compute potential. This identifies the need for better sharding strategies or communication-overlap techniques.

Hardware-aware model architecture

Microbenchmarks reveal the specific "personality" of the TPU hardware. For example, on the Ironwood (TPU7x) architecture, the 256x256 systolic array imposes a physical constraint on operand shapes. Benchmarking confirms that models using a head_dim of 128 (common in older Llama variants) achieve less optimal MXU utilization. This insight drives researchers to co-design frontier models with dimensions aligned to 256-byte boundaries to maximize hardware efficiency.

Software levers and optimization strategies

Once a bottleneck is identified, microbenchmarks can help you select the right software lever:

• Kernel selection : You can compare standard JAX operations against specialized Pallas kernels or Tokamax Splash Attention. Microbenchmarks provide the ground truth for which implementation minimizes the backward pass overhead.
• Sharding and mesh tuning : By measuring collective performance across different topologies (e.g., 2x2x2 vs 4x2x1), you can optimize Fully Sharded Data Parallel (FSDP) parameters to ensure that communication time is fully hidden behind computing, i.e., tasks complete fast enough so that the TPU does not need to stop and wait for a response.
• Rematerialization : Microbenchmarks help calculate the exact trade-off between compute cycles and memory pressure, guiding decisions on when to store calculated data in memory vs. “rematerialize” (recalculate) it, to avoid HBM overflows.

Predictive optimization

Finally, microbenchmark data allows for predictive modeling of large-scale deployments. Instead of running expensive tests on full TPU slices, you can use results to build analytical models that accurately predict training-throughput, TTFT, Time-Per-Output-Token (TPOT), and MFU before they scale up.

Case study: Ironwood (TPU 7x) performance analysis

In one example, we analyzed the performance tuning of a 110B Mixture-of-Experts (MoE) training workload deployed on a 4x4x4 TPU 7x configuration, leveraging microbenchmarks to establish hardware baselines and drive optimization for complex, sparse workloads. Here is what we found:

• Roofline analysis : Diagnostic profiling of the training workload revealed a clear division in performance bottlenecks. The forward-pass and dense-core operations were primarily compute-bound, reaching 1.85 PFLOPS. Conversely, the routing and attention primitives in the MoE architecture remained memory-bound, operating at 30–60% of the SOL due to HBM saturation, where sparse activation patterns become heavily IO-bound during expert dispatch, sending incoming requests to sub-networks.
• Correlation with microbenchmarks : Microbenchmark results directly guided interventions that reduced the training step time by 21.2%. ICI collective benchmarks identified communication stalls resolved by SparseCore collective offloading, while HBM bandwidth tests identified memory-bound Attention primitives, prompting the adoption of Tokamax Splash Attention.

Bringing discipline to TPU performance optimization

In conclusion, the microbenchmarks suite provides the structured validation necessary to bridge the gap between theoretical specifications and actual runtime performance, empowering you to fully maximize the potential of TPU architectures across large-scale deployments.

To learn more, check out the following resources:

1. Ironwood: The first Google TPU for the age of inference
2. Optimizing frontier model training on TPU v7x (Ironwood)
3. AI accelerator performance and benchmarking
4. A developer’s guide to training with Ironwood TPUs