Pallas:Mosaic GPU# Mosaic GPU 后端的特定文档。 参考文档 使用 Pallas 编写 Mosaic GPU 内核 什么是 GPU? 数组布局和内存引用转换 MMA (TensorCore) 使用 core_map 同步结构和原语 集群启动控制 异步拷贝 内联 Mosaic GPU 编译器参数 调试 从 PyTorch 调用内核 Mosaic GPU 流水线 Mosaic GPU 的流水线 GPU 内存空间 示例:Hopper GPU 上的矩阵乘法内核 Warp 特化 示例:带有 Warp 特化的矩阵乘法 为 Blackwell 编写高性能矩阵乘法内核 0. 基本内核 1. Warp 特化 2. Tiled epilogue 3. Collective (2CTA) MMA 4. Persistent kernel 5. Dedicated epilogue warpgroup 6. Grid tiling 最终内核 集体矩阵乘法 算法概览:Ring All-Gather 用于设备间通信的 Pallas 原语 使用 Pallas 实现 将内核集成到 JAX