CUDA 算子工程：手写 FlashAttention v2 之路

前言：为什么又一本 CUDA 书

作者杨艺韬 · 2,855 字

前言：为什么又一本 CUDA 书

"We do not need a new algorithm. We need to compute the same thing in a way that respects the memory hierarchy of the hardware we have." ——Tri Dao 等，《FlashAttention》, NeurIPS 2022

一篇没有新算法的论文

2022 年 5 月 27 日，斯坦福大学的博士生 Tri Dao 在 arXiv 上传了一篇论文，题目是《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》。

这是一篇有些"反常"的论文。它没有提出新的注意力机制，没有改变 softmax 的数学定义，也没有动 Q/K/V 三个矩阵的乘法本身——它做的事情，用一句话讲就是：把已经存在的、所有人都在用的、教科书级别的注意力计算，按 GPU 的内存层级重新组织了一遍。

结果呢？

                              FlashAttention v1（2022）
端到端 BERT-Large 训练:        加速 15%
端到端 GPT-2 训练:             加速 3 倍
端到端 LRA benchmark:          加速 2.4 倍
显存占用:                      从 O(N²) 降到 O(N)

数字来源：Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022, Table 5–6.

一篇没有新算法的论文，让大模型训练成本降低了 30%-70%。它在 NeurIPS 2022 收到的口头汇报和后续引用次数（截至 2026 年 5 月超过 4500 次），让 Tri Dao 直接拿到了 Together.ai 的联合创始人席位。一年后他发表了 FlashAttention v2，在 A100 上把性能再翻了一倍，达到了 230 TFLOPs/s，相当于 cuBLAS GEMM 性能的 ~73%。又过了一年，FlashAttention v3 论文在 H100 上跑到 740 TFLOPs/s（FP16）和 1.2 PFLOPs/s（FP8）——后者已经接近 H100 SXM5 FP8 Tensor Core 理论峰值 1979 TFLOPs/s 的 60%。

数字来源：Shah, Bikshandi, Zhang, Thakkar, Ramani, Re, Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, NeurIPS 2024.

这是一条非常清晰、非常陡峭的曲线：

                  FA1        FA2        FA3
                (A100)     (A100)     (H100, FP16)
TFLOPs/s:        ~120       ~230       ~740
% of cuBLAS:     ~38%       ~73%       ~75%

每一次跃迁，背后都不是"更聪明的 attention 算法"。是对 GPU 硬件结构、对内存层级、对异步指令调度更精细的理解。FA1 解决的是"HBM 访问太多"；FA2 解决的是"warp-level 并行没拆好"；FA3 解决的是"Hopper 上 TMA 异步拷贝和 Tensor Core 之间没流水起来"。

如果说 2017 年那篇《Attention Is All You Need》定义了现代大模型的架构，那么 2022 年这篇《FlashAttention》定义了现代大模型的实现范式——从那以后，每一个想认真做 LLM 推理或训练的工程团队都意识到：算法已经基本固定了，剩下的全是工程，而工程的核心就是把同一份计算用 GPU 喜欢的方式重新组织一遍。

这本书想做的事情，就是把这种"重新组织"的手艺，从头到尾教给读者。

一个数字：H100 的 989 TFLOPs/s

NVIDIA 官方公布的 H100 SXM5 规格（详见 NVIDIA H100 Tensor Core GPU Architecture Whitepaper, 2022）里，最显眼的数字是这一行：

FP16 Tensor Core 峰值算力:        989 TFLOPs/s
FP8 Tensor Core 峰值算力:         1979 TFLOPs/s
HBM3 显存带宽:                    3.35 TB/s
SMs 数量:                         132
寄存器文件总容量:                 16.5 MB
SMEM 总容量:                      33 MB
L2 cache:                         50 MB

把这些数字放在一起，能看出一个很值得关注的失衡：

flowchart LR
  COMP[计算: 989 TFLOPs/s]
  MEM[带宽: 3.35 TB/s]
  COMP -->|算力/带宽 = 295 FLOPs/byte| RATIO[算术强度阈值]
  style COMP fill:#76b900,color:#fff
  style MEM fill:#1a73e8,color:#fff
  style RATIO fill:#fbbf24,color:#000

这个 295 FLOPs/byte 的数字是什么意思？它说的是：在 H100 上，每从 HBM 读 1 字节，你必须做至少 295 次浮点运算，才不会被带宽卡住。

把这个阈值代回到典型的 LLM 算子：

算子	算术强度（FLOPs/byte）	在 H100 上能跑到峰值吗？
Element-wise（加法、ReLU）	0.25	不能，被带宽卡死
LayerNorm / RMSNorm	~1	不能
Attention（短序列）	~10	不能
Attention（长序列，FA2）	~50–100	接近峰值
GEMM（M=N=K=4096）	~1300	可以打到 80%+
GEMM（small batch）	~10–50	不能

所以一个残酷的事实是：LLM 推理时绝大多数 kernel 都跑不到 H100 的算力峰值——它们被 HBM 带宽卡住了。整个 vLLM、SGLang、TensorRT-LLM 团队这些年最大的工作，都是围绕"怎么把更多的算搬到 SMEM 和寄存器里、减少 HBM 往返"展开。FlashAttention 干的就是这件事。PagedAttention 干的也是这件事。INT4 GEMM、KV Cache 量化、投机解码——它们的本质都是同一件事：减少 HBM 访问。

不理解这个数字，就不理解为什么 LLM 推理性能优化是一门"内存科学"而不是"算力科学"。这本书第 4 章会反复回到这个数字上。

为什么又一本 CUDA 书

如果你今天搜 "CUDA tutorial"，能找到的资源大致分四类：

教科书类：以 PMPP（《Programming Massively Parallel Processors》, Hwu / Kirk / Hajj）为代表。覆盖面广，讲 stencil、scan、graph、sparse 等多种 pattern，适合系统性学习并行编程。但它不是为 LLM 时代写的——书里很多例子在 Pascal/Volta 上很经典，但放到 Hopper 上需要重新翻译。
官方文档类：NVIDIA CUDA C++ Programming Guide、PTX ISA、Nsight 文档。权威、准确，但没有故事线。读者在文档里找得到 cp.async.bulk.tensor 的语法，但找不到"为什么这条指令会存在"的工程背景。
博客碎片类：从 Mark Harris 当年的 reduction 系列博客，到近年各类"手写 GEMM 性能压榨"的中文博客。质量参差，时效性差，且互相之间没有体系——读者读完 reduction，不知道下一步该读什么。
论文与源码：FlashAttention 系列论文、CUTLASS 文档、cuBLAS 反汇编、各大推理引擎的 kernel。是一手资料，但门槛极高——读者没有第 1-13 章的铺垫，很难看懂第 14-18 章里的工程决定。

这本书想做的事情，是把第 4 类资料的内容，用第 1 类资料的体系组织起来，配上第 2 类资料的权威数字，并避开第 3 类资料的过时与碎片化。

具体来说：

它有一条贯穿全书的主线——手写 FlashAttention v2。每一章的每一个例子，都不是孤立的练习，而是 FA2 中真实存在的子问题。读者读完第 6 章 online softmax 的时候不会问"这个有什么用"，因为他们已经知道这是 FA2 第 15 章要用的核心组件。

它锚定一个具体的硬件代际——Hopper（H100）。每一个特性都基于 Hopper 的真实数字（132 SMs、989 TFLOPs/s、3.35 TB/s HBM3、228 KB SMEM/SM），而不是泛泛的"GPU"。当我们讲 TMA 时，会附 PTX cp.async.bulk.tensor.2d.shared::cluster.global 指令的真实编码；当我们讲 Tensor Core 时，会附 mma.sync.aligned.m16n8k16 在 SASS 上对应的具体延迟周期数。

它直接以 LLM 算子为例——而不是用 vector add、Mandelbrot 集这些经典 CUDA 教学例子。读者写的第一段 reduction 就是 LayerNorm 的统计量；写的第一段 softmax 就是 attention 中的 softmax；写的第一段 GEMM 就是 LLM 矩阵乘的一种 tile 配置。所有的例子都能直接对应到读者会在 vLLM / Megatron / FlashAttention 源码里看到的代码。

它讲源码而不是讲 API——尤其在 CUTLASS 那一章。CUTLASS 是 NVIDIA 官方维护的 CUDA C++ 模板库，30 万行代码，门槛极高，但它是现代 CUDA 工程的"事实标准库"，FlashAttention v3、cuBLAS 的部分实现、各大推理引擎的 GEMM 后端都建立在它之上。本书第 13 章不会教读者"怎么用 CUTLASS"——因为它的 API 在持续演化——而是教"CUTLASS 为什么这样设计"，让读者下一次升级到 CUTLASS 4.x 时，能自己读懂新的接口。

这本书的诚意

我得诚实地告诉读者两件事。

第一件事：这本书不会让任何人变成 CUDA 高手——只有对着真卡反复写、反复 profile、反复看 SASS，才能成为高手。书的作用是把"应该看哪些细节、用哪些工具、按什么顺序去理解"这条路径标清楚。读者读完后还需要自己花至少几百小时在 H100 / A100 上写 kernel，才能把书里的内容真正吸收。

第二件事：CUDA 的细节会变。Hopper 之后是 Blackwell，Blackwell 之后还会有新代际。本书选择的"主线 + 锚点"策略，是希望即使硬件换了，读者也能复用大部分思维框架——TMA 之后是第二代 TMA，Tensor Core 4 代之后是 5 代，FP8 之后是 FP4——细节会变，但"算力 vs 带宽的失衡"、"SMEM 是黄金"、"warp 协作是杠杆"、"指令异步性是关键"这些原则不会变。把握住这些原则，读者下一次面对新代际硬件时不会从零开始。

怎么读这本书

flowchart TB
  start[开始]
  q1{有 GPU 编程基础吗？}
  q2{熟悉 CUDA 语法吗？}
  q3{读过 CUTLASS 或 FA 源码吗？}
  read1[从第 1 章顺序读]
  read2[第 1-4 章快速过, 然后顺序读]
  read3[第 1-9 章选读,精读 GEMM 与 FA 部分]
  read4[直接进第 13 章 CUTLASS, 然后第 14-18 章 FA2 实战]

  start --> q1
  q1 -->|没有| read1
  q1 -->|有| q2
  q2 -->|不熟| read2
  q2 -->|熟悉| q3
  q3 -->|没读过| read3
  q3 -->|读过| read4

完全没有 GPU 经验的读者：从第 1 章开始顺序读。前 4 章会建立 GPU 心智模型，建议同时打开 NVIDIA H100 Whitepaper 第 2-4 章作为参考。第一遍读完后，建议拿一台 A100 / H100（云上租 30 分钟也行）把第 5、10、12 章的代码亲手跑一遍，再回头读第二遍。

有 OpenMP / MPI / OpenACC 经验，但没写过 CUDA 的读者：第 1-4 章可以扫读，重点关注与 CPU 多线程的差异——尤其是 SIMT 模型和 SMEM 的角色。从第 5 章 Reduction 开始正式精读。

有 CUDA 基础但停留在 Pascal/Volta 时代的读者：你的痛点很可能是"在 Hopper 上老写法跑不出峰值"。建议第 1-3 章扫读，第 4 章（内存层级）精读，第二、三篇按需选读，第四篇（FA2 实战）精读——这是 Hopper 时代写法和老写法差异最明显的地方。

已经读过 FlashAttention 论文 / CUTLASS 源码 / vLLM kernel 的高级读者：建议直接从第 13 章开始。前 12 章可以作为参考查漏补缺。第 17 章（TMA + Warp Specialization）和第 20 章（PTX/SASS）应该会让你拿到一些新东西。

无论从哪一章开始读，建议每读完一章都做一件事：打开你最熟悉的一份开源 CUDA 代码（vLLM、CUTLASS、Triton 内核、xformers 都行），找到这一章对应的代码片段，看看你现在能不能读懂之前读不懂的东西。这本书的价值不在它本身，而在它能让多少行你之前看不懂的代码变得可读。

致谢与免责

这本书的所有数字，能引用论文 / Whitepaper 的都引用了来源；不能引用的（比如某些工程经验值），会在脚注里标"经验值"。如果读者发现任何数字有误、任何描述与新代际硬件不符，欢迎在 GitHub 上提 issue。

CUDA 是 NVIDIA 的注册商标，本书所有讨论基于 NVIDIA 官方公开文档（CUDA C++ Programming Guide v12.x、PTX ISA v8.x、Hopper Whitepaper）以及公开论文与开源代码。本书不涉及任何 NVIDIA 内部信息。

最后，把这本书献给那些深夜对着 Nsight Compute 报告、试图把 kernel 性能再榨出 5% 的 AI 系统工程师。你们的每一行 kernel 优化，都在替整个 AI 产业省下数以亿计的算力成本。这本书是为你们写的。

下一章我们正式开始。第 1 章会先把"GPU 为什么不是更多核的 CPU"这件事讲清楚——这是后续所有内容的地基。