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