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.yamlShell
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 256×256 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:
