极客工场

Appearance

共 5,374 字预计阅读约 18 分钟

Attention 算子优化：从物化 N×N 矩阵到 FlashAttention 与 PagedAttention ​

更新时间：2026-09-25。前序文章 神经网络推理的算子分解 把一次推理拆成了算子账单，GPU 内存层级与带宽 解释了 HBM、L2、SRAM 之间的数量级落差，Tensor Core 与矩阵加速 讲了算力本体 GEMM 怎么跑满。本篇把这三条线拧到同一个 kernel 上：Transformer 的注意力层。它只占模型参数的一小部分，却长期是训练和长序列推理里最难优化的算子——原因不在算不动，而在中间结果太大、搬来搬去把时间全花在了 HBM 上。

本文要回答的问题 ​

  • Scaled Dot-Product Attention 的 QKᵀ 和 softmax(·)V 两段计算，维度和运算量各是多少？
  • 朴素实现为什么必须把 N×N 的注意力分数矩阵写到 HBM 再读回来？为什么说它是访存密集而非计算密集？
  • online softmax 是什么？怎样只凭一个 running max 和一个 running sum，就把 softmax 拆成分块单遍完成？
  • FlashAttention 的分块循环按什么顺序走？HBM 流量怎样从 O(N²) 降到 O(N²d²/M)？反向传播为什么要重算？
  • GQA/MQA 怎样压缩 KV Cache？PagedAttention 解决的是什么形态的访存与碎片问题？INT8/FP8 KV 量化又能再省多少？

一、先算清楚注意力这笔账：QKᵀ 与 softmax(·)V ​

标准的 Scaled Dot-Product Attention（缩放点积注意力）写成一行公式：

O = softmax(QKᵀ / √dₖ) · V

其中 Q、K、V 三个矩阵的形状都是 N×d：N 是序列长度（token 数），d 是每个注意力头的维度（head dimension，常见 64 或 128）。公式虽短，kernel 里实际是三段独立计算。第一段 S = QKᵀ/√dₖ，是两个 N×d 矩阵相乘得到 N×N 的分数矩阵 S，运算量约 2N²d 个 FLOP。第二段 P = softmax(S)，对每一行做「减最大值—指数—求和—归一化」，逐元素操作，运算量只有约 O(N²)，但它是串行归约：一行内必须先知道整行最大值，再算指数和分母。第三段 O = P·V，N×N 乘 N×d，又是约 2N²d 个 FLOP。

把数字代进去感受形状。N=4096、d=128 时，Q、K、V 各约 52 万个元素（FP16 下各 1MB），而中间的 S 和 P 是 N×N ≈ 1670 万个元素，FP16 下各 32MB——一个头就 32MB，几十个头叠起来，单层注意力的中间结果以 GB 计。更关键的是两端增长速度不同：计算量是 O(N²d)，中间数据量是 O(N²)。当 d 固定（它确实基本固定在 64~128）而 N 从 2048 涨到 32768，N 涨 16 倍，注意力矩阵的元素数涨 256 倍。长上下文场景下，注意力层就是被这张二次方膨胀的矩阵压垮的。

用一个真实配置核对总规模。一个 80 层、80 个查询头、头维度 128 的模型，N=8192 时单层 MHA 的注意力分数元素数是 80×8192² ≈ 54 亿，FP16 下超过 10GB——这还只是一层的一个中间矩阵，训练时若全部保存，显存根本装不下。算子账单里 FFN 的 FLOP 往往比注意力多，但 FFN 是普通 GEMM、算术强度高、没有巨型中间结果，优化成熟；注意力之所以成为专项难题，难点全部来自这张矩阵的形状，而不是算得少。

多头（MHA）只是把这套结构复制 h 份：Q、K、V 被 reshape 成 h 个 N×(d/h) 的小头并行计算，头之间互不通信。理解优化时盯住一个头即可，所有招法都是对单头形状做的，头数只影响并行规模，不改变单头的访存结论。

二、朴素实现：算子边界把 N×N 矩阵逼进 HBM ​

在深度学习框架里，最直白的实现就是把三段写成三个算子：一个 GEMM 算子产出 S，一个 softmax 算子逐行处理，再一个 GEMM 算子乘 V。每个算子各自启动、各自从 HBM 读输入、把完整输出写回 HBM。这在框架工程上完全正确，却恰好踩中了内存层级最痛的地方。

chip-design/expert/transformer-attention-kernel.md D2 示意图

图里的路径就是朴素 kernel 的真实数据流：S 被完整写出、被 softmax 完整读回，P 再被完整写出、被第二个 GEMM 完整读回。N×N 矩阵在 HBM 上进出至少四趟，而 Q、K、V 这些 N×d 的小矩阵反而只占零头。粗算一下 HBM 流量：S 和 P 各读写一次，约 4N² 个元素（FP16 下 8N² 字节），而两段 GEMM 的计算量只有 4N²d FLOP。A100 的 HBM 带宽约 1.5~2TB/s，FP16 矩阵算力（含 Tensor Core）约 312 TFLOPS，把 8N² 字节除以带宽、4N²d FLOP 除以算力，会发现 N=4096、d=128 时搬运时间和计算时间几乎同量级——算术强度被算子边界压到了约 d/2 个 FLOP/字节，恰好在 roofline 的脊点附近甚至更低，Tensor Core 大量时间在等数据。

这就是「attention 是访存密集算子」的确切含义：不是它没有矩阵乘，而是中间矩阵必须落 HBM，搬运规模随序列长度二次方膨胀。顺带一个连锁伤害：训练时要为反向传播保存 S 和 P，N×N 的中间显存直接限制了能跑的最大序列长度和 batch size。优化的靶子因此非常明确——让这张 N×N 矩阵永远不要离开片上 SRAM。

还有一笔常被忽略的小账：为了让矩阵形状对齐，batch 内不等长的序列要做 padding（补零到同一长度），朴素实现连补齐部分的 QKᵀ、softmax、·V 也照算不误，再用 mask 把无效位置在 softmax 前置为负无穷。N=2048 的 batch 里若实际长度平均只有 1024，算力和带宽各浪费约四分之三。这也是后面 tiling 方案「变长友好」这一项的现实来源。

三、Online softmax：把 softmax 改成可以分块的单遍算法 ​

阻碍分块的是 softmax 本身。一行的标准算法是：先求整行最大值 m 做数值稳定（减去 m 再取指数，防止 exp 溢出），再求指数之和 l，最后归一化。最大值和求和都是对整行的归约，看不到全行就拿不到 m 和 l——看起来必须等 QKᵀ 算完一整行才能往下走。

online softmax（流式 softmax）打破了这个依赖。核心是维护两个递推量：到当前块为止的行最大值 m 和指数和 l，每来一个新块就更新它们，并对旧结果做 rescale（重新缩放）。

chip-design/expert/transformer-attention-kernel.md D2 示意图

按图走一遍递推。设已经处理完前 j−1 个 K 块，手里有 m_old、l_old 和未归一化的输出累加 O_old。读入新块算出 S_ij 后，先取局部行最大值，合并出 m_new = max(m_old, rowmax(S_ij))。基变了，旧块所有指数都要乘以修正因子 e^(m_old−m_new)：分母变成 l_new = e^(m_old−m_new)·l_old + Σ exp(S_ij−m_new)，输出累加同样乘修正因子后再加上新块与 V_j 的乘积。等最后一块处理完，用 l_new 做一次归一化，得到的结果与「看完整行再算」逐位等价（仅有浮点舍入差异）。

用三个数验证递推确实自洽。假设一行分数分两块到达：第一块是 [1, 2]，第二块是 [3, 1]。处理第一块：m_old=2，l_old=e^(1−2)+e^0≈1.368。第二块出现更大的 3，m_new=3，修正因子 e^(2−3)≈0.368；l_new=0.368×1.368+e^0+e^(1−3)≈0.503+1+0.135≈1.638。而一次性对 [1,2,3,1] 减最大值 3 求指数和：e⁻²+e⁻¹+1+e⁻²≈0.135+0.368+1+0.135≈1.638，两数完全一致。旧输出块乘同一修正因子后，归一化结果也与一次性算法相同。

这个递推的价值在于：softmax 不再需要物化整行分数，最大值、分母、输出三个量都能随块增量更新，每块数据只读一次。它不是 FlashAttention 发明的（online softmax 的数值技巧此前已在 Milakov & Gimelshein 2018 等工作中出现），但 FlashAttention 第一次把它和矩阵乘分块严丝合缝地焊进了同一个 GPU kernel。注意 m 只会单调增大，修正因子恒不大于 1，数值上是稳定收敛的——这也是该递推可以放心放进 kernel 内层循环的原因。

四、FlashAttention：tiling、循环顺序、IO 复杂度与反向重算 ​

把分块落到内存层级上，就是 FlashAttention 的全部骨架。Q、K、V 按序列维切成能放进片上 SRAM（GPU 上即 shared memory + 寄存器）的小块：外循环遍历 K、V 的块，内循环遍历 Q 的块；每一对 (Q_i, K_j) 在 SRAM 内完成局部 QKᵀ、online softmax 更新、与 V_j 的乘加重缩放，累加结果 O_i 留在寄存器里持续更新。S 和 P 的分块（尺寸约 B_r×B_c，典型几百）生存在 SRAM，用完即弃，从头到尾不进 HBM。Q、K、V 各从 HBM 读一次，O 写回一次。

IO 复杂度的结论值得记住：设片上 SRAM 容量为 M，朴素实现的 HBM 流量是 Θ(N²)（注意力矩阵强制进出），FlashAttention 降到 Θ(N²d²/M)——分块越大（M 越大），每块 K/V 被复用的 Q 块越多，流量越小；额外显存从 Θ(N²) 降到 Θ(N)（只需存 O 和 m、l 等量）。这是 kernel fusion 与 tiling 联手的标准战果，与 Tensor Core 与矩阵加速 里 GEMM 分块的算术完全同源，只是 softmax 的存在让递推复杂了一层。

维度朴素三段实现FlashAttentionPagedAttention
HBM 流量Θ(N²)，S/P 各进出 HBMΘ(N²d²/M)，中间块不出 SRAM推理时 KV 按页读取，随机但聚合度高
中间显存Θ(N²)，需物化注意力矩阵Θ(N)，无 N×N 物化无注意力矩阵；KV Cache 按页分配
变长序列padding 浪费算力与显存按实际长度分块，天然支持为变长并发而生，按需分页
主要瓶颈HBM 带宽SRAM 容量、softmax rescale 开销页表间接访问、decode 阶段带宽
主战场短序列、框架原型训练与长上下文 prefill高并发自回归推理/decode

实际收益随场景变化，公开的典型量级是：训练侧端到端提速约 1.2~2 倍（纯 attention kernel 上可达 2~4 倍），显存从二次方变线性后可支撑的序列长度拉长数倍（GPT 类模型从 1k~2k 推到 4k~16k 量级）；FlashAttention-2 改进并行划分与工作分区后，A100 上 attention FLOPs 利用率（相对理论峰值）提升到约五成到七成，约为 v1 的两倍。序列越长、朴素实现越被带宽卡死，tiling 的剪刀差越大；N 很短（几百以内）时分块调度的固定开销反而会侵蚀收益。

循环顺序还有一个工程细节值得点破：K/V 在外层、Q 在内层，是因为这样 K_j、V_j 载入 SRAM 后能被所有 Q 块依次扫过，每块 K/V 的复用次数等于 Q 的块数；反过来则 K/V 要被反复载入。因果语言模型里还要加 causal mask（每个 token 只能看它之前的位置），分块后只需让块内乘加跳过未来位置，约一半的块整体在 mask 外可以直接跳过，计算和访存再省一半。FlashAttention-2 把工作按 Q 块在 warp 间重新切分、减少 warp 间通信与同步，FlashAttention-3 则针对 Hopper 的异步流水线和 Tensor Core 编排做指令级重叠，思路一脉相承：让 rescale、指数等非 GEMM 操作藏到矩阵乘的等待间隙里。

反向传播有个反直觉的设计：FlashAttention 前向不存 N×N 矩阵，反向需要 softmax 的中间统计量时，干脆用保存的 Q、K、V 把注意力分数重算一遍，而不是在前向时写回 HBM。重算增加了约一倍计算，但省掉的是二次方访存——在带宽受限的屋顶下这笔交易非常划算，本质是用过剩的算力换昂贵的带宽。训练吞吐因此同时受两件事支配：前向块大小要匹配 SRAM 容量，重算比例要在激活保存策略里登记清楚，这与梯度检查点（gradient checkpointing）是同一类取舍。

五、GQA/MQA 与 PagedAttention：推理侧的 KV Cache 账 ​

训练和 prefill（一次性处理整段输入）的优化集中在注意力矩阵；自回归 decode（逐 token 生成）阶段瓶颈换了位置：每生成一个新 token，都要拿它的 Q 去和之前所有 token 的 K、V 做注意力，而这些历史 K、V 构成 KV Cache，常驻 HBM。decode 时矩阵乘退化为「矩阵×向量」，计算量极小，每个 token 都要把整个 KV Cache 读一遍，是纯粹的带宽受限场景。于是优化目标从「别物化 N×N」变成「让 KV Cache 更小、读得更经济」。

MQA（Multi-Query Attention，多查询注意力）最激进：所有查询头共享唯一一组 K、V 头，KV Cache 直接降到 MHA 的 1/h，训练收敛性和生成质量在部分任务上会有可观察到的回退。GQA（Grouped-Query Attention，分组查询注意力）是折中：查询头分成若干组，每组共享一对 KV 头，比如 Llama-2 70B 用 64 个 Q 头配 8 个 KV 头，KV Cache 缩到 1/8，质量损失明显小于 MQA。当代主流大模型几乎都采用 GQA。每 token 的 KV Cache 字节数可以直接按 2 × 层数 × KV头数 × 头维度 × 每元素字节 估算，把「头数」这一项从 h 改成 g，就是 GQA 省掉的部分。

给一个 decode 带宽账的量级。N=4096、80 层、8 个 GQA 头、头维度 128、FP16 时，每个新 token 要读的 KV Cache 约为 2×80×8×128×4096×2 字节 ≈ 1.07GB；而这个 token 的注意力计算只有约 2×N×d 量级 FLOP，摊下来算术强度低到个位数，在 2TB/s 的 HBM 下每个 token 仅 KV 读取就耗时约 0.5ms 量级。上下文越长，单 token 越慢，这就是「长对话越聊越迟钝」的硬件原因，也是下面两类压缩手段的直接动机。

PagedAttention（vLLM 提出）解决的是另一个维度的问题：显存管理。传统推理框架为每个请求预分配连续的最大长度 KV 空间，实际序列长短不一、增长不定，产生大量内部碎片，KV Cache 有效利用率常只有两三成，并发能力被锁死。PagedAttention 把 KV Cache 切成固定大小的页（block，通常对应十几到几十个 token），用逻辑块到物理块的页表间接寻址，按需分配、不同请求间还可共享相同前缀（如系统提示词）的物理页。碎片趋近于零后，同一张 GPU 能并发的请求数可提升数倍，代价是注意力计算多一层页表间接寻址、kernel 访存模式变为离散聚集读取。页大小本身是个折中：页选大了则内部碎片回升、页选小了则页表项数量和调度次数都会膨胀，十几到几十个 token 一页是多数实现的经验区间。配合 continuous batching（连续批处理），这一帧跑完的请求立刻让出物理页、新请求当帧插入，GPU 不必等最慢的序列，decode 集群的吞吐利用率因此可以逼近带宽上限。它和 FlashAttention 不冲突：prefill 阶段用 tiling kernel，decode 阶段在分页的 KV 上做块读取，两者管的是生成流程的不同段落。

六、INT8/FP8 KV 量化：在带宽受限处再砍一半 ​

既然 decode 是带宽受限，缩小 KV Cache 的每个元素就直接等于提速。KV Cache 量化把缓存的 K、V 从 FP16 降到 INT8 或 FP8（E4M3 为主），容量和读取流量直接减半，与 GQA 的头数缩减还是乘法叠加：64 Q 头配 8 KV 头（1/8）再叠加 FP8（1/2），相对 MHA + FP16 的基线只占 1/16。这对能塞下多少并发、多长上下文往往是决定性的。

工程上的做法通常是 per-token 或 per-channel 缩放（存一组 scale，注意力计算时按需反量化或在低精度域累加），缓存侧单边量化即可拿到大部分收益，不要求权重和计算全量化。精度敏感的场景里，一个常见策略是只量化 V、保留 K 为 FP16，或对异常值通道做混合精度——softmax 对 K 的尺度变化相对更敏感。INT8 的硬件覆盖最广，FP8 则在 Hopper 及更新的 GPU、新一代 NPU 上有原生计算与访存通路，逐渐成为长上下文服务的默认配置。精度风险主要集中在两类数据：长序列里少数 K 通道会出现大动态范围的异常值（outlier channel），统一 scale 下容易被量化台阶吃掉；softmax 后的 P 也不宜低精度，因为指数分布本身尖锐。实践中常见的口径是「KV 缓存 FP8、softmax 与累加 FP32、GEMM 主体 FP16/BF16」，各段取各自性价比最高的精度，而不是全模型一刀切。FlashAttention 家族的较新版本也把 FP8 KV Cache 与分块 online softmax 合并进了同一 kernel：量化发生在块写回时，反量化发生在 SRAM 内，HBM 上流动的始终是压缩表示。

至此可以看到注意力优化的统一逻辑：先识别此刻被什么卡住——prefill 卡物化矩阵的访存，就 tiling + online softmax；decode 卡 KV Cache 带宽，就 GQA 减头、分页减碎片、低精度减字节——再让数据在尽可能靠近计算单元的存储层级里被反复使用。这与 GPU 内存层级与带宽 的分析框架完全一致，attention kernel 只是把所有招法浓缩到了一个算子内。

这一篇在知识体系里的位置 ​

本篇处于「AI 芯片与算子优化」支线的深水区。上游的算子账单见 神经网络推理的算子分解，内存层级与带宽墙见 GPU 内存层级与带宽，矩阵乘怎样吃满 Tensor Core 见 Tensor Core 与矩阵加速，片上数据流复用的硬件对应物见 AI 加速器架构：脉动阵列、数据流（WS/OS）、NoC 与内存层级。再往上走是长上下文系统架构、投机解码与推理服务调度，本篇提供的是它们共同依赖的 kernel 级账本。

总结 ​

Attention 算子优化的主线就是不让 N×N 注意力矩阵落到 HBM：FlashAttention 用分块 online softmax 的 running max/sum 单遍递推，把 softmax 焊进 QKᵀ 与 ·V 的 SRAM 分块循环，HBM 流量从 Θ(N²) 降到 Θ(N²d²/M)、中间显存降到线性并以反向重算换带宽；推理侧再用 GQA/MQA 减 KV 头、PagedAttention 分页消碎片、INT8/FP8 压缩 KV Cache，逐层拆掉 decode 阶段的带宽墙。