Pallas TPU# TPU 專用文件。 指南 使用 Pallas 撰寫 TPU 核心 什麼是 TPU? 值得注意的特性和限制 支援的運算 管線化 TPU 及其記憶體空間 使用 VMEM/SMEM 的限制 基礎知識:管線化 Pallas 中的管線化 處理縮減 Megacore 配置中的 TPU 結論 矩陣乘法 背景 您的第一個矩陣乘法核心 矩陣乘法效能 管線化核心的效能 矩陣乘法範本化 結論 純量預取和區塊稀疏計算 使用純量預取的動態區塊索引 範例:使用純量預取的區塊動態切片 稀疏核心:表示稀疏資料 範例:稀疏 @ 稠密矩陣乘法 稠密資料上的稀疏存取模式 範例:使用區塊稀疏輸出遮罩的稠密 @ 稠密矩陣乘法 Pallas 中用於 TPU 的分散式運算 TPU 拓撲 遠端直接記憶體存取 (RDMA) 模型 進階技術 最終筆記