Google 为 XProf 添加了一个内核性能分析套件。XProf 是其面向 TPU 工作负载的开源性能分析器。现在，开发者可以查看自定义 Pallas 内核中的周期级细节。此前，这些内核在跟踪记录中显示为单个不透明块。在 TPU v7（Ironwood）上，XProf 会在运行时对硬件性能计数器进行采样。在一个分块矩阵乘法示例中，Google 发现了内存停顿问题。通过添加三重缓冲，他们将内核执行时间从 125.5µs 缩短至 88µs，降幅约为 30%。
据 Google AI Infra 团队的 Yogesh SY 介绍，这一差距源于性能分析器处理自定义编译路径的方式。使用 Pallas、Mosaic 或 Triton 创建的内核会跳过标准 XLA Pass。这可能使编译期静态成本模型产生偏差。因此，“optimal FLOPs”和有效吞吐效率等指标可能不准确，甚至完全无法工作。静态工具可能会将一个 MXU 指令块标记为已充分利用，但该单元实际上正在空闲等待 HBM，因为静态分析忽略了时间因素。
该套件分三个层级工作。开发者可以传入以下标志进行编译器检查：--xla_enable_custom_call_region_trace=true 和 --xla_xprof_register_llo_debug_info=true。之后，Graph Viewer 会显示一个“Custom Call Text”面板，其中展示每个自定义调用降低后的 MLIR。
这让工程师能够检查操作是否按预期融合，以及内存分块的结构是否符合设计。Trace Viewer 会显示用于静态执行分析的低级操作（LLO）包数据。这些数据包括每个时钟周期执行的机器指令。它为 MXU、标量与向量 ALU、向量填充、加载、溢出、存储以及跨通道单元（XLU）提供时间对齐的轨道。对于运行时遥测，XProf 会定期对硬件计数器进行采样，其分辨率下限为 1µs，这一限制来自主机级定时器。
一种新的外部事件触发模式取消了这一限制。采样器会捕获 TPU 跟踪指令和边界触发器，包括自定义调用作用域的进入和退出。这可以实现亚微秒级捕获，并提高归因准确性。开发者可以在最多四个 SparseCore 上，为每个核心配置最多 28 个计数器，形成一个 4×28 矩阵。采集通过 jax.profiler.ProfileOptions 启用，使用 tpu_enable_periodic_counter_sampling 和 tpu_tc_perf_counter_sampling_options，并设置 is_external_trigger:true。对于周期采样模式，则使用 interval_us 取代触发器标志。
在矩阵乘法案例研究中，一个受内存限制的变体在 sync_wait 计数器中出现了大量峰值。通过三重缓冲让 HBM 加载与 MXU 计算重叠，减少了这些事件，并将执行时间从 125.5µs 缩短至 88µs。该数据来自 Google 编写的单个演示内核，而非广泛的基准测试，因此它展示的是分析流程，而不是其他场景中可以预期的性能收益。
该文章还提出了指标的“信任层级”。直接从硬件寄存器读取的值，例如 HBM 利用率和 TPO 指标，可视为自定义内核的事实依据，而 XLA 成本模型的估算结果则需要谨慎对待。新的 Perf Counters View 以表格形式列出了超过 16000 个原始计数器。跟踪轨道的高度反映某个时间区间内原始计数器的最大值，而不是经过归一化的百分比。例如，在主频约为 2.0 GHz 的核心上，500ns 时间窗口内增加 100 个周期，相当于该单元的利用率为 10%。
XProf 是 OpenXLA 项目的一部分，并与 JAX 性能分析集成。该文章没有说明内核性能分析套件的发布状态或版本。
计数器采样功能的文档面向 TPU v7（Ironwood），因此使用早期 TPU 世代的团队不应假定它们具有相同的覆盖范围。任何需要调优 Pallas 内核的开发者，都可以从 XProf 仓库中的 Pallas Matmul with Perf Counters Notebook 和 OpenXLA 内核性能分析说明入手。有限的 4×28 计数器预算要求团队针对每次调查选择相应的计数器。该设置旨在保障工作负载性能。工程师也不应再将 XLA 得出的效率数据视为自定义内核的可靠指标，而应改用寄存器级计数器作为优化依据。