Pallas: Mosaic GPU

针对 Mosaic GPU 后端的特定文档。

参考文档

• 使用 Pallas 编写 Mosaic GPU 内核
  • 什么是 GPU？
  • 数组布局与内存引用转换
  • MMA (TensorCore)
  • 使用 core_map
  • 同步结构与原语
  • 集群启动控制
  • 异步拷贝
  • 内联 Mosaic GPU
  • 编译器参数
  • 调试
  • 从 PyTorch 调用内核
• Mosaic GPU 流水线
  • Mosaic GPU 流水线
  • GPU 内存空间
  • 示例：Hopper GPU 上的矩阵乘法内核
  • 线程束专用化 (Warp Specialization)
  • 示例：使用线程束专用化的矩阵乘法
• 为 Blackwell 编写高性能矩阵乘法内核
  • 0. 基础内核
  • 1. 线程束专用化
  • 2. 分块尾声 (Tiled epilogue)
  • 3. 集体 (2CTA) MMA
  • 4. 常驻内核 (Persistent kernel)
  • 5. 专用尾声线程束组
  • 6. 网格分块
  • 最终内核
• 集体矩阵乘法
  • 算法概述：环形全聚合 (Ring All-Gather)
  • 用于设备间通信的 Pallas 原语
  • 使用 Pallas 实现
  • 将内核与 JAX 集成