Pallas TPU# TPU 特定文档。 指南 使用 Pallas 编写 TPU 内核 什么是 TPU？ 值得注意的特性和限制 支持的操作 TPU 流水线 TPU 及其内存空间 TPU 特定的流水线（Pipelining）功能 矩阵乘法 背景 您的第一个矩阵乘法内核 矩阵乘法性能 流水线化内核的性能 矩阵乘法模板化 结论 标量预取和块稀疏计算 带标量预取的动态块索引 示例：带标量预取的块动态切片 稀疏内核：表示稀疏数据 示例：稀疏 @ 稠密矩阵乘法 稠密数据上的稀疏访问模式 示例：带块稀疏输出掩码的稠密 @ 稠密矩阵乘法 Pallas 在 TPU 上的分布式计算 TPU 拓扑结构 远程直接内存访问 (RDMA) 模型 进阶技巧 最终说明 Pallas Core 特定编程 环境设置 简单的单核内核 使用 core_map 进行流水线化 标量预取 在 SparseCore 上进行映射 SparseCore 内核编写 硬件概述 操作与工作负载 Express SparseCore 硬件 基础 SparseCore 内核 SparseCore 内核中的流水线化 TensorCore 与 SparseCore 的重叠执行 Gather 和 Scatter 操作 与 TensorCore 的性能对比 伪随机数生成 使用 jax.random API 使用硬件 PRNG 块不变采样
