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

Name: CUDA 算子工程：手写 FlashAttention v2 之路
Author: 杨艺韬

一本写给 AI 系统工程师的 CUDA 专著：不教"怎么用 CUDA"，教"怎么把 LLM 推理算子写到 SOTA"。

打开 vLLM、SGLang、TensorRT-LLM 任何一个推理引擎的源码，往最深处走，你最终都会撞上同一堵墙：自定义的 CUDA kernel。PagedAttention 的访存核、FlashAttention 的 IO-aware 实现、AWQ/GPTQ 的反量化算子、Marlin 的 INT4 GEMM、CUTLASS 的 Hopper Kernel Schedule——这些 kernel 是大模型推理性能的最后一公里，也是大多数 AI 工程师从未真正读懂的一公里。

flowchart LR
  A[PyTorch 模型] --> B[算子库 cuBLAS/cuDNN]
  B --> C[自定义 CUDA Kernel]
  C --> D[PTX 中间表示]
  D --> E[SASS 机器码]
  E --> F[Hopper SM/Tensor Core/TMA]
  style C fill:#76b900,color:#fff
  style F fill:#1a73e8,color:#fff

绝大多数 CUDA 教程止步在 vectorAdd / matrixMul，最多再加一个朴素的 reduction。它们解释了语法，但没解释为什么 cuBLAS 的 SGEMM 比你写的快 30 倍、为什么 FlashAttention 不是"更聪明的算法"而是"更懂 GPU 内存层级"、为什么 Hopper 上一个 kernel 不用 TMA 就拿不到峰值。

九年前 NVIDIA 推出 Volta 架构时，第一代 Tensor Core 让深度学习算力翻倍；六年前 Ampere 引入异步拷贝，让 GEMM 的 prologue 和 mainloop 终于可以重叠；三年前 Hopper 推出 TMA + 分布式共享内存 + Warp Specialization，把 GPU 编程模型彻底改写。今天再用 Volta 时代的写法去写 Hopper kernel，性能差距是 5-10 倍——这不是夸张，FlashAttention v2 在 H100 上从首版到 SOTA 的优化路径，公开数据就是 5 倍以上。

这本书的目标是把这条路径走完。

这本书在做什么

这本书有一条贯穿始终的主线：手写 FlashAttention v2 到 SOTA。

我们不会一上来就开始写 attention。前 13 章会把所有需要的"手艺"逐一磨出来——GPU 编程模型、Hopper 微架构、内存层级、Reduction、Softmax、LayerNorm、量化解码、朴素 GEMM、Tiled GEMM、Tensor Core GEMM、CUTLASS 设计哲学。每一章的内容都不是孤立的练习题，而是一个 FA2 中真实存在的子问题。

flowchart TB
  subgraph 第一篇 [第一篇 · 全景]
    P1[第1章 SIMT 思维]
    P2[第2章 Hopper 架构]
    P3[第3章 编程模型]
    P4[第4章 内存层级]
  end
  subgraph 第二篇 [第二篇 · 基础算子]
    A1[第5章 Reduction]
    A2[第6章 Online Softmax]
    A3[第7章 LayerNorm]
    A4[第8章 算子融合]
    A5[第9章 量化 Kernel]
  end
  subgraph 第三篇 [第三篇 · GEMM 之路]
    G1[第10章 朴素 GEMM]
    G2[第11章 Tiled GEMM]
    G3[第12章 Tensor Core]
    G4[第13章 CUTLASS]
  end
  subgraph 第四篇 [第四篇 · 手写 FA2]
    F1[第14章 IO-Aware]
    F2[第15章 FA2 前向]
    F3[第16章 FA2 反向]
    F4[第17章 TMA + WS]
    F5[第18章 Persistent]
  end
  subgraph 第五篇 [第五篇 · 性能工程]
    Q1[第19章 Nsight]
    Q2[第20章 PTX/SASS]
    Q3[第21章 反模式]
  end
  第一篇 --> 第二篇
  第二篇 --> 第三篇
  第三篇 --> 第四篇
  第四篇 --> 第五篇

到第 14 章我们正式开始写 FA2 的时候，online softmax 已经在第 6 章被推导清楚、Tensor Core 的 mma 指令已经在第 12 章被驯服、CUTLASS 的 CollectiveOp 设计已经在第 13 章被剖开——所有的零件都准备好了。第 14-18 章这五章会把 FA2 从"能跑"一步步推进到"接近 SOTA"，每一步的优化都对应到一篇真实论文（Dao, 2022; Dao, 2023; Shah et al., 2024）或一个真实的 commit。

为什么是 Hopper

写 CUDA 书绕不开"硬件代际"这个选择题。我选 Hopper（H100/H800/H200）为主线，Ampere（A100）作为对照，理由有三：

Hopper 是 LLM 时代的事实标准。截至本书写作时（2026 年），全球部署的 LLM 推理与训练算力，70% 以上跑在 Hopper 上。读者今天能摸到的卡——无论是云上的 H100、企业内部的 H200，还是国内市场的 H800——都是 Hopper。Blackwell（B200）虽然是最新代际，但生态、驱动、CUTLASS 支持还在快速演化，写出来的内容半年就要重写一遍。
Hopper 是 CUDA 编程范式的转折点。Ampere 之前，写 CUDA 主要是"组织 thread + 用好 shared memory"；Hopper 引入 TMA 异步拷贝、分布式共享内存（Cluster）、Warp Specialization、第二代 setmaxnreg——整个 kernel 结构从"对称工作的 thread"变成"分工合作的 warp 组"。这是几十年一遇的范式跃迁，错过它就等于跳过了现代 CUDA 的核心。
Ampere 仍然有教学价值。每讲一个 Hopper 特性，我们会先用 Ampere 的写法做对照，让读者看清楚"演化"本身——为什么 TMA 取代了手写的异步拷贝、为什么 Cluster 让 reduce 比 grid-level 快 3 倍、为什么 Warp Specialization 比传统 occupancy 调优更彻底。

读者读完本书，Blackwell 的相关特性可以平滑迁移——FP4 是 FP8 的延伸，第二代 TMA 是第一代的扩展，CTA Pair 是 Cluster 的进化。我们会在第 17 章末尾给出 Hopper → Blackwell 的迁移指南。

这本书面向谁

AI 推理 / 训练框架工程师：你在 vLLM / SGLang / TensorRT-LLM / Megatron-LM 团队工作，或者准备投这些岗位。你需要看懂 PagedAttention kernel、需要给 Llama 写一个量化 GEMM、需要给客户的特殊模型 fine-tune 一个 fused kernel。本书第三、四篇是为你写的。
CUDA 自学者：你看完了 PMPP（《Programming Massively Parallel Processors》第 4 版）或者 NVIDIA 官方的 CUDA C++ Programming Guide，但发现书里讲的都是 Pascal/Volta 时代的写法，而你手上的 H100 跑出来的性能远低于 Whitepaper 上的数字。本书的现代化更新是为你写的。
HPC / 数值计算工程师：你做过 OpenMP / MPI / OpenACC，现在被推到 CUDA 上做 LLM 相关的工作。本书第一、二篇会帮你建立 GPU 的心智模型，第三、四篇会让你看到 LLM 算子和传统 HPC 算子（FFT、Stencil、SpMV）的本质区别。
想读 CUTLASS 和 FlashAttention 源码的人：你对着 GitHub 上 30 万行的 CUTLASS 代码不知道从哪入手，对着 FA2 论文的伪代码看懂了大意但写不出 kernel。本书第 13 章和第四篇是你的源码导读。

前置知识：本书假设读者熟悉 C++（能看懂模板、constexpr、RAII），具备线性代数基础（矩阵乘法、向量内积），了解神经网络的基本概念（attention、layernorm、softmax 是什么）。不要求读过 CUDA 之前的内容——第 1-4 章会从零建立 CUDA 心智模型，但节奏比 PMPP 快，且直接以 LLM 算子为例。如果读者完全没碰过 GPU，建议先扫一遍 NVIDIA 官方的 CUDA C++ Programming Guide 第 1-3 章。

这本书不是什么

这本书不是 PMPP 的替代品。PMPP 是经典的并行编程教材，覆盖面比本书广（包括 stencil、scan、sparse、graph 等多种 pattern）。本书更窄、更深，专注在"现代 GPU 上写 LLM 算子"这一个赛道。如果读者关心通用并行编程，PMPP 是更合适的入门书。

这本书不是 cuDNN / cuBLAS 调用手册。我们会拿 cuBLAS 作为性能基准，会读 cuBLAS 的部分实现，但不会教读者怎么调用 cuBLAS。"调用 API" 是 5 分钟就能学会的事，"理解 API 内部为什么这么快" 才是本书要做的事。

这本书不是 Triton 教程。Triton 是 OpenAI 出品的 GPU 编程 DSL，用 Python 语法写 GPU kernel，是另一种范式。附录 B 会专门对比 CUDA C++ 与 Triton 的取舍——大多数生产级算子最终还是要回到 CUDA C++（CUTLASS、cuBLAS、TensorRT 都是 C++），但 Triton 在原型期非常高效。

这本书不是模型训练 / 部署教程。怎么用 PyTorch 训一个 LLM、怎么用 vLLM 部署一个推理服务——这些在《Transformer 解剖》《vLLM 内核探秘》《PyTorch 训练框架内核》里有完整论述。本书是它们的地基：读完它，你看那几本时会清楚为什么 vLLM 选择手写 PagedAttention kernel 而不是用 cuBLAS。

一份诚意

GPU 编程是少数几个**"门槛极高、回报极大"**的工程领域。门槛在哪里？在于硬件细节——Tensor Core 的 fragment layout、TMA 的 swizzle 模式、SMEM 的 bank 排布、L2 的 set-associative cache 行为——这些细节没有捷径，只能一行行硬啃。回报在哪里？在于一行 kernel 代码可能让 1 万张卡的集群利用率提升 20%，那就是几亿美元的算力。

这本书不会替读者跳过门槛，但会把每一道门后面值得看的风景都标出来。读完它，希望读者下一次看到 cuBLAS / CUTLASS / FlashAttention 的源码时，眼里看到的不再是密密麻麻的模板和 PTX，而是一个个清晰的"为什么这么写"的工程决定。

那时这本书的目的就达到了。

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

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

这本书在做什么

为什么是 Hopper

这本书面向谁

这本书不是什么

目录

开篇

第一篇 · CUDA 与 Hopper 全景

第二篇 · 把基础算子写到极致

第三篇 · GEMM 之路

第四篇 · 手写 FlashAttention v2

第五篇 · 性能工程

附录

一份诚意