Google Adds Cycle-Level Kernel Profiling to XProf

Google has introduced an enhanced Kernel Profiling feature to XProf, an open-source profiler tailored for Tensor Processing Unit (TPU) workloads. This update allows developers to examine cycle-level information within custom Pallas kernels, a significant advancement from previous iterations where these kernels were represented as singular opaque blocks in trace captures.

During runtime, XProf samples hardware performance counters on TPU v7 (Ironwood) models. A real-world example from Google's AI Infra team demonstrated this by identifying and subsequently addressing a memory stall in a tiled matrix-multiply (matmul) scenario. By implementing triple buffering, they managed to decrease kernel execution time from 125.5 microseconds to 88 microseconds, marking a substantial 30% improvement.

The enhancement hinges on how profilers process custom compilation paths. Kernels constructed using Pallas, Mosaic, or Triton bypass standard XLA (Accelerated Linear Algebra) passes, which can skew compile-time static cost models. Consequently, quantitative metrics like optimal FLOPs (floating-point operations per second) and goodput efficiency might prove unreliable without proper adjustments.

The new Kernel Profiling suite operates at three distinct levels. Developers activate it by incorporating two compiler inspection flags: --xla_enable_custom_call_region_trace=true and --xla_xprof_register_llo_debug_info=true. Following this activation, the Graph Viewer presents a Custom Call Text panel, illustrating the lowered Machine Learning Inferencing Representation (MLIR) for each custom call.

This feature enables engineers to verify the fusion of operations and the structuring of memory tiles as per the intended design.

The Trace Viewer exhibits Low-Level Operations (LLO) bundle data, encompassing machine instructions per clock cycle. It tracks time-aligned events for various components, including the MXU (Matrix Multiply Unit), scalar and vector ALUs (Arithmetic Logic Units), vector fills, loads, spills, stores, and the cross-lane unit (XLU). For real-time insights, XProf captures hardware counters at intervals of at least 1 microsecond, with an optional external event-triggered mode offering sub-microsecond capture precision.

Engineers can configure up to 28 counters per core, dispersed across up to four SparseCores, totaling a 4x28 matrix. This collection process is facilitated through jax.profiler.ProfileOptions, leveraging tpu_enable_periodic_counter_sampling and tpu_tc_perf_counter_sampling_options with is_external_trigger:true. In periodic mode, the interval_us flag supplants the trigger flag.

As demonstrated in a matmul case study, the memory-bound variant exhibited significant surge in sync_wait counters upon HBM (High Bandwidth Memory) loads. The application of triple buffering to overlap HBM loads with MXU computation resulted in the reduction of these events, thereby boosting performance from 125.5 microseconds to the reduced 88 microseconds.

While this improvement is based on a single Google-authored kernel demo, it outlines the workflow and potential gains that could be extrapolated across other custom kernels. The announcement also introduces a "hierarchy of trust" framework for interpreting metrics. Hardware registers, such as HBM utilization and TPO (Tensor Processing Optimization) metrics, are considered the gold standard for custom kernel performance, whereas XLA cost-model estimates should be treated with caution.

A new Perf Counters View offers a tabular listing of over 16,000 raw counters, with trace track height representing the maximum raw counter value in a given interval, rather than a normalized percentage.