第 2 章 核心模块：FlashAttention-2 与 PagedAttention 实现解析
第 2 章 核心模块：FlashAttention-2 与 PagedAttention 实现解析
本章导读：第一章我们把瓶颈拆成了「算子效率」和「显存管理」两半。这一章是全书的技术心脏，逐一拆解两个核心模块的内部机理。FlashAttention-2 的精髓在于「在线 softmax」——不物化完整分数矩阵也能算出正确结果；PagedAttention 的精髓在于「像操作系统管内存一样管 KV Cache」。读完你应该能分别说出：FlashAttention-2 靠哪三个设计把 GPU 利用率拉满，以及 PagedAttention 如何用块表（block table）把碎片几乎消灭。本章不含完整源码逐行，但会给出关键数据流的等价伪代码与原理图，并在第 3 章把「为什么这样快」的原理讲深。
本章解决什么、不解决什么
在动手读原理之前，先划清边界，免得你带着错预期进来：
- 本章解决：两个核心模块「各自怎么工作、为什么有效、以及在引擎里怎么协同」。读完你能独立解释 FA-2 的三个设计、PagedAttention 的分页思想，以及它们为什么是搭档。
- 本章不解决：可运行的端到端部署配置、具体框架的命令行参数、以及「在你的硬件上到底能快多少」。前者留到第 4 章实战案例，后者留到第 3 章的性能测试方法论，并在第 4 章用真实配置落地。
一句话：第 1 章教你「怎么诊断」，第 2 章教你「两个引擎的内部构造」，第 3 章教你「为什么这样快、怎么测真」，第 4 章教你「怎么把它跑起来」。
阅读建议
- 如果你更关心「能不能直接上手」，建议第 2 章和第 4 章对照读：原理看这里，配置看实战。
- 如果你更关心「凭什么相信这些加速数字」，请务必读第 3 章的 Roofline 推导与压测方法论，别把本文的「数量级」当承诺。
- 每节开头都有一句「读者读完应带走的一句话」，读前先扫一眼，带着问题读，效率更高。
本章路线一览
两条优化路线的分工
第一章给出的判读是：推理加速有两条互补的战线。一条在算子层，解决「注意力算得慢」的带宽瓶颈；一条在系统层，解决「KV Cache 放不下、放不满」的显存碎片。本章的两个主角恰好一人守一条线，而且职责几乎不重叠：
- FlashAttention-2（算子层）：不让
[n, n]分数矩阵落 HBM，全程在 SRAM 内分块计算，把显存从 O(n²) 降到 O(n)，且数学结果无损。详见 2.1。 - PagedAttention（系统层）：把 KV Cache 像内存页一样分页，用块表映射离散物理块，几乎消灭碎片，并支持前缀共享。详见 2.2。
一次推理请求里，两者分别站在哪一层
在真实引擎（以 vLLM 形态为典型）中，它们是一前一后、各管一段的齿轮：
- 请求进入，调度器分配逻辑序列；
- PagedAttention 按块表按需分配物理 KV 块；
- FlashAttention-2 取一块 K/V，在 SRAM 内算注意力；
- 新 token 的 KV 写回新块，循环至生成结束。
读这一章最容易掉的三个坑
基于大量初学者和一线工程师的反馈，提前帮你排雷：
- 把「加速倍数」当成承诺。本文给出的都是「数量级与趋势」，具体数字随硬件、序列长度、精度、batch 剧烈变化。永远以你自己的压测为准，见第 3 章方法论。
- 以为两者二选一。FlashAttention-2 治带宽、PagedAttention 治碎片，职责不重叠。只上一个，另一个瓶颈照样卡你。见 2.3 协同分析。
- 混淆「省分数矩阵显存」与「省 KV Cache 显存」。前者是 FA-2 的功劳（O(n²) 临时显存变 O(n)），后者是 PagedAttention 的功劳（消除碎片、提升利用）。别张冠李戴。
版本演进的提醒
两个模块都在快速演进：FlashAttention 已有多个版本，注意力内核还在持续被新的融合实现（以及不同硬件厂商的专属内核）取代；PagedAttention 的分页思想也被更多引擎采纳并扩展（如加入前缀缓存的更精细管理）。本文聚焦「原理不变的那部分」，具体 API 与默认参数请结合你所用引擎的官方文档核实，不要凭记忆套用旧版本配置。
本章图示约定（先读 30 秒，省你后面反复猜）
本章配了多张原理图，统一约定如下，方便你快速抓重点：
- 蓝色框通常表示 FlashAttention-2 相关（算子层、SRAM 内计算）；
- 绿色框通常表示 PagedAttention 相关（系统层、显存分页）；
- 橙色框通常表示两者协同或对外入口（推理引擎、调度）；
- 箭头表示数据或控制流向，标注文字是这一步的作用，不是代码片段。
所有图都强调「分层」与「数据流」，建议你边看图边在脑子里走一遍：一个 token 从进入到写出 KV，经过了哪几层、每层谁负责。能顺畅走完一遍，说明原理已经吃透。
三节对应关系速查
| 小节 | 主角 | 解决的核心问题 | 关键图 |
|---|---|---|---|
| 2.1 | FlashAttention-2 | 注意力算得慢、O(n²) 临时显存 | Tiling / Online Softmax / HBM 对比 |
| 2.2 | PagedAttention | KV Cache 碎片、并发容量受限 | 连续分配 vs 分页 / 块表 / 共享 |
| 2.3 | 两者协同 | 怎么组合起来拿到端到端吞吐 | 分层栈 / 处理流 / 误区 |
带着这张表读，不容易迷路。
本章的前置知识与学完你会什么
读本章前，建议你已经掌握（若生疏，先回看第 1 章）：
- 标准注意力公式
softmax(QK^T/√d)V，以及它为何物化[n, n]分数矩阵； - KV Cache 为什么随序列和并发增长（第 1 章 1.2 的显存公式）；
- Roofline 的基本直觉：带宽受限 vs 计算受限。
读完本章，你应该能独立做到：
- 向同事讲清楚 FlashAttention-2 的三个设计，并说清「为什么它无损」；
- 用操作系统分页的比喻，解释 PagedAttention 怎么消灭碎片、怎么共享前缀；
- 在一张白纸上画出「一次推理请求里，两者各自站在哪一层」；
- 识别「只上其中一个就够了」之类的常见误区，并在方案评审里指出遗漏。
这四项不是背概念，而是能「讲出来、画出来、用得上」——这才是把原理真正变成工程能力的标志。
第二章我们把两个核心模块的原理和分工讲透了，带着「模块怎么搭」进入第三章，我们去看「为什么这样搭才快、怎么验证它真的快」。
如果你是按「问题驱动」的方式学习，这里给一条捷径：被「长序列显存爆」困扰，优先精读 2.1 与第 3 章的 Roofline 推导；被「并发上不去、显存利用率低」困扰，优先精读 2.2 与第 4 章的部署配置。两条线最终会在第 4 章实战里汇合。也提醒一句：本章所有性能表述都是「数量级与趋势」，不是可承诺的基准数字，落地前请用第 3 章的方法论在你的硬件上亲自测一遍。当你能不翻书就把这一章的三张核心图重画出来，原理就算真正内化了。
你应该从本章带走的三件事
- FlashAttention-2 的三个设计：Tiling（分块）、Online Softmax（在线归一化）、以及重新划分的 GPU warp 分工——三者合力把注意力留在 SRAM、压到最少 HBM 往返，做到无损加速。
- PagedAttention 的分页思想：固定大小块 + 块表映射 + 按需分配，像操作系统虚拟内存那样消灭碎片，并顺手用 copy-on-write 实现前缀共享。
- 两者是搭档不是替代：一个管「KV 存在哪、怎么取」，一个管「取到后怎么算」，叠加才同时拿到「显存省 + 算得快」。
下一章（第 3 章，进阶原理）将把「为什么 FA-2 对长序列收益更大」「PagedAttention 在共享场景下的正确性边界」「数值稳定性」与「性能测试方法论」讲透。
读到这里，你已经握住了两把钥匙：知道注意力为什么慢（第 1 章），也知道两个核心模块各自怎么动手（本章）。下一章我们把「为什么这样快」讲透，并给你一套能在自己硬件上复现的性能测试方法。