Google XProf Adds Cycle-Level Kernel Profiling for TPU Workloads

XProf Enhances TPU Profiling Capabilities

Google has introduced a Kernel Profiling suite to XProf, its open-source profiler designed for TPU workloads. This addition allows developers to gain cycle-level insights into custom Pallas kernels, which previously appeared as opaque blocks in trace captures. The new functionality is particularly relevant for TPU v7 (Ironwood), where XProf samples hardware performance counters during runtime.

Addressing Static Analysis Limitations

According to Yogesh SY of Google's AI Infra team, custom compilation paths used by kernels developed with Pallas, Mosaic, or Triton bypass standard XLA passes. This can lead to distortions in compile-time static cost models, making metrics like "optimal FLOPs" and goodput efficiency inaccurate. The new profiling suite aims to provide a more accurate view of actual runtime performance, as static analysis can incorrectly flag units as utilized when they are idle, waiting on HBM.

Multi-Level Inspection and Telemetry

The suite operates at three levels. Developers can enable compiler inspection flags to view a "Custom Call Text" panel in the Graph Viewer, which displays the lowered MLIR for each custom call. This allows engineers to verify operation fusion and memory tile structuring. The Trace Viewer provides Low-Level Operations (LLO) bundle data for static execution analysis, including machine instructions per clock cycle and time-aligned tracks for various processing units.

For runtime telemetry, XProf periodically samples hardware counters with a 1µs resolution floor. A new external event-triggered mode removes this floor by capturing TPU trace instructions and boundary triggers, such as custom call entry and exit, enabling sub-microsecond capture and improved attribution. Developers can configure up to 28 counters per core across up to four SparseCores.

Practical Application and Optimization

In a case study involving a tiled matmul, Google identified a memory stall using the new profiling tools. By implementing triple buffering, they reduced the kernel time from 125.5µs to 88µs, representing approximately a 30% improvement. The memory-bound variant showed large spikes in sync_wait counters, which were mitigated by overlapping HBM loads with MXU compute through triple buffering.