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

附录 C · 与 vLLM·Transformer 衔接路径

作者杨艺韬 · 1,367 字

C.1 三个专栏的分层

LLM 系统工程是一个三层栈：

flowchart TB
  subgraph App [应用层]
    A1[LangGraph / Agent]
    A2[RAG / 工具使用]
  end
  subgraph Engine [引擎层 · 《vLLM 推理内核深度解析》]
    E1[PagedAttention]
    E2[Scheduler]
    E3[KV Cache 管理]
    E4[多卡通信]
  end
  subgraph Model [模型层 · 《Transformer 解剖》]
    M1[Attention 数学]
    M2[Position Encoding]
    M3[FFN / MoE]
    M4[LayerNorm]
  end
  subgraph Kernel [Kernel 层 · 本专栏]
    K1[GEMM]
    K2[FlashAttention]
    K3[量化 Kernel]
    K4[元素级算子]
  end
  App --> Engine --> Model --> Kernel

每一层的 kernel 都建立在下一层之上：

vLLM 的 PagedAttention 用本专栏第 14-18 章讲的 FA 思想。
Transformer 的 Attention 数学公式落到本专栏的 Tensor Core GEMM + Online Softmax。
vLLM 的 INT4 推理 用本专栏第 9 章讲的量化 kernel。

C.2 Transformer 解剖：模型层的视角

《Transformer 解剖：从 Attention 到推理系统》（19 章）回答的问题是：模型本身是怎么运作的？

它的章节与本专栏的对应关系：

Transformer 章节	主题	对应到本专栏
第 2 章 Self-Attention	Q/K/V 的数学推导	本专栏第 14 章（attention 访存）
第 3 章 Multi-Head	头维度的并行	本专栏第 15 章（head 维 grid 划分）
第 4 章 Position Encoding	RoPE 的旋转	本专栏第 8 章（element-wise fusion）
第 6 章 LayerNorm	归一化的意义	本专栏第 7 章（kernel 实现）
第 12 章 KV Cache	推理时的状态	本专栏第 9 章（量化 KV Cache）
第 13-14 章 Flash Attention	论文阅读	本专栏第 14-18 章（手写 FA2）
第 16 章 Quantization	INT8/INT4 算法	本专栏第 9 章（量化 kernel 实现）

阅读建议：

从未读过 Transformer 的工程师：先读《Transformer 解剖》第 1-12 章建立模型层认知，再回到本专栏把每个算子的实现细节填进去。
已经熟悉 Transformer：直接用《Transformer 解剖》作为参考——本专栏讲到 LayerNorm 时翻看 Transformer 第 6 章的数学，讲到 FA 时翻看 Transformer 第 13-14 章的论文导读。

C.3 vLLM 内核探秘：引擎层的视角

《vLLM 推理内核深度解析》（19 章）回答的问题是：怎么把模型跑得既快又稳？

它的章节与本专栏的对应关系：

vLLM 章节	主题	对应到本专栏
第 4 章 PagedAttention	KV cache 分页管理	本专栏第 18 章（persistent kernel）
第 5 章 KV Cache	多 batch 的 KV 排布	本专栏第 4 章（内存层级）
第 8 章 Model Runner	CUDA Graph 集成	本专栏附录 A（Graph）
第 9 章采样	logits 处理	本专栏第 8 章（element-wise）
第 10 章前缀缓存	长 prompt 优化	本专栏第 4 章（L2 cache）
第 13 章量化引擎	AWQ/GPTQ/Marlin	本专栏第 9 章（量化 kernel）
第 14 章张量并行	多卡 GEMM	本专栏第 13 章（CUTLASS）

阅读建议：

想成为 LLM 推理工程师：先读本专栏（kernel 层），再读 vLLM（引擎层）。本专栏提供"为什么 PagedAttention 那样写"的根本理解；vLLM 提供"怎么把它工程化成生产级"的工程经验。
已经在用 vLLM 但想理解内部：用本专栏填补 vLLM 文档没讲到的 kernel 层细节。

C.4 推荐阅读顺序

针对不同背景的读者：

路径 A：完全新手（从未碰 LLM）

《动手学深度学习》前 5 章（PyTorch 基础）
《Transformer 解剖》第 1-12 章（模型层）
本专栏第 1-9 章（GPU 基础 + 小算子）
《Transformer 解剖》第 13-19 章（推理系统）
本专栏第 10-21 章（GEMM + FA2 + 性能工程）
《vLLM 推理内核深度解析》全部

路径 B：会用 PyTorch 但不懂底层

《Transformer 解剖》第 1-10 章（建立模型直觉）
本专栏全部（Kernel 层完整训练）
《vLLM 推理内核深度解析》第 1-10 章（推理工程）

路径 C：已经写过 CUDA，想转向 LLM

本专栏第 1-2 章（Hopper 架构补课）
本专栏第 14-18 章（FA2 实战）
《Transformer 解剖》第 13-14 章（FA 论文导读）
本专栏第 9 章（量化）
《vLLM 推理内核深度解析》第 4 章（PagedAttention）

路径 D：已经在做 LLM 推理优化

直接精读：

本专栏第 17 章（TMA + WGMMA + Warp Spec）
本专栏第 13 章（CUTLASS 设计哲学）
《vLLM 推理内核深度解析》第 13 章（量化引擎）
FA3 论文 (Shah et al., 2024) 原文

C.5 推荐工程实践项目

如果读者想在读完本专栏后做一个能放在简历上的项目：

入门项目（1-2 周）

实现一个 fused LayerNorm + 残差加 kernel，对比 PyTorch 默认实现的延迟和带宽。
在 H100 上把 reduce 从 100 GB/s 调到 3+ TB/s（本专栏第 5 章的演进）。
写一个 INT8 dequant + GEMM fused kernel，对比朴素拆分。

中级项目（1-2 月）

实现一个能跑 LLaMA-7B 的简化推理引擎，包含 attention、FFN、LayerNorm 三大算子。
把第 11-12 章的 GEMM 调到 cuBLAS 90% 性能。
复现 FA2 forward + backward，在 H100 上达到 70%+ 算力利用。

高级项目（3-6 月）

复现 FA3 forward (TMA + WGMMA + Warp Spec)，目标 80%+ 算力。
实现一个 INT4 GEMM (Marlin-style)，达到 cuBLAS FP16 的 90%+ 性能。
给开源项目（vLLM、SGLang、Megatron-LM）贡献一个 kernel 优化 PR。

C.6 结语

CUDA 算子工程是少数几个门槛极高、回报极大的工程领域。

门槛在于硬件细节——Tensor Core 的 fragment layout、TMA 的 swizzle 模式、SMEM 的 bank 排布、L2 的 set-associative cache 行为——这些细节没有捷径，只能一行行硬啃。

回报在于一行 kernel 代码可能让 1 万张卡的集群利用率提升 20%——那就是几亿美元的算力。

这个专栏希望读者读完后，下一次看到 cuBLAS / CUTLASS / FlashAttention 的源码，眼里看到的不再是密密麻麻的模板和 PTX，而是一个个清晰的"为什么这么写"的工程决定。

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

祝读者在 GPU 工程的路上走得远。