第 2 章核心模块：FlashAttention-2 与 PagedAttention 实现解析

文档摘要

第 2 章核心模块：FlashAttention-2 与 PagedAttention 实现解析本章导读：第一章我们把瓶颈拆成了"算子效率"和"显存管理"两半。这一章是全书的技术心脏，逐一拆解两个核心模块的内部机理。FlashAttention-2 的精髓在于"在线 softmax"——不物化完整分数矩阵也能算出正确结果；PagedAttention 的精髓在于"像操作系统管内存一样管 KV Cache"。读完你应该能分别说出：FlashAttention-2 靠哪三个设计把 GPU 利用率拉满，以及 PagedAttention 如何用块表（block table）把碎片几乎消灭。本章不含完整源码逐行，但会给出关键数据流的等价伪代码与原理图。 2.

第 2 章核心模块：FlashAttention-2 与 PagedAttention 实现解析

本章导读：第一章我们把瓶颈拆成了"算子效率"和"显存管理"两半。这一章是全书的技术心脏，逐一拆解两个核心模块的内部机理。FlashAttention-2 的精髓在于"在线 softmax"——不物化完整分数矩阵也能算出正确结果；PagedAttention 的精髓在于"像操作系统管内存一样管 KV Cache"。读完你应该能分别说出：FlashAttention-2 靠哪三个设计把 GPU 利用率拉满，以及 PagedAttention 如何用块表（block table）把碎片几乎消灭。本章不含完整源码逐行，但会给出关键数据流的等价伪代码与原理图。

2.1 FlashAttention-2 的核心思想

FlashAttention-2 的目标只有一个：让注意力计算全程待在 GPU 的 SRAM（片上高速缓存）里，尽量不碰 HBM。它依赖三个关键设计。

设计一：Tiling（分块）。把 Q、K、V 按行/列切成小块，每次只把一块载入 SRAM 计算，避免一次性物化 [n,n]。

设计二：Online Softmax（在线归一化）。标准 softmax 需要整行的最大值和总和才能归一化，所以必须看到全部元素。FlashAttention 引入了"递增式"归一化：每读入一块新的 K 分块，就更新运行最大值 m 和指数和 l，并对已累积的输出 O 做修正。数学上可以证明，这种增量更新与一次性 softmax 完全等价。

FlashAttention-2 Tiling 与 Online Softmax

设计三：减少非矩阵乘的 GPU 操作。FlashAttention-1 中 softmax 的多次 HBM 往返是瓶颈；FA-2 把"偶数次"的 HBM 访问进一步压到最少，并把工作负载在 warp 间重新划分，减少共享内存的通信与同步。

关键收益对比（以 A100 为例，理论数量级）：

显存占用：从 O(n²) 降到 O(n)。
实际加速：相比 PyTorch 朴素实现，FA-2 在长序列上可达数倍提速，且序列越长优势越大。
注意：FA-2 不改变数学结果，只是换了一种数值稳定的计算顺序，所以它是"无损加速"。

2.2 PagedAttention 的分页式 KV Cache

PagedAttention 是 vLLM 的核心创新，灵感直接来自操作系统的虚拟内存分页。

问题回顾：传统实现给每个请求预分配一整段连续的 KV Cache（按 max_length）。但真实生成长度千差万别，于是出现两类浪费：

内部碎片：预留了 4K，实际只用 200。
外部碎片：连续空闲块凑不成一个大请求。

PagedAttention 的解法：把 KV Cache 切成固定大小的块（block，如每块存 16 个 token 的 K/V），像内存页一样用**块表（block table）**把逻辑序列映射到物理块。新 token 来了就分配一个新块，不再连续预分配。

核心优势：

几乎消除碎片：块可离散存放，利用率大幅提升。
天然支持共享：并行采样、beam search 时，多个序列可共享同一前缀块（copy-on-write），KV Cache 占用再降一截。
更大吞吐：同样显存能塞下更大 batch，GPU 利用率上升。

等价伪代码（注意力按块取 K/V）：


# block_table: 逻辑块 -> 物理块
# 对每个 query 位置 i:
for logical_block in range(num_blocks_of_seq):
    phys = block_table[logical_block]
    K_blk = kv_cache_K[phys]   # 只取这一块
    V_blk = kv_cache_V[phys]
    scores = Q[i] @ K_blk.T
    # 与已累积的 running max/sum 合并
    accumulate_output(O[i], scores, V_blk)

2.3 两者如何协同

在工程落地中，二者是搭档而非替代：

PagedAttention 负责"KV Cache 存在哪、怎么分页取"，解决显存管理与调度。
FlashAttention-2 作为底层算子，负责"拿到一块 K/V 后怎么高效算注意力"，解决带宽瓶颈。
vLLM 默认组合：用 PagedAttention 管理显存，注意力内核可走 FlashAttention 或其他融合内核。

常见误区提醒：

认为开了 FlashAttention 就不需要 PagedAttention——错，前者不解决显存碎片。
认为 PagedAttention 会拖慢计算——实测相反，显存省下来换成了更大 batch 与更高吞吐。

下一章进阶原理，我们会把"为什么 FA-2 对长序列收益更大""PagedAttention 在共享场景下如何做写时复制"讲深，并给出可操作的性能测试方法论。

第 2 章 核心模块：FlashAttention-2 与 PagedAttention 实现解析

文档摘要