How to use Google microbenchmarks for evaluating TPU performance

Summary

Evaluating TPU performance requires measurements beyond product specifications, because real-world behavior depends on architecture and workload constraints. The suite measures network collectives, compute, HBM bandwidth, host transfers, and Ragged-Paged Attention primitives, using metrics such as GB/s, latency, TFLOPs, MFU, and TTFT-related throughput. Its src and configs directories provide implementations and tunable parameters, while sample Kubernetes recipes connect measurements to TPU architecture and runtime behavior. Roofline analysis uses peak MXU throughput and HBM bandwidth to distinguish compute-, memory-, and network-bound workloads and guide kernel, sharding, or rematerialization choices. In a 110B MoE workload on a 4x4x4 TPU 7x configuration, microbenchmark-guided changes addressed ICI stalls and memory-bound attention, reducing training step time by 21.2%.

Takeaways

1. 1

   Network tests cover all-gather, all-reduce, reduce-scatter, and all-to-all, with throughput and latency used to characterize ICI collective performance.

2. 2

   The configs directory exposes buffer sizes, iteration counts, and supported data types including bf16, fp8, and fp32, allowing benchmark runs to be tuned.

3. 3

   In the case study, SparseCore collective offloading addressed communication stalls, while Tokamax Splash Attention targeted HBM-bound attention primitives.