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

一个写给 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。

读完你能做到什么

建立现代 GPU 的心智模型。 SIMT 思维、Hopper 微架构、编程模型与内存层级——这四章决定了后面每一次优化你能不能预判收益（第1章 SIMT 范式、第2章 Hopper 架构、第3章编程模型、第4章内存层级）。
把基础算子写到极致。 Reduction、Online Softmax、LayerNorm、Elementwise 融合、量化解码——每一个都是 FA2 里真实存在的子问题，不是练习题（第5章 Reduction、第6章 Online Softmax、第7章 LayerNorm、第8章融合、第9章量化算子）。
走完 GEMM 之路。 从朴素实现到 Tiled、再到 Tensor Core，最后读懂 CUTLASS 的设计哲学——也就明白了 cuBLAS 凭什么比你快 30 倍（第10章朴素 GEMM、第11章 Tiled、第12章 Tensor Core、第13章 CUTLASS）。
手写 FlashAttention v2 到 SOTA。 IO-aware 的思想来源、前向与反向的完整实现，再用 Hopper 的三件武器（TMA、Warp Specialization、Persistent Kernel）把它推到峰值（第14章 IO-aware、第15章 FA2 前向、第16章 FA2 反向、第17章 TMA 与 Warp Specialization、第18章 Persistent Kernel）。
用工具定位瓶颈，而不是猜。 Nsight Compute / Systems 实战、读懂 PTX 与 SASS（编译器到底干了什么），以及一份性能陷阱与反模式清单（第19章 Nsight、第20章 PTX 与 SASS、第21章陷阱）。
接上下游。 CUDA Graph 与 Stream、CUDA C++ 与 Triton 的取舍、以及与 vLLM / Transformer 两个专栏的衔接路径（附录 A、附录 B、附录 C）。

为什么是 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 解剖：从 Attention 到推理系统》《vLLM 推理内核深度解析》《PyTorch 训练框架内核深度解析》里有完整论述。本专栏是它们的地基：读完它，你看那几个专栏时会清楚为什么 vLLM 选择手写 PagedAttention kernel 而不是用 cuBLAS。

源码版本

这个专栏本身不逐行解析某一个仓库，但正文在讲到具体实现时会指向真实的上游代码。下面是这些引用对应的仓库与版本。前三个仓库的路径是在本地 checkout 上逐条核对的，后两个（apex、cutlass）钉到具体 tag 后核对了路径存在——每一条都是照着能找到文件，不是凭印象写的路径：

仓库	版本	文中引用到的位置
Dao-AILab/flash-attention	`0251105`（2026-08-15）	`csrc/flash_attn/src/flash_fwd_kernel.h`、`csrc/flash_attn/src/flash_bwd_kernel.h`、`hopper/flash_fwd_kernel_sm90.h`
vllm-project/vllm	`v0.8.5.post1`	`csrc/attention/attention_kernels.cuh`、`csrc/layernorm_kernels.cu`、`csrc/quantization/gptq_marlin/gptq_marlin.cu`
pytorch/pytorch	`v2.11.0`	`aten/src/ATen/native/cuda/layer_norm_kernel.cu`
NVIDIA/apex	`25.09`	`csrc/layer_norm_cuda_kernel.cu`
NVIDIA/cutlass	`v4.7.0`	`include/cutlass/gemm/threadblock/mma_pipelined.h`

有两处目录结构容易踩空，这里说明白：FlashAttention 2 的内核在 csrc/flash_attn/src/ 下（不是 csrc/flash_attn/ 本身），而 FA3 的 Hopper 内核在独立的 hopper/ 目录，不在 csrc/ 里。

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention && git checkout 0251105

一份诚意

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

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

那时这个专栏的目的就达到了。

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

这个专栏在做什么

读完你能做到什么

为什么是 Hopper

这个专栏面向谁

这个专栏不是什么

目录

开篇

第一篇 · CUDA 与 Hopper 全景

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

第三篇 · GEMM 之路

第四篇 · 手写 FlashAttention v2

第五篇 · 性能工程

附录

源码版本

一份诚意