Pallas TPU# TPU 特定文档。 指南 使用 Pallas 编写 TPU 内核 什么是 TPU？ 值得注意的属性和限制 支持的操作 TPU 流水线 TPU 及其内存空间 TPU 特有的流水线特性 矩阵乘法 背景 您的第一个矩阵乘法内核 矩阵乘法性能 流水线内核的性能 矩阵乘法的模板化 结论 标量预取和块稀疏计算 带标量预取的动态块索引 示例：带标量预取的块动态切片 稀疏内核：表示稀疏数据 示例：稀疏 @ 稠密矩阵乘法 稠密数据上的稀疏访问模式 示例：带块稀疏输出掩码的稠密 @ 稠密矩阵乘法 Pallas 在 TPU 上的分布式计算 TPU 拓扑 远程直接内存访问 (RDMA) 模型 高级技术 最后说明
