Google has extended XProf, its open-source profiler for TPU workloads, with cycle-level kernel profiling that exposes the execution details of custom Pallas kernels. Before this addition, custom kernels appeared as opaque blocks in trace captures; now developers can see hardware performance counters sampled at microsecond resolution and, in external-trigger mode, at sub-microsecond granularity.
The profiling suite operates at three levels of abstraction. To access the Graph Viewer's lowered MLIR for each custom call, developers must first pass compiler flags --xla_enable_custom_call_region_trace=true and --xla_xprof_register_llo_debug_info=true, after which a "Custom Call Text" panel displays the lowered MLIR, letting engineers verify whether operations are fused and memory tiles are structured as intended. The Trace Viewer shows Low-Level Operations (LLO) bundle data with machine instructions per clock cycle, presenting time-aligned tracks for the MXU, scalar and vector ALUs, vector fills, loads, spills, stores, and the cross-lane unit. For runtime telemetry, XProf samples hardware counters periodically with a 1µs resolution floor in standard mode, or removes that floor entirely in external-trigger mode by capturing TPU trace instructions and boundary triggers at custom call entry and exit.
Counter collection is enabled through jax.profiler.ProfileOptions using tpu_enable_periodic_counter_sampling and tpu_tc_perf_counter_sampling_options. For external-trigger mode, developers set is_external_trigger:true for sub-microsecond capture; for periodic mode, interval_us replaces the trigger flag. Developers can configure up to 28 counters per core across up to four SparseCores. A new Perf Counters View lists over 16,000 raw counters in tabular form, with track height reflecting the maximum raw counter value in an interval rather than a normalized percentage.
Google's case study on a tiled matmul kernel demonstrates the workflow. The memory-bound variant showed large spikes in sync_wait counters. By overlapping HBM loads with MXU compute through triple buffering, the team reduced those events and cut kernel time from 125.5µs to 88µs, approximately 30% improvement. The post notes this figure comes from a single Google-authored demo kernel, not a broad benchmark, so it illustrates the optimization process rather than typical gains elsewhere.
The addition addresses a fundamental gap in profiling custom kernels. According to Yogesh SY of Google's AI Infra team, kernels made with Pallas, Mosaic, or Triton skip standard XLA passes, which can distort compile-time static cost models. Metrics like "optimal FLOPs" and goodput efficiency may be inaccurate or absent entirely. 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 post establishes 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 require caution.
The profiling suite 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. Counter sampling is documented for TPU v7 (Ironwood), so teams on earlier TPU generations should not assume the same coverage. The limited 4x28 counter budget requires teams to select counters for each investigation, and engineers should anchor optimization on register-level counters rather than treating XLA-derived efficiency numbers as reliable for custom kernels.
For teams tuning Pallas kernels on TPU v7, the workflow starts with the Pallas Matmul with Perf Counters notebook in the XProf repository and the OpenXLA kernel profiling instructions. The takeaway: cycle-level visibility into custom kernel execution is now available on Ironwood, but only if you explicitly enable it and know which counters to watch.