Google 替 XProf 推出核心分析套件：支援 TPU v7 暫存器取樣，矩陣乘法示範核心耗時縮短 30%

據報導，Google 為其開源 TPU 效能分析器 XProf 新增核心效能分析套件，讓開發者得以檢視自訂 Pallas 核心的時鐘週期級細節，並支援在 TPU v7（Ironwood）上直接取樣硬體效能計數器。Google AI Infra 團隊的 Yogesh SY 指出，使用 Pallas、Mosaic 或 Triton 建立的核心會跳過標準 XLA Pass，使得編譯期靜態成本模型（如 optimal FLOPs 或吞吐效率）產生偏差，靜態工具常誤將正在等待 HBM 而閒置的 MXU 判定為充分利用。 該分析套件分為三個運作層級：在編譯器檢查層，開發者可啟用 --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 週期取樣下限，新增外部事件觸發模式（external event trigger mode），可擷取自訂呼叫範圍進出等邊界觸發器，實現次微秒級採集。開發者可在最多 4 個 SparseCore 上，為每個核心配置最多 28 個計數器（形成 4×28 矩陣），並透過 jax.profiler.ProfileOptions 啟用。 在 Google 針對 TPU v7 進行的分塊矩陣乘法示範案例中，分析流程透過 sync_wait 計數器識別出記憶體停頓，在加入三重緩衝以重疊 HBM 載入與 MXU 計算後，核心執行時間從 125.5µs 縮短至 88µs，降幅約 30%。報導強調該數據來自單一示範核心而非廣泛基準測試，主要用於展示分析流程。 套件內的 Perf Counters View 以表格形式列出超過 16,000 個原始計數器，軌道高度代表時間區間內的原始計數器最大值（例如 2.0 GHz 核心在 500ns 窗口內增加 100 週期，即代表 10% 利用率），建議開發者以硬體暫存器數值取代 XLA 靜態估算。XProf 隸屬於 OpenXLA 專案並整合 JAX 分析，但報導未說明該套件的正式發布版本或釋出時程，且計數器取樣文件目前僅針對 TPU v7，早期 TPU 架構未必支援相同涵蓋範圍。