Google has added a Kernel Profiling suite to XProf. This is its open-source profiler for TPU workloads. Now, developers can see cycle-level details in custom Pallas kernels. Before, these kernels appeared as single opaque blocks in trace captures. On TPU v7 (Ironwood), XProf samples hardware performance counters during runtime. In a tiled matmul example, Google found a memory stall. They cut kernel time from 125.5µs to 88µs, about 30%, by adding triple buffering.
According to Yogesh SY of Google's AI Infra team, the gap comes from how profilers treat custom compilation paths. Kernels made with Pallas, Mosaic, or Triton skip standard XLA passes. This can distort the compile-time static cost models. As a result, metrics like "optimal FLOPs" and goodput efficiency may be inaccurate or not work at all. A static tool can flag an MXU instruction block as fully utilized while the unit sits idle waiting on HBM, because static analysis ignores time.
The suite works at three levels. Developers pass these flags for compiler inspection: --xla_enable_custom_call_region_trace=true and --xla_xprof_register_llo_debug_info=true. After this, the Graph Viewer displays a "Custom Call Text" panel. It shows the lowered MLIR for each custom call.
This lets engineers check whether operations are fused and memory tiles are structured as intended. The Trace Viewer shows Low-Level Operations (LLO) bundle data for static execution analysis. This data includes machine instructions per clock cycle. It presents time-aligned tracks for the MXU, scalar and vector ALUs, vector fills, loads, spills, stores, and the cross-lane unit (XLU). For runtime telemetry, XProf samples hardware counters periodically, with a 1µs resolution floor driven by a host-level timer.
A new external event-triggered mode removes that floor. The sampler captures TPU trace instructions and boundary triggers. These include the entry and exit of custom call scopes. This enables sub-microsecond capture and improved attribution. Developers can configure up to 28 counters per core across up to four SparseCores, a 4x28 matrix. Collection is enabled through jax.profiler.ProfileOptions, using tpu_enable_periodic_counter_sampling and tpu_tc_perf_counter_sampling_options with is_external_trigger:true. For periodic mode, interval_us replaces the trigger flag.
In the matmul case study, a memory-bound variant showed large spikes in sync_wait counters. Overlapping HBM loads with MXU compute through triple buffering reduced those events and produced the 125.5µs to 88µs improvement. The figure comes from a single Google-authored demo kernel, not a broad benchmark, so it shows the workflow rather than the gains to expect elsewhere.
The post also sets out a "hierarchy of trust" for metrics. Values read directly from hardware registers, such as HBM utilization and TPO metrics, count as ground truth for custom kernels, while XLA cost-model estimates need caution. A new Perf Counters View lists over 16,000 raw counters in tabular form. Trace track height reflects the maximum raw counter value in an interval, not a normalized percentage. For example, an increment of 100 cycles over a 500ns window on a roughly 2.0 GHz core equals 10% utilization of that unit.
XProf is part of the OpenXLA project and integrates with JAX profiling. The post does not state a release status or version for the Kernel Profiling suite.
The counter sampling is documented for TPU v7 (Ironwood), so teams on earlier TPU generations should not assume the same coverage. Anyone tuning Pallas kernels can start with the Pallas Matmul with Perf Counters notebook in the XProf repository and the OpenXLA kernel profiling instructions. The limited 4x28 counter budget requires teams to pick counters for each investigation. The setup aims to safeguard workload performance. Engineers should also stop treating XLA-derived efficiency numbers as reliable for custom kernels and anchor optimization on register-level counters instead.