读书笔记-《面向机器学习系统的现代 GPU 编程》

循序渐进式讨论 Blackwell GPU 编程，建立初步宏观理解。

7 月偶然看到陈天奇大佬开源的新书《面向机器学习系统的现代 GPU 编程》，基于 CMU 的课程 15-442/15-642: Machine Learning Systems。深入浅出讲解了 CUDA 新一代 Blackwell 架构 GPU 上的特点，并从朴素矩阵运算 GEMM 一直推演到几乎极限 SOTA。

整本书行文流畅，视野宏观，加上可交互式动画，非常值得推荐仔细阅读（但是国内很难租到 B200 的 GPU，代码跑起来很麻烦……）。这里笔者忙碌工作中进行学习记录，非算法 AI 背景尽可能分析，避免低级错误。

让我们带着好奇与问题开始吧。

带着问题看书

在翻开这本书之前，先抛出几个问题。它们不急于回答，但会贯穿整篇笔记。本书与本文都不会赘述 transformer 架构和最基本的 GPU/CUDA 编程概念。本文只会在必要比对的地方抛出并深入核心的概念，解释清楚「是什么」「为什么」「怎么做」三步宏观理念，帮助笔者和读者跳出细节，避免在阅读时迷失方向。

  1. LLM 在算什么？为什么矩阵乘法是核心计算？

  2. GPU 为什么适合做矩阵乘法？从 Hopper 到 Blackwell，硬件到底变了什么？

  3. 怎么判断一个矩阵运算算得快还是慢？瓶颈在哪里？

  4. 怎么让它更快？优化不是堆技巧，而是一步一步推过去的——这条路径是什么样的？

  5. 业界在做什么？这本书覆盖了什么，没覆盖什么？

认识对象：LLM 在算什么，GPU 凭什么能算

LLM 推理本质上是一张有向无环图

暂时先不讨论 LLM 的训练，以及分布式情况（笔者也暂时不太懂），让我们看下本书讨论的重点单卡上 LLM 推理，从基本计算原理和公式来看：

一次 LLM 推理的前向传播，本质上是一张有向无环图（DAG）。图中的节点分为两类：

  • 线性运算：矩阵乘法。Q、K、V 投影，Attention 中的 $S = QK^T$ 和 $O = PV$，FFN 中的两层 Linear。它们占据了绝大部分计算量。
  • 非线性运算：softmax、layer norm、激活函数（GeLU / SiLU）。计算量不大，但 softmax 涉及跨行求 max 和 sum，对访存和数值精度有特殊要求。

无论从理论 FLOPs 还是实际延迟统计来看，矩阵乘法都是 LLM 推理的绝对瓶颈。所以这本书的核心命题也很直接：怎么把矩阵乘法做快？

GPU 的物理构成：存储与计算

要理解 GPU 怎么加速矩阵乘法，需要先弄清楚两件事：数据放在哪里（存储），数据被谁处理（计算）。本书没有说明 GPU 的一些基本概念，以及 Hopper 架构到 Blackwell 的主要区别，这里从存储和计算两方面补充下必要知识（这里也是笔者比较缺少硬件基本，简单补充点背景）。了解这些基础概念也是为后续我们讨论矩阵运算优化指明方向，当然如果你已经了解了可以跳过。

存储层次

GPU 的存储是严格分层的。离计算单元越近，速度越快，容量越小。从远到近：

HBM（High Bandwidth Memory）。物理介质是 DRAM，通过硅中介层（silicon interposer）以 3D 堆叠方式焊接在 GPU 芯片旁边。容量大（数十 GB），但物理距离决定了延迟高、带宽有限。所有输入数据和最终结果都必须经过 HBM。程序员通过 cudaMalloc 分配的全局内存（GMEM）就是 HBM 上的一块区域。（这也是存储巨头们暴涨的核心业务）

L2 Cache。物理介质是 SRAM，在 GPU 芯片内部，所有 SM 共享。作用是缓存 HBM 的热数据，减少 HBM 访问次数。（硬件自动管理，程序员不可控。和 CPU 的 L2 一样，能做的只有尽量提高数据在 L2 的命中率。）

以上两级是芯片外（off-chip）或芯片级共享的。

SM（Streaming Multiprocessor）。GPU 的基本计算单元，一块 Blackwell GPU 有上百个 SM。每个 SM 内部有自己的寄存器、SMEM、TMEM、Tensor Core、CUDA Core 和 TMA，相当于 GPU 内部的一个独立”小处理器”。程序以 CTA（Cooperative Thread Array，即 CUDA 的线程块 block）为单位被调度到 SM 上执行，同一 CTA 内的线程共享这个 SM 的 SMEM、可以互相同步。下面进入 SM 内部：

SMEM（Shared Memory）。物理介质是 SRAM，位于每个 SM 内部。容量极小（Blackwell 上约 228KB per SM），但带宽远高于 HBM——因为数据不需要离开 SM。关键是：SMEM 完全由程序员显式管理，什么时候搬入、什么时候搬出，都是代码决定的。这是 GPU 优化中程序员最重要的”可控变量”。

TMEM（Tensor Memory）。物理介质也是 SRAM，但比 SMEM 更靠近 Tensor Core，专为 tcgen05.mma 指令设计 D = A * B + C 的矩阵累加乘法运算数据存储内存。Blackwell 新增。容量比 SMEM 更小，但带宽更高。程序员通过 MMA 指令间接使用。

Register（寄存器）。物理介质是 SM 内部的触发器（flip-flop），每个线程独享。速度最快、容量最小。编译器自动分配，程序员无法显式控制——但寄存器压力（一个线程用太多变量导致部分变量溢出到 SMEM）是实际优化中经常遇到的硬限制。

计算单元

Tensor Core。专为 $D = A \times B + C$ 设计的固定功能硬件。不是通用 ALU，只做矩阵乘加，但一个时钟周期能算一个完整的小矩阵乘法。这是 GPU 矩阵性能的核心来源。

CUDA Core。通用计算单元，处理非矩阵运算（如 softmax 中的加减乘除、地址计算等）。

TMA（Tensor Memory Access）。独立的硬件数据搬运引擎。在 HBM 和 SMEM 之间异步搬数据，不占用 CUDA Core 或 Tensor Core。相当于 GPU 上的 DMA 控制器。

核心矛盾

把存储和计算放在一起看，矛盾就出来了：Tensor Core 算得极快，但 HBM 喂数据的速度远远跟不上。计算吞吐远超内存带宽。

几乎所有 GPU 优化的本质都是在解决这个矛盾——让数据尽量留在离计算近的地方（tiling、data layout），让搬运和计算同时进行（TMA pipelining、warp specialization）。

Hopper 到 Blackwell：架构演进的趋势

特性 Hopper (SM90, H100) Blackwell (SM100, B100/B200)
Tensor Core 指令 wgmma tcgen05.mma
片上存储 SMEM SMEM + TMEM
同步原语 mbarrier mbarrier（增强）
线程组织 Warp, CTA Warp, Warpgroup, CTA
Cluster 调度 固定 CLC（可编程）
数据搬运 TMA TMA（增强）

趋势很明确：每一代都在重新划分计算和数据的边界。计算更专用（tcgen05.mma），存储层次更细（TMEM），数据搬运更异步（TMA 增强），线程组织更灵活（Warpgroup、CLC）。每一步都在缓解同一个核心矛盾：计算太快，数据跟不上。

认识标准：什么叫快

上一节的核心矛盾很明确：Tensor Core 算得快，HBM 喂得慢。这就引出一个问题——一段代码跑起来，到底卡在计算还是卡在带宽？Roofline 就是用来回答这个问题的。

Roofline：找到天花板

Roofline 的核心是”一算一搬”两个上限：

  • 算（计算峰值）：Tensor Core 每秒最多算多少 FLOP，比如 B200 的 FP16 约 2000 TFLOPS。
  • 搬（内存带宽）：HBM 每秒最多搬多少字节，比如 B200 约 8 TB/s。

任何 kernel 都要同时”算”和”搬”，所以性能上限取两者中较小的那个：

$$
\text{可达 FLOPS} \le \min(\text{峰值算力},\ \text{带宽} \times \text{算术强度})
$$

这里的算术强度 = 每从 HBM 搬 1 字节，能做多少次 FLOP。它衡量数据”耐不耐算”：同样搬 1 字节，大矩阵能换来上千次运算，decode 只能换来 1 次。

两者的分界线叫脊点（= 峰值算力 ÷ 带宽，B200 约 250 FLOP/byte）：

  • 算术强度 < 脊点：数据不耐算，很快就搬完了，算力闲着 → 卡带宽，优化方向是减少搬运、提高复用。
  • 算术强度 > 脊点：数据太耐算，搬得再快也算不完 → 卡计算，优化方向是让 Tensor Core 跑满。

用两个例子验证（B200，脊点 ≈ 250）：

  • 两个 4096×4096 的 FP16 矩阵：算术强度 ≈ 1370 ≫ 250 → 卡计算。
  • decode 阶段的 1×4096×4096（向量 × 矩阵）：算术强度 ≈ 1 ≪ 250 → 卡带宽。

所以 Roofline 的价值不是算出具体数字，而是先确认卡在哪一边，再决定优化方向。它只回答”理论上限在哪”；能不能摸到上限，正是后面九步优化要解决的问题。

认识方法：从朴素 GEMM 到 SOTA

在我们讨论如何在 Blackwell GPU 上进行矩阵运算优化前先建立三个核心概念和视角，帮助我们从最主要的维度去一步一步理解优化的方向以及每次优化差异聚焦在哪里。

三个核心概念：scope、layout、dispatch

后面每一步优化，本质上都是在三个层面上做文章：

  • scope：一项操作由哪一层线程组织执行——thread、warp、warpgroup、CTA（即 CUDA 的 block）、cluster 级别。决定并发度和同步粒度。
  • layout：数据块（tile）放在哪、元素怎么排——放在哪级存储、如何映射到 lane / warp / 寄存器，swizzle、k-major、mn-major 都属此类。决定访存效率和 Tensor Core 能不能直接用。
  • dispatch：选哪条硬件路径来执行——比如搬运走 TMA 引擎还是普通访存指令，编译器根据 scope 和 layout 决定。决定走专用引擎还是通用路径。

9 步优化

原书的 GEMM 推演一共九步，跨三章连续编号。每一步解决一个具体问题，同时明确地动 scope / layout / dispatch 三个层面中的一处或几处：

第一阶段 · 把 GEMM 立起来（Building a Tiled GEMM，第 1–3 步）——先算对，不谈性能。

  • 第 1 步 · 单 tile GEMM：一个 CTA 算一个 128×128 输出 tile，走通 GMEM → SMEM → TMEM → 写回的完整数据通路，建立三个层面的基线。
  • 第 2 步 · K 循环累加：K 大于 64 时沿 K 分块循环，所有部分积累加到同一个 TMEM 累加器。（动 layout：同一块 SMEM 暂存区反复覆盖 + 同一个 TMEM 累加器原地累加）
  • 第 3 步 · 空间切分（Multi-CTA）：M、N 超出 128 时把输出切成 tile 网格，每个 CTA 负责一个 tile。（动 scope：从单 CTA 到二维 grid）

到第 3 步，覆盖完整矩阵的正确 GEMM 成立，但性能还不是目标——搬运同步、搬算不重叠、tile 之间零协作。

第二阶段 · 让搬和算重叠（Pipelining GEMM with TMA，第 4–6 步）——动数据通路。

  • 第 4 步 · TMA 异步搬运：把 GMEM ↔ SMEM 的搬运从”线程逐元素拷”换成 TMA 引擎一次搬一整块，硬件负责寻址和传输，mbarrier 跟踪完成。（动 dispatch：走专用搬运引擎）
  • 第 5 步 · 两级软件流水（PIPE_DEPTH=2）：之前只有一块 SMEM 暂存区，只能”搬一块、算一块、再搬下一块”——搬早了会覆盖正在算的那块。这步改成两块暂存区交替用：算第 i 块时，第 i+1 块已搬好。重叠有了可能，但还没真正并行——同一个线程组必须等当前这块算完，才能发下一次搬运；真正的并行等第 7 步。（动 layout：一块暂存 → 两块轮换）
  • 第 6 步 · Persistent kernel + tile 调度器：之前每个 tile 启动一个 CTA、各自初始化一遍；改为固定启动 ≈SM 数个 CTA，算完一个 tile 就从调度器领下一个：摊销初始化，调度顺序还能提升 L2 复用。（动 scope：CTA 池 + 调度器；调度器分软件和硬件 CLC 两种，后者见 Part I）

第三阶段 · 分工与协作（Scaling GEMM with Warp Specialization and Clusters，第 7–9 步）——动线程组织。

  • 第 7 步 · Warp Specialization：一个 warpgroup 既要发 TMA 又要发 MMA 又要写回，只能串行。拆成 producer（搬）、MMA consumer（算）、writeback（写回）三个角色，用 mbarrier 交接数据。持续重叠从这里开始。（动 scope：角色分工）
  • 第 8 步 · 2-CTA Cluster：单个 CTA 算不动更大的 tile。两个 CTA 组成 cluster 共算一个 256×256 tile：cooperative MMA 可以读对方 SMEM，累加器横跨两个 TMEM，TMA 一次把数据多播给两个 CTA，每块搬进来的数据参与约两倍的运算。（动 scope + layout + dispatch：执行范围跨 CTA）
  • 第 9 步 · Multi-Consumer Warp Specialization：tile 变大后一个 MMA consumer 消化不完。每个 CTA 加第二个 consumer：两个 consumer 共享同一个 B tile、各算一半输出行，输出 tile 从 256×256 长到 512×256，而每级 B 的搬运量不增。这一步后性能追平 cuBLAS。（动 scope + layout：consumer 再分细）

九步推演的本质：先算对（1–3），再让数据通路异步化（4–6），最后让执行主体分工协作（7–9）。每一步解决一个问题、暴露下一个瓶颈；优化不是”把所有技巧堆上去”，而是”找到当前瓶颈，针对性地消除它”。

同步的本质：顺序和互斥

这里笔者简单补充一个小思考，为什么在看这本书中的异步编程和我们之前写的 CPU 程序十分不同？

回头看九步推演，同步几乎全部由 mbarrier 这类硬件屏障完成：TMA 搬完、MMA 算完，各自把屏障推进一个 phase，等待方等到对应 phase 才继续。这和 CPU 编程的习惯很不一样——CPU 上保证并发正确性主要靠锁，底层由硬件原子指令（CAS 一类）支撑。

根本差异不在工具，而在并发模式：锁解决互斥——多个线程改同一块共享状态时，同时只放一个进去；屏障解决顺序——数据什么时候就绪、缓冲区什么时候能复用。GEMM 流水线是单向数据流（TMA → SMEM → MMA → TMEM → 写回），里面没有临界区，只有先后关系，所以互斥锁在这套编程模型里几乎没有位置。

锁缺席还有更硬的原因：GPU 线程以 warp 为单位集体执行，且没有 CPU 那样的线程级抢占调度。一个 warp 自旋等锁，会占住发射槽和寄存器，甚至可能饿死持锁的 warp 导致死锁。GPU 的原子指令并没有消失，只是用在了别处——计数器、任务分配（比如 tile 调度器领活），而不是做锁。

最有趣的一点是：屏障的”到达”可以是硬件做的。TMA 搬运完成、tcgen05 计算完成，硬件会自动推进 mbarrier 的 phase，线程只需要声明预期字节数（expect_tx）然后等待。CPU 上唤醒一个等待者，永远需要某个线程主动去做。

笔者觉得这是 GPU 异步编程和 CPU 并发最核心的不同：把”谁先谁后”交给硬件屏障和计数器，而不是让线程去争一把锁。这要求对硬件能力和并发逻辑有更精细的理解——也正是本书全程在做的事。

总结：这本书在说什么，没说什么

回顾这本书，笔者一开始读的时候也是非常痛苦，因为本质上它已经默认你了解了原来 Hopper 架构下的 CUDA 各种概念，什么是 SM，什么是 wrap……，对于 LLM 的 transformer 架构也没有展开， flash attention 这章因为篇幅原因笔者也不想再展开，后面在分析其他 vllm，llama.cpp 等推理框架时候再回来详细讲解其原理和实现。

以上通过不断和 agent 反复拉扯有了大致的理解，然后再来回顾下本书的核心理念，说明究竟想要表达什么，以及看清楚其没有覆盖到哪些方面，笔者和读者自身角度能够继续往哪些方向去延展和学习。

它说了什么

一本单 kernel 优化方法论的入门书。它以 Blackwell 为具体载体，展示了一条从硬件认知到性能模型到优化路径到完整 kernel 的推演过程。

它没有说什么

  • 没有 CUDA , Trix 等基本科普手册。
  • 没有说 LLM 基本原理和概念，没有训练/推理更全景的视角，没有讨论分布式 GPU 的情况。
  • 没有讨论其他芯片，以及不同设备上的情况，专注在 CUDA 生态。

清楚一本书的边界，笔者认为也非常有意义。了解它没说什么，才知道下一步该读什么。具体的实现和硬件可能不同，但是背后沉淀的建模思想，异步理念是可以迁移通用的。

后面可以考虑比如 kernel 算子，推理框架，分布式，端侧都可以考虑继续探究。

风起于青萍之末～