面向機器學習系統的現代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調度等。