Google XProf 新增 Pallas 内核周期级分析工具并警示静态模型偏差

Google 为面向 TPU 工作负载的开源性能分析器 XProf 添加了针对自定义 Pallas 内核的性能分析套件，使开发者能够查看周期级细节而非仅显示不透明块。该工具在 TPU v7（Ironwood）上通过硬件计数器采样发现内存停顿问题，并在分块矩阵乘法示例中将执行时间从 125.5µs 缩短至 88µs。文章强调由于 Pallas、Mosaic 或 Triton 创建的内核跳过标准 XLA Pass，导致编译期静态成本模型产生偏差，使得 optimal FLOPs 等指标可能完全无法工作。

新套件通过三个层级工作：编译器检查标志启用后，Graph Viewer 展示降低后的 MLIR 以确认操作融合与内存分块结构；Trace Viewer 提供时间对齐的低级操作包数据，覆盖 MXU、ALU 及各类加载存储单元；运行时遥测支持亚微秒级捕获，允许在最多四个 SparseCore 上配置 4×28 个计数器矩阵。文章提出信任层级概念，指出直接从硬件寄存器读取的 HBM 利用率等值应视为事实依据，而 XLA 成本模型的估算结果需谨慎对待，新的 Perf Counters View 列出了超过 16000 个原始计数器且轨道高度反映最大值而非归一化百分比。

该功能目前文档面向 TPU v7（Ironwood），早期世代团队不应假定具有相同覆盖范围。文章未说明内核性能分析套件的发布状态或具体版本，相关设置需通过 jax.profiler.ProfileOptions 启用。对于需要调优 Pallas 内核的开发者，建议放弃依赖 XLA 得出的效率数据，转而使用寄存器级计数器作为优化依据，并可参考 OpenXLA 仓库中的演示 Notebook 和说明文档进行实践。