前些日子,卡内基梅隆大学助理教授陈天奇推出了一本免费的在线书《Modern GPU Programming For MLSys(面向机器学习系统的现代 GPU 编程)》。这本书源自他在 CMU 讲授的「机器学习系统」课程中的 GPU 编程模块,核心内容是当下大模型训练和推理中的关键难点——高性能 GPU kernel 的编写技巧,整体围绕 Blackwell 架构,将矩阵乘法(GEMM)和 FlashAttention 作为全书实践范例。书籍链接为:https://mlc.ai/modern-gpu-programming-for-mlsys/
陈天奇的背景 陈天奇现在在卡内基梅隆大学任机器学习系与计算机科学系助理教授,同时是英伟达的杰出工程师职位。他在华盛顿大学完成博士学业,师从 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 编构实践同样重要。本书正是为此空缺而著。
全书分为四大部分: 第一部分探讨 GPU 的基本运作原理,覆盖 GPU 执行模型、kernel 性能评估法,以及影响全书的重点概念,诸如数据布局(data layout)、异步数据传输(TMA)、张量核心(Tensor Core / tcgen05)、特殊内存(TMEM)、异步同步(mbarrier),为后续篇章建立硬件理解基础。
第二部分介绍 TIRx,这是本书用来编写案例的 Python DSL,其设计理念是"贴近硬件",让读者既可精确掌控底层运作,又能通过运行代码学习,而非仅停留在静态伪代码的层面。第三部分用矩阵乘法完整铺陈性能优化的演进路线:从单 tile 顺序执行,到 K 维度循环累加、多 CTA 空间分块,再到引入 TMA 异步载入、软件流水线(pipeline)、持久化 kernel 与 tile 调度器,最终借助 warp 专精化(warp specialization)和双 CTA 集群(2-CTA cluster),把性能推向当前架构下的较高水准。





