面向机器学习系统的现代GPU编程

现代GPU编程对于机器学习系统至关重要。在这些系统中，性能往往取决于少数GPU内核的质量。注意力内核、LLM预填充和解码内核、低精度块缩放GEMM、融合MoE层以及其他大型融合内核直接影响训练和推理的端到端速度。然而，要使这些内核快速运行，需要的不仅仅是一系列优化技巧。现代GPU不再是同一旧设计的简单变体。最近的架构引入了更丰富的内存空间、新的访问模式和日益专门的执行单元。为了良好地编程这些GPU，我们既需要清晰的硬件心智模型，也需要理解高性能内核是如何构建的实用知识。本书旨在培养这两方面的能力。

本书遵循简单的递进顺序：首先理解GPU硬件，然后学习我们将使用的编程模型，最后逐步构建最先进的内核。我们的主要目标是Blackwell代，主要运行示例是快速矩阵乘法（GEMM）和FlashAttention。在此过程中，我们还将研究GPU优化的核心要素：数据布局、异步数据移动和异步协调。

本书内容源于卡内基梅隆大学的机器学习系统课程系列。为了便于研究和运行，本书使用TIRx Python DSL逐步构建真实的GPU内核示例。TIRx贴近硬件，使我们能够在通过可运行代码学习的同时，推理低层控制。

第一部分：理解GPU：介绍GPU的整体组织、编写快速内核的通用方法以及关键概念，如数据布局、异步内存操作和协调。建立本书其余部分依赖的硬件直觉。

第二部分：TIRx概述：介绍TIRx的关键元素，作为全书代码示例的基础。

第三部分：GEMM：从分块到SOTA：通过TMA流水线、持久调度、warp specialized和2-CTA集群，完整指导优化分块GEMM。

第四部分：Flash Attention 4：使用第三部分的技术构建完整的注意力内核：两个MMA，中间有softmax、在线softmax重新缩放、因果掩码和GQA。

附录：TIRx语言参考和编译器内部。

本书详细介绍了优化路径，包括9个步骤，从顺序单块GEMM到多消费者warp specialized，以及Flash Attention 4的具体实现，如Tile-Primitive图、warp角色和作用域、两MMA阶段、TMEM布局和重用、屏障连接、流水线结构、重新缩放和写回、因果掩码、GQA支持、Tile调度等。