編輯｜Panda

前些天，CMU 助理教授、TVM/XGBoost/MLC-LLM 的創造者陳天奇發布了一本免費在線書籍《Modern GPU Programming For MLSys（面向機器學習系統的現代 GPU 編程）》。

這本書脫胎于他在 CMU 開設的「機器學習系統」（Machine Learning Systems）課程系列中新增的 GPU 編程專題課，內容聚焦當前大模型訓練和推理中最核心、也最難啃的部分——高性能 GPU kernel 的編寫方法，主線圍繞 Blackwell 架構展開，并以矩陣乘法（GEMM）和 FlashAttention 作為貫穿全書的實戰案例。

• 書籍地址：https://mlc.ai/modern-gpu-programming-for-mlsys/

陳天奇簡介

陳天奇目前是卡內基梅隆大學機器學習系與計算機科學系的助理教授，同時擔任英偉達的杰出工程師（Distinguished Engineer）。他博士畢業于華盛頓大學，師從 Carlos Guestrin，本科和碩士就讀于上海交通大學 ACM 班。在創業方向，他曾是 OctoAI（后被 NVIDIA 收購）的首席技術官。

陳天奇在開源機器學習系統領域有著相當深的積累，是多個被廣泛使用的開源項目的發起人或共同創建者，包括深度學習編譯棧 Apache TVM、樹模型庫 XGBoost、通用大模型部署引擎 MLC-LLM，以及 Apache MXNet（共同創建者）。

他同時擔任 MLSys 學術會議的董事會主席，并曾任 MLSys 2023 的程序委員會主席，在機器學習系統這一交叉領域具有較高的學術與產業影響力。

新書寫了什么

《Modern GPU Programming For MLSys》瞄準的問題很具體：在大模型的訓練和推理鏈路里，端到端速度往往取決于少數幾類 GPU kernel 的實現質量，包括注意力機制 kernel、LLM 的 prefill/decode kernel、低精度的 block-scaled GEMM、融合 MoE 層等。

而隨著 GPU 架構持續演進，新的內存層級、訪問模式和專用執行單元不斷出現，僅靠「優化技巧清單」已經不足以寫出真正高效的 kernel，需要同時具備對硬件的清晰認知和對高性能 kernel 構建方式的實操理解。這正是本書試圖填補的部分。

全書結構分為四大部分：

Part I：理解 GPU。介紹 GPU 的整體執行模型、判斷 kernel 性能優劣的方法，以及數據布局（data layout）、異步數據搬運（TMA）、張量核心（Tensor Core / tcgen05）、特殊內存（TMEM）、異步協調（mbarrier）等貫穿全書的核心概念，為后續內容打下硬件直覺基礎。

Part II：TIRx 概覽。介紹本書用于編寫示例代碼的 Python DSL——TIRx。這套 DSL 的設計目標是"貼近硬件"，讓讀者既能精確控制底層細節，又能通過可運行的代碼學習，而不只是閱讀靜態的偽代碼。

Part III：從基礎 Tiled GEMM 到 SOTA。以矩陣乘法為例，完整展示性能優化的遞進路徑：從單 tile 的順序實現，到 K 維度循環累加、多 CTA 空間分塊，再引入 TMA 異步加載、軟件流水線（pipeline）、持久化 kernel 與 tile 調度器，最終通過 warp 專精化（warp specialization）和雙 CTA 集群（2-CTA cluster）將性能推向當前架構下的較優水平。

Part IV：FlashAttention 4。在前面所有技術積累的基礎上，構建一個完整的注意力機制 kernel，涵蓋雙階段 MMA 計算、softmax 之間的銜接、在線 softmax 重縮放（online-softmax rescaling）、因果掩碼（causal masking）以及分組查詢注意力（GQA）支持等工程細節。

書末附有參考章節，包含 TIRx 的完整語言參考手冊和編譯器內部實現細節，供有意深入研究底層機制的讀者查閱。

陳天奇在發布動態中提到，這套內容源自他們今年在 CMU 機器學習系統課程中新開設的迷你系列課程，覆蓋了數據布局 swizzling、3D TMA 使用方式等具體問題，并直接對標 Blackwell 這一代 GPU 的編程實踐。由于本次課程未能錄制公開課視頻，他表示希望書中的圖文材料和交互式演示能夠彌補這一缺憾。