Google Brings Deep Hardware Insight to TPU Profiling with XProf Kernel Suite

Google has rolled out a kernel profiling suite for XProf, its open-source profiler for TPU workloads. The enhancement gives developers visibility into Pallas, Mosaic, and Triton kernels at cycle-level granularity—a stark upgrade from the past, when custom kernels appeared as undifferentiated blocks in performance traces.

The Problem with Custom Kernels#

Custom compiled kernels bypass standard XLA optimization passes, creating a blind spot for traditional profilers. According to Yogesh SY of Google's AI Infra team, this gap leads to unreliable metrics. A static analysis might report an MXU (matrix unit) as fully utilized when hardware counters reveal it's actually stalled waiting on memory access. The mismatch arises because static cost models ignore timing dynamics.

Three Layers of Profiling#

The new suite operates at three increasingly detailed levels:

• Graph-level: Compiler flags (--xla_enable_custom_call_region_trace=true and --xla_xprof_register_llo_debug_info=true) unlock a "Custom Call Text" panel in the Graph Viewer, showing lowered MLIR code for each custom call. Engineers can verify whether operations are fused and memory tiles are laid out as intended.

• Low-level static: The Trace Viewer displays Low-Level Operations (LLO) bundle data, including machine instructions per clock cycle. Time-aligned tracks show activity across the MXU, scalar and vector ALUs, vector fills, loads, spills, stores, and cross-lane unit.

• Runtime telemetry: XProf samples hardware performance counters with sub-microsecond resolution, driven by an event-triggered capture mode that monitors TPU trace instructions and custom call entry/exit boundaries. Developers can configure up to 28 counters per core across up to four SparseCores.

A Real Example: 30% Speed Gain#

In a tiled matrix-multiplication demo, Google engineers identified a memory stall that was throttling kernel execution. By adding triple buffering to overlap HBM loads with MXU compute, they trimmed kernel time from 125.5µs to 88µs—roughly a 30% improvement. The sync_wait counter revealed large spikes that disappeared once the buffering strategy was in place.

Google cautions that this single demo illustrates the workflow rather than expected gains for all workloads. The company also published a "hierarchy of trust": hardware register values (HBM utilization, TPO metrics) are ground truth for custom kernels, while XLA cost-model estimates should be treated with skepticism.