Lecture 4 · 高效注意力内核：FlashAttention 与 FlashInfer · Efficient Attention

对应 Maps to：Stanford CS349D AI Inference Infrastructure（Spring 2026），#10 Wed Apr 29 Efficient Attention。必读：T. Dao, D. Y. Fu, S. Ermon, A. Rudra, C. Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（NeurIPS 2022）；Z. Ye, L. Chen, R. Lai, W. Lin, Y. Zhang, S. Wang, T. Chen, B. Kasikci, V. Grover, A. Krishnamurthy, L. Ceze, FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving（MLSys 2025）。选读：J. Shah et al., FlashAttention-3（2024）；T. Zadouri et al., FlashAttention-4（2026）。 依据说明：原课没有公开 slides。本讲依据两篇必读论文的全文，以及 FlashAttention-3、FlashAttention-4 的摘要和引言来写。§3.5 关于 FlashAttention-2 的内容来自它的论文（Dao, 2023）。§6”之后的发展”是我补充的。 配套代码 Code：code/lecture04/（numpy + 标准库；4 个脚本 + 1 个模块 attn_kernels.py：标准注意力、分块注意力与在线 softmax、注意力状态的合并、split-KV、级联注意力、块稀疏格式、注意力变体、负载均衡调度；每个脚本几秒以内） 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 解释注意力在推理里为什么是一个特殊的算子：prefill 时计算量随长度平方增长、中间结果巨大；decode 时它的主要工作是读 KV。
2. 说出 GPU 的两级存储（HBM 与片上 SRAM）的容量和带宽差多少，以及”IO 感知”（IO-aware）的意思：少搬数据，而不只是少算。
3. 推导在线 softmax，解释 FlashAttention 怎样用分块 + 在线 softmax + 算子融合，让 $N \times N$ 的分数矩阵从不写回显存。能说出它的 IO 复杂度 $\Theta(N^2 d^2 / M)$ 与标准实现的 $\Theta(Nd + N^2)$，以及这个比值在实际参数下大约是多少。
4. 概括 FlashAttention-2、-3、-4 各自针对的问题，理解”内核要跟着硬件一起演化”。
5. 分析 decode 注意力：它的算术强度约等于 GQA 的组大小，批处理提高不了它；批小、上下文长时 GPU 的并行度不够，split-KV 怎样解决；注意力状态的合并运算 ⊕ 为什么精确、为什么合并顺序会影响可复现性。
6. 解释 FlashInfer 的三个设计：用块稀疏格式统一各种 KV 布局、用可组合格式和级联注意力利用共享前缀、用 JIT 模板支持注意力变体，以及它的负载均衡调度怎样处理长短不一的请求、怎样与 CUDA Graph 兼容。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  为什么注意力需要专门的内核：它在推理里的分量、GPU 的两级存储、算子融合
§2  标准注意力：搬运的账                                                     ← 代码 01 演示 3、5
§3  FlashAttention：IO 感知、在线 softmax、分块算法、IO 复杂度、FA2→FA4、评测   ← 代码 01
§4  decode 阶段的注意力：强度、并行度与 split-KV、⊕ 合并、级联注意力           ← 代码 02
§5  FlashInfer：块稀疏格式、可组合格式、JIT 模板、负载均衡与 CUDA Graph、评测    ← 代码 03、04
§6  之后的发展
§7–11  小结、常见误区、练习、延伸阅读、下一讲预告

0.4 记号 Notation

0.5 关于数字 · About the numbers

本讲的数字有三类：(1) 按算法的数据流记账的 HBM 读写量（代码 01）：numpy 没有真正的 SRAM，“读写多少”是按 FlashAttention 论文的分析方式数出来的，小规模时用实现里的计数器数、并验证与公式完全一致，大规模时用公式；(2) 简化的 GPU 模型（代码 02、04）：H100 的 132 个 SM 平分显存带宽，一个工作块的时间 = 它要读的字节 ÷ 每个 SM 的带宽，只用来说明并行度和负载均衡；(3) CPU 实测（代码 04 演示 3 的调度耗时）。所有正确性检查都是真实计算。论文里的数字写成”论文报告”。

1. 为什么注意力需要专门的内核 · Why attention needs special kernels

1.1 注意力在推理里的分量

前三讲一直把注意力当成理所当然的一步，可它恰恰是推理里最难做快的算子：

• prefill：注意力的计算量是 $O(N^2 d)$（每个 query 和它之前的每个 key 做点积），而线性层是 $O(N d^2)$。上下文很长时，注意力的计算量开始占主导（第 1 讲代码 02 演示 3：32K 的提示，平均每 token 的时间明显上升）。更麻烦的是中间结果：$N \times N$ 的分数矩阵，$N = 32\text{K}$ 时一个头就是 2 GB（FP16）。
• decode：每个新 token 的 query 要和它的整段 KV 做注意力。注意力本身的计算量不大，但要把整段 KV 从显存读一遍。第 1 讲代码 02 演示 4 看到，批 256、上下文 1024 时，读 KV 已经占了全部搬运量的 70%。

而且，推理系统对注意力还有特殊的要求：一批里的请求长短不一（第 2 讲），KV 散落在分页的块里（第 3 讲），多个请求共享前缀（第 3 讲），新模型还在不断引入注意力的变体。

1.2 GPU 的两级存储

要理解高效的注意力内核，需要比第 1 讲”计算器 + 水管”再细一层。GPU 上的数据主要待在两个地方（前置讲 P1 会详细讲）：

（数字来自 FlashAttention 论文对 A100 的描述。）GPU 上的一个内核（kernel）就是一次运算的程序：它从 HBM 把输入读进片上（寄存器、SRAM），算完，把输出写回 HBM。

1.3 大多数算子受限于搬运

第 1 讲的 roofline 告诉我们，算术强度低于屋脊点的计算受限于带宽。矩阵乘法（每个元素被复用很多次）可以做到 compute-bound，但 softmax、加掩码、逐元素运算这类算子，每个元素读进来只做几次运算就写回去，永远是 memory-bound 的。如果一个计算由好几个这样的算子组成，每个算子都把中间结果写回 HBM、下一个再读出来，时间就几乎全花在搬运上。

解决办法叫算子融合（kernel fusion）：把几个算子合并成一个内核，数据读进片上一次，在片上做完所有步骤，只把最终结果写回。注意力（矩阵乘法 → 缩放 → 掩码 → softmax → 矩阵乘法）正是最需要融合的地方。困难在于：softmax 需要整行分数才能算，而一整行分数（$N$ 个数，乘以很多行）放不进 SRAM。FlashAttention 就是解决这个困难的。

2. 标准注意力：搬运的账 · Standard attention’s IO

标准实现（FlashAttention 论文的 Algorithm 0）分三步，每一步都是一个独立的内核：

1. 从 HBM 读 Q、K，算 $S = QK^\top / \sqrt{d}$，把 $S$ 写回 HBM；
2. 读 $S$，算 $P = \text{softmax}(S)$，把 $P$ 写回 HBM；
3. 读 $P$、$V$，算 $O = PV$，把 $O$ 写回 HBM。

$Q$、$K$、$V$、$O$ 各是 $N \times d$，$S$、$P$ 各是 $N \times N$。总的 HBM 读写是 $\Theta(Nd + N^2)$。$N$ 比 $d$ 大得多时，$N^2$ 那一项占主导，系数约为 4（写 S、读 S、写 P、读 P）。而且 $S$ 和 $P$ 要存在显存里：显存占用也是 $O(N^2)$。

代码 01 演示 5 把这笔账放到真实模型上：如果 Llama-3-8B 的 prefill 用标准注意力，它比 FlashAttention 多出来的读写（每个头、每层；FP16）按 H100 的带宽换算成时间：

提示短时，分数矩阵不大，差别很小；提示变长，多出来的读写按 $N^2$ 增长，32K 时超过了模型本身的计算时间。这还没算显存：标准实现要同时存下一层所有头的 $S$ 矩阵，32K 时每个头 2 GB、32 个头 64 GB，一块 H100 根本放不下。

3. FlashAttention

3.1 关键洞察：IO 感知

FlashAttention 论文开篇就指出：当时有大量”近似注意力”的工作（稀疏、低秩……）把计算量降到线性或近线性，可它们大多没有带来实际的加速，因为它们只盯着 FLOPs，忽略了搬运（IO）的开销。论文提出的原则是IO 感知（IO-aware）：设计算法时，仔细计算在不同层级的存储之间搬了多少数据。

FlashAttention 是精确的注意力（不是近似），它的目标是：在片上把整个注意力算完，$S$ 和 $P$ 从不写回 HBM。障碍是 softmax 需要整行，下一节的在线 softmax 解决它。

3.2 在线 softmax

数值稳定的 softmax 先减去最大值：对向量 $x$，$m(x) = \max_i x_i$，$f(x) = [e^{x_1 - m(x)}, \dots]$，$\ell(x) = \sum_i f(x)_i$，$\text{softmax}(x) = f(x) / \ell(x)$。

关键在于：把 $x$ 拆成两段 $x^{(1)}$、$x^{(2)}$，整体的统计量可以由两段各自的统计量算出来：

也就是说，只要记住到目前为止的最大值 $m$ 和以它为基准的和 $\ell$，就可以一块一块地看数据：新来一块，算出它的局部最大值和局部和；如果整体最大值变大了，就把旧的 $\ell$ 乘上 $e^{m_\text{旧} - m_\text{新}}$ 重新缩放，再加上新块的贡献。

代码 01 演示 1 手算 $x = [1, 3, 2, 5, 4]$，每次只看 2 个数：

最后 $e^{x_i - m} / \ell$ 与一次看完整行的 softmax 只差 $10^{-16}$。

注意力的输出 $O = \text{softmax}(s) V$ 也可以同样地累加：维护一个未归一化的输出累加 $\tilde O$，新块来了，旧的 $\tilde O$ 乘上同样的缩放因子，再加上新块的 $e^{s - m} V_\text{块}$；最后除以 $\ell$。这就是 FlashAttention 能把一整行注意力拆成很多块、逐块在片上完成的数学基础。（这个技巧本身更早就有：FlashInfer 论文把它归功于 Milakov 和 Gimelshein 2018 年关于在线 softmax 的工作，FlashAttention 论文也引用了先前分解 softmax 的工作。）

3.3 分块算法

FlashAttention（论文 Algorithm 1）把 $Q$ 切成大小为 $B_r$ 的行块，$K$、$V$ 切成大小为 $B_c$ 的块，块大小由 SRAM 的大小决定（$B_c = \lceil M / 4d \rceil$，$B_r = \min(\lceil M/4d \rceil, d)$），保证一块 Q、一块 K、一块 V 和中间结果能同时放进 SRAM：

for j in K/V 的每一块:                      # 外层：K/V 块
    把 K_j、V_j 从 HBM 读进 SRAM
    for i in Q 的每一块:                    # 内层：Q 块
        把 Q_i 和当前的 O_i、ℓ_i、m_i 从 HBM 读进 SRAM
        在片上：S_ij = Q_i K_jᵀ；按在线 softmax 更新 m_i、ℓ_i、O_i
        把 O_i、ℓ_i、m_i 写回 HBM

整个过程在一个融合的 CUDA 内核里完成：读入、矩阵乘、缩放、掩码、softmax、矩阵乘、写回，$S$ 和 $P$ 只在片上存在。因果掩码时，完全在对角线以上的块（全部被掩掉）直接跳过，计算量和读写都省一半。

训练时还有一个技巧：反向传播需要 $S$、$P$，FlashAttention 不存它们，而是在反向时用保存的 $m$、$\ell$ 重新算一遍（重算 recomputation）。多算了一些，但少了大量搬运，反向反而更快。推理只有前向，用不到这个技巧。

我们的 attn_kernels.flash_attention 实现了这个算法的两种循环顺序（见 §3.5）。代码 01 演示 2：长度 64、257、1000，几种块大小，含 / 不含因果掩码，FP32 输入下与标准注意力的最大差都在 $10^{-7}$ 量级。FlashAttention 只改变计算的顺序，不改变数学。

3.4 IO 复杂度

论文的定理 2：标准注意力需要 $\Theta(Nd + N^2)$ 次 HBM 访问，FlashAttention 需要 $\Theta(N^2 d^2 / M)$ 次。直觉是：每块 K/V 大小是 $\Theta(M)$，对每一块 K/V，都要把整个 Q 扫一遍（$\Theta(Nd)$），一共 $\Theta(Nd / M)$ 块，乘起来就是 $\Theta(N^2 d^2 / M)$。论文还证明了一个下界（命题 3）：对所有的 SRAM 大小，精确注意力都不可能渐近地做到更少的 HBM 访问。

$d^2 / M$ 有多小？论文说，$d$ 在 64–128、$M$ 约 100 KB 时，$d^2$ 比 $M$ 小好几倍。代码 01 演示 3 把常数也算上（SRAM 取 10 万个元素），先用计数器在小规模上验证公式与实现完全一致，再算到 $N = 32\text{K}$：

（FA1 是论文的循环顺序，FA2 是 FlashAttention-2 的循环顺序，见 §3.5。）

• $d = 64$ 时，FlashAttention 的读写比标准实现少约 8 倍。论文在 GPT-2（$d = 64$）上报告的读写量少了约 9 倍，与这一行吻合。
• $d = 128$ 时只少约 2 倍：$d$ 越大，SRAM 里一块能放的行越少，要扫的遍数越多。
• 读写的减少是一个常数倍，不是渐近的改进，两者都随 $N^2$ 增长。FlashAttention 更根本的好处是显存占用从 $O(N^2)$ 降到 $O(N)$（只需要存每行的 $m$、$\ell$），以及把多个 memory-bound 的内核融合成一个。
• 真实 GPU 上，被反复读的 K/V 块有相当一部分能命中 L2 缓存（H100 有 50 MB），实际的 HBM 读写比我们的记账更少。

代码 01 演示 4 验证了 $1/M$ 的关系：$N = 8192$、$d = 128$，SRAM 从 2.5 万增加到 40 万个元素时，FA1 的读写量从 536 M 降到 37 M，近似按 $1/M$ 下降。

3.5 从 FlashAttention-2 到 FlashAttention-4

FlashAttention 之后，同一个作者团队又发表了三个版本。每一版要解决的问题都来自新的硬件：

• FlashAttention-2（Dao，2023）：把两层循环交换，外层循环 Q 块，内层循环 K/V 块。这样 $O_i$、$\ell_i$、$m_i$ 一直留在片上，最后只写一次；不同的 Q 块交给不同的线程块，在序列长度这个维度上也能并行，线程块之间不需要同步。同时减少了非矩阵乘法的运算（Tensor Core 做矩阵乘法比普通运算快得多，每一次非矩阵乘法运算都很”贵”）。论文报告比 FlashAttention 快约 2 倍，在 A100 上达到理论峰值的 50%–73%。注意代码 01 演示 3：按我们的块大小规则，$d = 64$ 时 FA2 顺序的 HBM 读写量反而比 FA1 多（K/V 要被扫 $N/64$ 遍），可它在 GPU 上更快。读写量只是性能的一个方面，并行划分、片上同步、非矩阵乘法的开销同样重要。
• FlashAttention-3（Shah 等，2024）：针对 Hopper（H100）。FlashAttention-2 在 H100 上只能用到 35% 的峰值算力，而优化好的矩阵乘法能到 80%–90%。原因之一是没有用上 Hopper 的新特性：异步的 Tensor Core 和 TMA（张量内存加速器，专门搬数据的硬件单元）。FA3 用 warp 专门化（一部分 warp 只负责搬数据，一部分只负责计算）让计算和搬运重叠，让分块的矩阵乘法和 softmax 交错进行，并支持 FP8（块量化 + 用随机正交变换减小量化误差）。论文报告比 FA2 快 1.5–2.0 倍，FP16 达到 740 TFLOP/s（75% 利用率），FP8 接近 1.2 PFLOP/s，FP8 的数值误差比基线 FP8 注意力低 2.6 倍。
• FlashAttention-4（Zadouri 等，2026）：针对 Blackwell（B200 / GB200）。这一代硬件的扩展是不对称的：Tensor Core 的吞吐大约翻倍，可共享内存带宽和计算指数函数的单元几乎没有增长。于是 softmax 里的指数运算和数据搬运成了新瓶颈。FA4 用全异步的矩阵乘法流水线和更大的块，用软件模拟指数函数（在通用的乘加单元上用多项式近似，分担专用单元的压力），有条件地跳过不必要的 softmax 重缩放，并用 Blackwell 的张量内存和双 CTA 矩阵乘法模式减少共享内存的流量。它完全用 Python 里的 CuTe-DSL 编写。论文报告在 B200 上 BF16 最高 1613 TFLOP/s（71% 利用率），比 cuDNN 9.13 快最多 1.3 倍、比 Triton 快最多 2.7 倍，编译时间比传统的 C++ 模板方式快 20–30 倍。

四个版本合起来是一个很好的案例：同一个算法思想（IO 感知的分块 + 在线 softmax）不变，但每一代硬件都要重新设计内核，因为瓶颈每一代都在移动：先是 HBM 带宽，然后是并行划分，然后是异步搬运，然后是指数单元。

3.6 评测与论文卡片

论文报告的结果（训练为主，因为 2022 年的重点是训练）：BERT-large（序列长度 512）的训练比 MLPerf 1.1 的记录快 15%；GPT-2（序列长度 1K）快 3 倍；长序列基准 Long-Range Arena（1K–4K）快 2.4 倍。更长的上下文带来了更好的模型：GPT-2 的困惑度好 0.7，长文档分类提高 6.4 个点；第一次有 Transformer 在 Path-X（16K）和 Path-256（64K）上做到高于随机猜测（61.4%、63.1%）。论文还把思路推广到块稀疏注意力，比当时的近似注意力方法都快。

4. decode 阶段的注意力 · Attention in decode

FlashAttention 为 query 很多的情形（prefill、训练）而设计。decode 时每个请求每个头只有一个 query，情况完全不同。

4.1 算术强度：约等于组大小

考虑一个 KV 头和共用它的 $g$ 个 query 头（GQA 的组大小），一个请求这一步有 $l_{qo}$ 个 query、KV 长度为 $l_{kv}$：

• 运算：$4 g\, l_{qo}\, l_{kv}\, d$（$QK^\top$ 和 $PV$ 各 $2d$ 次）；
• 搬运：读 Q、写 O 各 $g\, l_{qo}\, d$ 个元素，读 K、V 各 $l_{kv}\, d$ 个元素。

算术强度化简后是

与 FlashInfer 论文给出的 $O\big(1 / (1/l_{qo} + 1/l_{kv})\big)$ 乘以 $g$ 一致。代码 02 演示 1（BF16，H100 的屋脊点约 295 FLOP/byte）：

三点结论：

1. decode 注意力的强度约等于组大小 $g$：每个 KV 元素读进来以后只被 $g$ 个 query 头各用一次。MHA 时约为 1，和 batch = 1 的线性层一样低。decode 注意力永远受限于带宽。
2. 批处理不能提高注意力的强度：每个请求的 KV 只属于它自己，批里请求再多，每份 KV 也只被它自己的 query 用（第 1 讲的”权重能分摊、KV 不能分摊”）。
3. 能提高强度的，是让更多 query 共用同一份 KV：GQA / MQA / MLA（模型结构；MLA 见第 8、10 讲）、一次验证多个 token（第 5 讲的投机解码）、多个请求共享前缀（§4.4）。

4.2 并行度的问题与 split-KV

decode 还有一个 prefill 没有的问题：活不够分。一个注意力内核通常让一个线程块负责一个（请求，KV 头）。批小的时候，这样的工作块只有几个，而 H100 有 132 个 SM。

代码 02 演示 2 用一个简化的模型（132 个 SM 平分显存带宽，一个工作块的时间 = 它要读的 KV 字节 ÷ 每个 SM 的带宽）估算 Llama-3-8B（8 个 KV 头）一层 decode 注意力的时间：

（括号里是 SM 利用率。）batch = 1、上下文 32K 时只有 8 个工作块，利用率 6%。split-KV 把每条 KV 切成几段，交给不同的 SM 并行处理，最后把各段的部分结果合并，利用率接近 100%，这一层快了 16 倍。这就是 Tri Dao 等人 2023 年在博客里提出的 Flash-Decoding。批大的时候本来就有足够的工作块，再切分只会多出合并的开销（切 64 段时 batch = 256 反而变慢）。切几段要根据批大小和上下文长度来定，这是 §5.6 负载均衡调度要解决的问题。

4.3 注意力状态与合并运算 ⊕

各段的部分结果怎样合并？FlashInfer 论文把它形式化成一个运算（思路来自 Blockwise Parallel Transformer 等工作）。对一个 query $q$ 和一个 key 集合 $I$，定义注意力状态

两个不相交的集合 $I$、$J$ 的状态可以合并成 $I \cup J$ 的状态：

$\oplus$ 满足结合律和交换律，所以多段可以按任意顺序合并。这和在线 softmax 是同一件事，只是把”一块一块地往前走”换成了”任意分组再合并”。FlashInfer 把注意力状态作为注意力运算的标准输出，把 $\oplus$ 当作标准的归约运算（就像矩阵乘法里的加法）。

代码 02 演示 3：一个 query、1 万个 key，切成 2、8、64 段。FP64 下合并结果与一次算完的差在 $10^{-16}$ 量级。可在 FP32 下，同样的段按正序合并和倒序合并，结果的最后几位不同（8 段时差 $7.5 \times 10^{-8}$，64 段时 $3.7 \times 10^{-7}$）：浮点加法不满足结合律。如果合并顺序取决于运行时的偶然因素（例如用原子加法，谁先算完谁先加），同样的输入就会得到略有不同的输出。这就是第 1 讲 §6.5 提到的”批不变性”问题在内核层面的根源，§5.6 会看到 FlashInfer 怎样避免它。

4.4 共享前缀：级联注意力

很多请求共享同一段前缀（系统提示、同一篇文档）时，朴素的 decode 注意力让每个请求各自把整段 KV 读一遍，前缀被读了 $n$ 次。级联注意力（cascade attention）把注意力拆成两部分：

1. 前缀部分：把 $n$ 个请求的 query 拼成一个 $n \times d$ 的矩阵，对前缀的 KV 一次算完。前缀只读一遍，而且变成了矩阵乘矩阵。
2. 独有部分：每个请求对自己的那一段 KV 单独算。

最后用 $\oplus$ 把两部分的状态合并。代码 02 演示 4（前缀 4096 + 每个请求独有 256）：

结果与逐个整段计算相同。前缀部分的 $n$ 个 query 共用一份 KV，相当于把组大小乘上了 $n$，强度从个位数提高到两百多，接近屋脊点。这需要 KV 存储能表达”前缀是共享的”：第 3 讲的写时复制和基数树正好提供了这个结构。

5. FlashInfer：为推理服务设计的注意力引擎 · FlashInfer

5.1 两个挑战

FlashAttention 解决了”怎样在片上算完注意力”，可推理服务对注意力内核还有别的要求。FlashInfer 论文把它们归结为两个挑战：

1. 负载多样而且动态：prefill、decode、增量 prefill（分块 prefill、多轮对话）的注意力形状不同；请求之间共享前缀；投机解码带来树形注意力；一批里请求的 query 数和 KV 长度各不相同，而且每一步都在变。朴素的实现会负载不均衡。
2. 硬件和模型都在变：KV 的存储方式多种多样（vLLM 的页表、SGLang 的基数树）；每一代 GPU 都需要专门的流水线和模板；新模型不断引入注意力的变体（GQA、特殊的掩码、logits 的变换……）。

当时的状况是每个推理系统各写一套注意力实现，只覆盖一部分情况，维护成本高，也常常不够快。FlashInfer 是一个基于代码生成的注意力引擎，用三个设计回应这些挑战：块稀疏格式、可定制的模板（JIT 编译）、负载均衡的动态调度。它已经被集成进 SGLang、vLLM 和 MLC-Engine。

5.2 块稀疏格式：统一各种 KV 布局

FlashInfer 的第一个观察是：页表就是一个块稀疏矩阵。把”query 行 × 物理 KV 块”看成一个大矩阵，请求 $r$ 用到了物理块 $b$，第 $r$ 行第 $b$ 列就是一个非零块。存稀疏矩阵的标准格式之一是 BSR（block compressed sparse row，块压缩稀疏行），它用两个数组：

• indices：把所有行的非零块的列号首尾相接；
• indptr：第 $r$ 行的非零块在 indices 里的范围是 [indptr[r], indptr[r+1])。

代码 03 演示 1 用第 3 讲的分页存储（block_size = 4，物理块打乱）为 3 个请求做 prefill，把块表写成 BSR：

请求 0：长度 19，块表 [13, 12, 15, 23, 0]
请求 1：长度 10，块表 [17, 3, 22]
请求 2：长度 29，块表 [8, 9, 18, 6, 1, 10, 2, 11]

BSR：indptr  = [0, 5, 8, 16]
     indices = [13, 12, 15, 23, 0, 17, 3, 22, 8, 9, 18, 6, 1, 10, 2, 11]

块表拼起来就是 indices。按 BSR 计算的注意力与”按块表拼成连续数组再算”的结果完全相同（最大差 0）。

一旦统一成 BSR，各种布局就只是 indptr / indices 不同：vLLM 的页表（块大小 $(1, 16)$）、SGLang 的基数树（$(1, 1)$，页大小为 1）、投机解码的树形注意力、只看部分 token 的稀疏注意力，都能用同一套块稀疏注意力内核处理。块的形状 $(B_r, B_c)$ 可以任意：$B_c$ 由 KV 的管理方式决定（页大小），$B_r$ 是一次处理几个 query 行。

实现上的难点是：Tensor Core 要求输入是规整的小矩阵（最小 16 行），可分页的 KV 散落在显存各处，块的大小（例如 1 个 token）也不一定对齐。FlashInfer 的做法是先按 indices 算出地址，把散落的 KV 行收集（gather）进连续的共享内存，再在共享内存上用普通的稠密 Tensor Core 指令计算。从共享内存开始，稀疏和稠密的实现就完全一样了。（Hopper 的 TMA 只支持仿射的地址模式，所以只有连续的 KV 才用 TMA，分页的 KV 退回到 Ampere 式的异步拷贝。）

5.3 可组合格式：共享前缀只读一遍

单一的块格式有一个限制：块越大（$B_r$ 越大），同一块里的几个 query 行越能共用读进共享内存的 KV，但不同请求的行很少恰好用同样的块；块越小越灵活，但每行都要单独去显存读。

FlashInfer 的可组合格式（composable formats，思路来自它的作者之前的 SparseTIR）是：同一个稀疏矩阵可以拆成几个子矩阵，每个子矩阵用最适合它的块格式。例如几个请求共享前缀时，“这几行 × 前缀的块”是一块稠密的子矩阵，用 $(n, B_c)$ 的大块存储，一次读进共享内存给 $n$ 个请求共用；各请求独有的部分仍用 $(1, B_c)$。两部分分别算出注意力状态，再用 $\oplus$ 合并。这其实就是 §4.4 的级联注意力。拆分不需要移动任何数据，只是为同一块显存多算一组索引。

代码 03 演示 2：4 个请求从同一个前缀 fork 出来（第 3 讲的写时复制让前缀的物理块天然被共享，块表的前 15 个块号完全相同），再追加各自的内容。单一格式下，读取的 KV 块数是 76；可组合格式（前缀用 $(4, B_c)$ 读一次）是 31；结果相同（最大差 $1.8 \times 10^{-7}$）。论文报告，在并行生成（多个请求共享前缀）的场景下，这类优化带来 13%–17% 的端到端加速。

5.4 JIT 模板：支持注意力的变体

新模型不断引入注意力的变体：

• logits soft-cap（例如 Gemma 2）：$s \to c \cdot \tanh(s / c)$，防止分数过大；
• 滑动窗口（例如 Mistral）：每个 query 只看最近 $w$ 个 key；
• sigmoid 注意力（FlashSigmoid）：不用 softmax，权重是 $\sigma(s + b)$；
• 把 RoPE（旋转位置编码）、归一化、投影直接融合进注意力内核……

为每一种变体手写一个高度优化的内核是不可持续的。FlashInfer 受 PyTorch 的 FlexAttention 启发，把注意力拆成一个固定的骨架（分块、在线 softmax、读写 KV、流水线）和几个可定制的函子（functor）：QueryTransform、KeyTransform、ValueTransform、OutputTransform、LogitsTransform、LogitsMask，还可以选择是否使用 softmax。用户用一小段 CUDA 代码写这些函子，JIT 编译器（即时编译，运行时才生成和编译代码）把它们填进模板，生成专门的内核。

代码 03 演示 3 用 numpy 演示这个思路：variant_attention 的骨架是分块 + 在线 softmax，用户只提供一两个小函数。三种变体（soft-cap、滑动窗口、sigmoid 注意力）的结果都与按定义直接实现的一致（差在 $10^{-14}$ 以内），而且都和普通注意力明显不同，说明不能拿普通内核凑合。论文的例子：给 StreamingLLM（一种用”注意力汇聚点 + 最近窗口”实现超长上下文的方法）生成一个把 RoPE 融合进注意力的内核，只需要约 20 行额外代码，端到端的 token 间延迟降低 28%–30%。

5.5 块大小的选择

FlashAttention-2 只提供少数几种块大小（例如 128 × 64），它们对 A100 上的 prefill 是最优的，对 query 很少的 decode 却很浪费：Tensor Core 指令一次处理至少 16 行，只有 1 个 query 时，其余行都在空转。FlashInfer 提供多种块大小：query 方向 1、16、32、64、128，KV 方向 32、64、128。它按”这一批平均每个请求有几个 query”选最小的够用的 query 块（GQA 时把一个 KV 头对应的 $g$ 个 query 头合并到 query 维度上，相当于 $g$ 个 query），再按寄存器和共享内存的限制选 KV 块，让 SM 的占用率最高。query 块大小为 1 时用普通的 CUDA 核心而不是 Tensor Core。论文的内核实验显示，正是这个灵活的块大小让 FlashInfer 的 decode 内核优于 FlashAttention 的。

5.6 负载均衡调度与 CUDA Graph

§4.2 看到，decode 注意力需要把长 KV 切开以获得足够的并行度，可切多少要看这一批的情况。FlashInfer 的负载均衡调度（论文 Algorithm 1，受矩阵乘法的 Stream-K 调度启发）：

1. 定义一个工作块的代价 $\alpha \cdot l_q + \beta \cdot l_{kv}$；
2. 根据这一批的总工作量算出每段 KV 的最大长度 $L_{kv} = \lceil \sum_i \lceil l_{qo}(i) / T_q \rceil \cdot l_{kv}(i) \,/\, \#\text{CTA} \rceil$，让每个 CTA 平均分到的工作量差不多；
3. 把每个请求的 KV 切成长度不超过 $L_{kv}$ 的段，每段是一个工作块；
4. 工作块按长度从大到小排序，用一个优先队列每次把下一个工作块分给当前总代价最小的 CTA（经典的”最长处理时间优先”贪心）；
5. 记下哪些部分结果属于同一个请求，算完以后用 $\oplus$ 合并（一个单独的”收缩”阶段）。

代码 04 演示 1 比较三种调度（每请求一块、固定切成 1024 一段、负载均衡）在长尾长度分布的批上的表现（Llama-3-8B 的 8 个 KV 头，132 个 CTA）：

每请求一块时，批小则活不够分，批大则最长的那个请求决定完成时间（长尾分布下，一个请求可以比平均长十几倍）。固定切分解决了一部分问题，但段长不随批变化。负载均衡调度在各种批上都把利用率提到八成到接近 100%。论文的内核实验（批大小 16，序列长度为常数、均匀分布、Zipf 偏斜分布）也报告：在均匀和偏斜分布下，FlashInfer 的内核明显优于 FlashAttention 的，原因正是这个调度。

两个容易被忽略的工程细节：

• 确定性：Stream-K 用原子加法合并部分结果，合并顺序取决于哪个 CTA 先算完。FlashInfer 刻意不用原子加法：调度只由序列长度决定，合并的顺序也就确定了，同样的输入总是得到逐位相同的输出。代码 04 演示 2：按固定顺序合并 3 次，输出逐位相同；按随机顺序合并（模拟原子加法），每次的最后几位都不一样（差 $7.5 \times 10^{-8}$）。可复现对调试、回归测试和 RL 训练（第 15 讲，推理和训练算出的概率要对得上）都很重要。
• CUDA Graph 兼容：CUDA Graph 把一整次 forward 的所有内核启动录下来、之后整体重放，能消除大量的 CPU 启动开销（第 1 讲的 overhead-bound），但它要求每次重放时内核的配置（网格大小、指针）都不变，可每一步的请求长度都在变。FlashInfer 的办法是：调度结果（每个 CTA 处理哪些块、部分结果怎样合并）每一步在 CPU 上算一次，异步拷到 GPU 上一块地址固定的工作区；注意力内核和合并内核都是持久内核（网格大小固定为 CTA 数，每个 CTA 去工作区读自己的任务）。这样，同一个 CUDA Graph 可以一直重放。同一步里所有层的序列长度相同，调度结果在所有层之间共用，开销被摊薄。代码 04 演示 3 在 CPU 上实测了 Python 版调度的耗时：批 256 时约 0.8 ms，摊到 32 层每层约 26 µs，只是这一步 decode GPU 时间的百分之几（真实的实现用 C++，快得多）。

5.7 评测与论文卡片

• 端到端（论文报告）：集成进 SGLang 后，与 SGLang 原来的 Triton 后端相比，在 Llama 3.1 8B（1 块 H100）和 70B（4 块 H100）的服务基准上 token 间延迟降低 29%–69%；长上下文（StreamingLLM，融合 RoPE 的内核）延迟降低 28%–30%；并行生成（可组合格式）加速 13%–17%。
• 内核：在均匀和偏斜的长度分布下优于 FlashAttention（负载均衡）；decode 内核优于 FlashAttention（灵活的块大小）。

6. 之后的发展 · Since then

（这一节是我的补充，不是原课内容。）

• 注意力后端成了推理引擎的可插拔部件：vLLM 和 SGLang 都支持多种注意力后端（FlashInfer、FlashAttention 2/3、Triton 实现、厂商库），按模型和硬件选择。
• 为新的注意力结构写专门的内核：DeepSeek 在 2025 年开源了 FlashMLA，一个面向 MLA（多头潜在注意力，把 KV 压缩到低维潜在空间）decode 的 Hopper 内核；第 10 讲的稀疏注意力、线性注意力也都需要专门的内核才能把理论上的节省变成实际的加速。
• 编程方式的变化：FlexAttention（PyTorch）让研究者用几行 Python 写注意力变体，自动生成融合内核；FlashAttention-4 用 Python 里的 CuTe-DSL 编写，编译快得多。“用高层语言写高性能内核”正在成为趋势。
• 硬件继续变化：Blackwell 这一代 Tensor Core 的增长快于其他单元（§3.5），注意力里非矩阵乘法的部分（指数运算、归约）越来越重要。低精度（FP8、FP4）注意力也在发展，需要特别处理数值误差。

7. 小结 · Summary

1. 注意力是推理里最难做快的算子：prefill 时中间结果按 $N^2$ 增长，decode 时它的主要工作是读 KV。GPU 有两级存储：HBM 大而慢，片上 SRAM 小而快（快一个数量级）。softmax 这类算子永远受限于搬运，解法是算子融合。
2. 标准注意力把 $N \times N$ 的分数矩阵写回、读出好几遍，HBM 读写 $\Theta(Nd + N^2)$，显存 $O(N^2)$。Llama-3-8B 的 32K prefill 若用标准注意力要慢 2.5 倍，而且显存放不下（代码 01 演示 5）。
3. FlashAttention 的关键是 IO 感知：用在线 softmax（维护 $m$、$\ell$，逐块缩放、累加）把注意力拆成小块，在一个融合内核里算完，$S$ 从不写回 HBM。它是精确的；IO 是 $\Theta(N^2 d^2 / M)$，并有匹配的下界；按实际常数，HBM 读写在 $d = 64$ 时少约 8 倍、$d = 128$ 时少约 2 倍，更根本的是显存从 $O(N^2)$ 降到 $O(N)$（代码 01）。
4. FlashAttention-2/3/4 每一代都在追赶新的硬件瓶颈：并行划分、异步搬运与 FP8、不对称的硬件增长。算法思想不变，内核要随硬件重写。
5. decode 注意力的强度约等于 GQA 组大小，批处理提高不了它；批小、上下文长时并行度不够，split-KV（Flash-Decoding）把 KV 切段并行、再用 ⊕ 合并（代码 02：batch = 1、32K 上下文快 16 倍）。⊕ 在数学上精确，在浮点上与合并顺序有关。共享前缀时，级联注意力让前缀只读一遍（64 个请求省 93%）。
6. FlashInfer：页表就是块稀疏矩阵，BSR 统一了各种 KV 布局；可组合格式让共享前缀只读一遍；JIT 模板 + 函子支持注意力变体；确定性的负载均衡调度在长尾的批上把利用率提到八成到接近 100%，并通过持久内核与固定工作区兼容 CUDA Graph（代码 03、04）。

8. 常见误区 · Common pitfalls

1. “FlashAttention 是一种近似注意力。” 不是。它是精确注意力，只改变计算顺序（代码 01 演示 2：与标准实现的差在浮点误差范围内）。
2. “FlashAttention 减少了计算量。” 不减少（不算因果掩码时跳过的块），训练时反向的重算还会多算一些。它减少的是 HBM 的读写和显存占用。
3. “FlashAttention 让注意力的复杂度变成线性的了。” 计算量仍是 $O(N^2 d)$，HBM 读写仍随 $N^2$ 增长，只是常数小得多；变成线性的是额外的显存。
4. “读写量少的实现就一定更快。” FlashAttention-2 的循环顺序在某些参数下读写量更多，在 GPU 上却快约 2 倍：并行划分、同步、非矩阵乘法的开销同样重要。
5. “批处理能提高 decode 注意力的效率。” 不能。批处理分摊的是权重；每个请求的 KV 只给它自己用，注意力的强度只取决于有几个 query 共用一份 KV（组大小、投机解码的 token 数、共享前缀的请求数）。
6. “把 KV 切得越碎越好。” 切分是为了凑够并行度；活已经够分时，再切只会增加合并的开销（代码 02 演示 2：batch = 256 切 64 段反而变慢）。
7. “分块计算的结果和不分块完全一样，所以输出总是可复现的。” 数学上一样，浮点上不一定：合并顺序不同，最后几位就不同。要可复现，合并顺序必须是确定的（代码 02 演示 3、代码 04 演示 2）。
8. “有了 FlashAttention，注意力就不再是瓶颈了。” 长上下文的 decode 仍然要把整段 KV 读一遍；这正是第 9 讲（KV 放到哪里）和第 10 讲（让 KV 变小、只读一部分）要继续解决的问题。

9. 练习 · Exercises

练习 6 的参考答案里的数字是用代码 02 的函数实际算出来的。

练习 1（手算在线 softmax）。$x = [2, 1, 4]$，每次只看一个数。写出每一步的 $m$ 和 $\ell$，并用最后的 $m$、$\ell$ 算出 softmax。

看到 2：$m = 2$，$\ell = e^0 = 1$。看到 1：最大值不变，$\ell = 1 + e^{1-2} = 1.3679$。看到 4：最大值变大，旧的 $\ell$ 先乘 $e^{2-4} = 0.1353$，得 $0.1851$，再加 $e^{4-4} = 1$，$\ell = 1.1851$，$m = 4$。softmax $= [e^{-2}, e^{-3}, e^{0}] / 1.1851 = [0.1142, 0.0420, 0.8438]$。直接算：$e^{x - 4} = [0.1353, 0.0498, 1]$，和 $1.1851$，结果相同。

练习 2（显存）。Llama-3-8B 有 32 个 query 头、32 层。用标准注意力（FP16）做 32K 的 prefill，一层所有头的 $S$ 矩阵一共多大？FlashAttention 需要的额外显存（每个 query 行存 $m$、$\ell$ 两个 FP32）呢？

一个头：$32768^2 \times 2$ 字节 ≈ 2.15 GB；32 个头 ≈ 68.7 GB，加上 $P$ 矩阵还要翻倍，一块 80 GB 的 H100 放不下（还要放权重和 KV）。逐头计算可以省显存，但要多次启动内核、更慢。FlashAttention：每个头每行 2 个 FP32，$32768 \times 2 \times 4$ 字节 = 256 KB，32 个头共 8 MB。

练习 3（核对 IO 公式）。用 §3.3 的块大小规则（SRAM = 10 万个元素，$d = 64$），手算 $N = 4096$ 时 FA1 的 HBM 读写量，与代码 01 演示 3 的 9.4 M 比较。

$B_c = \lceil 100000 / 256 \rceil = 391$，$B_r = \min(391, 64) = 64$，$T_c = \lceil 4096/391 \rceil = 11$。FA1 = 读 K、V 一次（$2 \times 4096 \times 64 = 524{,}288$）+ $T_c$ 遍、每遍对 $N$ 行读 $Q_i, O_i, \ell_i, m_i$、写回 $O_i, \ell_i, m_i$（$4096 \times (3 \times 64 + 4) = 802{,}816$）$\times 11 = 8{,}830{,}976$，合计约 9.36 M，与表中的 9.4 M 一致。标准实现约 68.2 M，比值约 7.3。

练习 4（decode 注意力的强度）。(a) Llama-3-70B（$g = 8$）、上下文 8192 时，decode 注意力的强度是多少？(b) DeepSeek 的 MLA 让一层的 128 个 query 头共用同一份压缩后的潜在 KV（可以粗略地看成 $g = 128$），强度变成多少？离 H100 的屋脊点还有多远？

用 §4.1 的公式 $I = g \cdot l_{kv} / (g + l_{kv})$（$l_{qo} = 1$，BF16 下单位是 FLOP/byte；代码 02 的 attn_intensity 给出同样的数）：(a) $g = 8$ 时 $8 \times 8192 / 8200 \approx$ 8.0 FLOP/byte；(b) $g = 128$ 时 $128 \times 8192 / 8320 \approx$ 126 FLOP/byte，提高了十几倍，但仍低于屋脊点 295，decode 注意力依然受限于带宽，只是读同样多的字节能服务多得多的 query 头。MLA 的细节（潜在向量的维度、解压的计算）见第 10 讲。

练习 5（为什么推理用不到重算）。FlashAttention 在训练的反向传播里”重算”$S$ 和 $P$ 而不是把它们存下来。为什么推理时不需要这个技巧？推理时 FlashAttention 还需要保存什么？

推理只有前向，没有反向传播，不需要 $S$、$P$ 来算梯度，所以既不用存也不用重算。前向结束时只需要输出 $O$；如果注意力被拆成几部分（split-KV、级联注意力、分块 prefill 的跨块合并），还要保留每部分的 LSE 以便用 ⊕ 合并。

练习 6（切几段）。用代码 02 的 decode_time，求 batch = 2、上下文 16384、8 个 KV 头时，把 KV 切成 1、2、4、8、16、32、64 段的时间和利用率。最优的段数是多少？为什么？

最优是 8 段：$2 \times 8 \times 8 = 128$ 个工作块，刚好不超过 132 个 SM，一轮就做完；再切下去工作块超过 132 个，要分两轮，每块也更小，总时间几乎不变，只多了合并的开销。经验法则：让”请求数 × KV 头数 × 段数”略多于 SM 数（或它的整数倍）。负载均衡调度（§5.6）自动做的就是这件事，而且还考虑了长度不同的请求。

练习 7（级联注意力的账）。8 个请求共享 1000 个 token 的前缀，每个请求独有 50 个 token。朴素做法和级联做法各要读多少个 token 的 KV？节省多少？

朴素：$8 \times 1050 = 8400$；级联：$1000 + 8 \times 50 = 1400$；节省 $1 - 1400/8400 \approx 83\%$。前缀部分的强度相当于组大小乘以 8。

练习 8（为什么 FlashAttention-4 要软件模拟指数函数）。§3.5 说 Blackwell 的 Tensor Core 吞吐大约翻倍，计算指数函数的专用单元却几乎没有增长。用 roofline 的思路解释，为什么这会让 softmax 成为新的瓶颈？FlashAttention-4 的两个对策各在解决什么？

注意力的矩阵乘法部分（$QK^\top$、$PV$）跑在 Tensor Core 上；softmax 里每个分数都要算一次指数，跑在专门的特殊函数单元上。矩阵乘法变快了一倍、指数没变快，“每个分数要做的指数运算”就从被掩盖的小开销变成了关键路径。这就像 roofline 里多了第三根屋顶：一个只有指数单元吞吐的上限。对策一：用通用的乘加单元做多项式近似来计算一部分指数，把指数运算分摊到两类硬件上。对策二：有条件地跳过 softmax 的重缩放（只有最大值变化足够大时才缩放旧的结果），减少非矩阵乘法的运算。这正是”内核要随硬件的不对称变化重新设计”的例子。

练习 9（思考：块稀疏格式还能表达什么）。§5.2 说页表、基数树、树形注意力都能写成块稀疏矩阵。(a) 第 3 讲 SGLang 的基数树缓存，用 BSR 表示时块大小是多少？(b) 一个只看”最开头 4 个 token + 最近 1024 个 token”的注意力（StreamingLLM 式），它的掩码怎样用块稀疏表示？

(a) SGLang 的页大小是 1 个 token，所以块大小是 $(1, 1)$（每个非零块是一个 query 行 × 一个 token），indices 就是每个 token 在显存池里的位置。这种”向量级稀疏”对 Tensor Core 不友好，FlashInfer 的做法是先把这些 token 的 KV 收集进共享内存，再用稠密指令计算。(b) 每个 query 行的非零块是”前 4 个 token 所在的块”加上”最近 1024 个 token 所在的块”。块大小选得和窗口对齐时，非零块是规整的；不对齐时，边界上的块用掩码函子（LogitsMask）处理部分位置。这样一个”注意力汇聚点 + 滑动窗口”的注意力，计算量和读写量都与总上下文长度无关。

10. 延伸阅读 · Further reading

本讲的必读论文

• T. Dao, D. Y. Fu, S. Ermon, A. Rudra, C. Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
• Z. Ye et al. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. MLSys 2025.

选读

• T. Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023.
• J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, T. Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. 2024.
• T. Zadouri, M. Hoehnerbach, J. Shah, T. Liu, V. Thakkar, T. Dao. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. 2026.

相关工作

• M. Milakov, N. Gimelshein. Online Normalizer Calculation for Softmax. 2018（在线 softmax）。
• T. Dao, D. Haziza, F. Massa, G. Sizov. Flash-Decoding for Long-Context Inference（博客，2023）。
• PyTorch 团队. FlexAttention（博客与文档，2024）：用 Python 写注意力变体、自动生成融合内核。
• M. Osama et al. Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU. 2023（FlashInfer 负载均衡调度的灵感来源）。

背景

• 前置讲 P1（GPU 与性能的第一性原理）：SM、warp、共享内存、Tensor Core、算子融合。

11. 下一讲预告 · Next lecture

到这里，单卡推理引擎的骨架已经齐了：连续批处理（第 2 讲）、分页 KV 与前缀复用（第 3 讲）、高效的注意力内核（本讲）。可 decode 的根本限制还在：每生成一个 token 都要把权重和 KV 读一遍，而这一遍读出来的数据只为每个请求产出一个 token，算力大量闲置。第 5 讲：推测解码读两篇论文：Speculative Decoding（Leviathan 等）用一个小模型先”猜”几个 token，再让大模型一次 forward 验证它们，并用一种巧妙的拒绝采样保证生成的分布与大模型完全相同；SpecInfer 把猜测扩展成一棵 token 树，用树形注意力一次验证多条候选。我们会在迷你模型上实现推测解码（包括第 3 讲 KVCache.truncate 预留的”回滚”），在小词表上精确验证”分布不变”，并在纸上 GPU 上算出接受率和加速比的关系。本讲的 §4.1 已经埋下伏笔：一次验证 8 个 token，注意力的强度提高了 8 倍。