Pallas TPU# TPU 特定文档。 指南 使用 Pallas 编写 TPU 内核 什么是 TPU? 值得注意的属性和限制 支持的操作 流水线 TPU 及其内存空间 使用 VMEM/SMEM 的约束 入门:流水线 Pallas 中的流水线 处理约简 Megacore 配置中的 TPU 结论 矩阵乘法 背景 您的第一个矩阵乘法内核 矩阵乘法性能 流水线内核的性能 矩阵乘法的模板 结论 TPU 的 Pallas 分布式计算 TPU 拓扑 远程直接内存访问 (RDMA) 模型 高级技术 最终说明