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.
The microbenchmark suite provides standardized performance evaluations across five core areas:
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.
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:
kubectl apply -f Ironwood/guides/collectives/tpu7x-2x2x1-ici-all-gather-microbenchmark.yaml
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.
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 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:
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.
Once a bottleneck is identified, microbenchmarks can help you select the right software lever:
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.
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:
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: