Attention 流水线：单 Q 双 S 与双 Q 的计算重叠

发表于 分类于 阅读次数： 本文字数： 6.7k 阅读时长 ≈ 24 分钟

Attention 前向由两次矩阵乘和一段 softmax 交替组成。串行实现里，Tensor Core 算完 QK 要等 softmax，softmax 做完又要等 PV，两类运算单元轮流出空。把不同数据块上的工作交错排开，让一边等待时另一边继续算，是这个算子提高吞吐的主要来源。

在单个 CTA 内安排这种交错，有两种常见做法。单 Q 双 S 从同一个 query 块的相邻 KV 轮次里找重叠；双 Q 把两个 query 块放进同一个 CTA，让两条计算链交替推进。两者要回答的问题相同：哪些数据能提前算，哪些状态必须按顺序更新，多出来的中间结果放在哪里，为此占用的片上空间又是否划算。

本文假设读者了解 FlashAttention-2 的 dense 前向计算和在线 softmax，讨论范围限于单个 CTA（CUDA 线程块）内的流水线。硬件背景以 Hopper 和 Blackwell 的异步矩阵指令为主，片上存储的例子按 Blackwell 的口径给。文中的 softmax 阶段包括 max、exp、sum、类型转换和 P 的交接，exp 只是其中一段。时间线里的格子是相对时间单位，字节数是布局算例，都不是某个 kernel 的实测值。

一、一轮 attention 的工作量与硬件分工

一次 QK、一次 PV、一段 softmax

固定一个 query 块 Q，沿序列方向依次读 K₀/V₀、K₁/V₁……。这里的 tile 是矩阵块；“单 Q”指一个 query 块，可以包含 128 个 token。

第 j 轮的主要工作是：

Sj=QKjTD,Pj=exp⁡(Sj−mj),Uj=αjUj−1+PjVjS_j=\frac{QK_j^{\mathsf T}}{\sqrt D},\qquad P_j=\exp(S_j-m_j),\qquad U_j=\alpha_jU_{j-1}+P_jV_j

m 是每行的运行最大值，α 是最大值变化时对已有累加结果使用的重缩放系数。U 是尚未归一化的输出累加器；ℓ 是这一行已处理的全部 KV 块的指数和，以当前 m 为基准在线累积并重缩放，最后用 O=U/ℓ 逐行归一化。

下文把 QKᵀ 简写为 QK，把 P×V 简写为 PV。

QK 与 PV 都落到 MMA（Matrix Multiply-Accumulate，矩阵乘加）上。softmax 不同，它由逐元素运算、最大值与求和归约、指数计算组成，跑在通用运算单元和 SFU 上。

一轮 KV：数据从左向右流动
① QK → S矩阵乘加Tensor Core
② S → Pmax · exp · sum通用运算单元 + SFU
③ P × V矩阵乘加Tensor Core
另有输出重缩放、数据搬运和同步；图中先突出两次 MMA 与 softmax。
图 1：同一份 P 必须等 S 算出后才能生成，PV 也必须等 P 准备好。可重叠的工作要从其他轮次或其他 Q 中寻找。

指数一般用 exp2 实现，因为 exp(x)=2^(x·log₂e)，硬件指数路径通常编译到 MUFU.EX2。MUFU／SFU 是执行这类特殊函数的单元，吞吐上限和 Tensor Core 完全不同。不过一段源码里的 exp 未必只对应一条机器指令，具体数学函数可能额外包含几条。FlashAttention-3 对硬件分工的解释、PTX 的 ex2 定义

运算量差 512 倍，时间却可能一样

设 Q、K、V 的块尺寸都是 128×128，head dimension D=128，一次乘加按 2 FLOPs 计：

每轮工作 运算量
QK 2 × 128³ = 4,194,304 FLOPs
PV 2 × 128³ = 4,194,304 FLOPs
逐元素 exp 128² = 16,384 次

两次 MMA 的 FLOPs 是 exp 次数的 512 倍。这个比值本身没有多少意义，FLOPs 和指数次数不是同一种量，能比的只有各自除上硬件吞吐之后的时间。

按 FlashAttention-4 设计说明中的 B200 每 SM 理想吞吐口径，BF16 矩阵运算为 8,192 FLOPs/cycle，指数为 16 次/cycle。把上面的工作量除下去：

TMMA≈8,388,6088192=1024 cycles,Texp⁡≈16,38416=1024 cyclesT_{\mathrm{MMA}}\approx\frac{8{,}388{,}608}{8192}=1024\ \text{cycles},\qquad T_{\exp}\approx\frac{16{,}384}{16}=1024\ \text{cycles}

两次 MMA 合计约 1024 周期，指数运算也约 1024 周期，量级上正好相当。这个估算只做了工作量除以峰值吞吐，没有计入访存、归约、重缩放、同步和指令调度，不能当实际耗时。FlashAttention-4 的吞吐分析

低精度优化里的一个现象可以由此理解：MMA 加快之后，exp 的时间未必同比缩短，它在总耗时中的占比反而上升。

二、串行安排：两类硬件轮流空闲

最直接的写法是每轮做完 QK 做 softmax，做完 softmax 做 PV，然后进入下一轮。按上一节的比例，一轮里 QK 与 PV 合起来约等于一段 softmax 的长度，把每次 QK、PV 各记 1 格，E 记 2 格，轮流出空的位置就看得出来。

串行安排：只画前两轮 KV
时间 →
0
1
2
3
4
5
6
7
Tensor Core
Softmax
QK₀
PV₀
QK₁
PV₁
E₀
E₁
QK / PV：矩阵乘E：含 exp 的 softmax 阶段灰色空档：等待
图 2：人为设定每次 QK、PV 各占 1 个时间格，E 占 2 格。格子是相对时间单位，省略搬运与输出校正，不是周期级测量。

图 2 里 E₀ 执行期间，Tensor Core 在等 P₀。但 QK₁ 只需要 Q 和 K₁，不需要 P₀ 或 U₀，本来可以提前算出来。

要让这种提前发生，还需要异步指令。Hopper 的 wgmma.mma_async、Blackwell 的 tcgen05.mma 属于这一类：发起矩阵运算之后，线程可以继续推进其他允许的工作，到读取结果时再确认完成。源码里把两段代码相邻写出并不会自动带来并行，这些异步能力才是前提。FlashAttention-3 论文的异步流水线、PTX 的异步 MMA 完成机制

有了异步指令，还要有人把不同的工作分头推进。常见做法是 Warp Specialization（Warp 分工）：不同 Warp 或 Warp 组分别负责加载、发起 MMA、处理 softmax 和校正输出，通过缓冲区和屏障交接数据。这样同一个 SM 上，异步矩阵硬件和处理逐元素运算的线程就能交错工作。Hopper 与 Blackwell 在发射粒度、分组方式和存储位置上都有差别，不能把一代的角色分配直接搬到另一代。

图中的 E 一直是完整的分数处理阶段：从可以读取的 S，到可以交给 PV 的 P。能重叠的是这一段工作与另一个 tile 的 MMA；其中 exp、归约、类型转换和写 P 各能藏起来多少，取决于实际调度和资源争用。

只要 S 有第二处可写，QK₁ 就能在 E₀ 期间启动。剩下的问题是怎么管这个多出来的槽，以及输出状态要怎样跟着更新。

三、单 Q 双 S：让下一轮的分数提前算

两份 S，一个 Q

如果只留一个 S 槽，softmax 还在读 S₀ 时，QK₁ 没有地方写新结果。再加一个槽，问题就解决了。

单 Q 双 S（single-Q + S ping-pong）保留一个 Q、两份逻辑分数缓冲、一套在线 softmax 与输出状态。两个槽交替接收相邻轮次的分数：

同一个 Q，沿 KV 方向往前推进
Q固定的 128 行 query
S 槽 A第 0、2、4……轮QK₀ → QK₂ → QK₄
S 槽 B第 1、3、5……轮QK₁ → QK₃ → QK₅
一套状态每行 m、ℓ输出累加器 U
A / B 是存储位置；0 / 1 / 2 是 KV 轮次。两个槽反复使用，不保留完整注意力矩阵。
图 3：两份 S 都服务于同一个 query 块。图中的“每行”很关键：128 行各有自己的 max 和 sum，不把整块分数混成一个 softmax。

下面按一种保守实现来画：S 和 P 复用同一区域，必须等 PV 完成才能释放槽。

单 Q 双 S：E₁ 与 PV₀、QK₂ 重叠
时间 →
0
1
2
3
4
5
6
7
8
9
Tensor Core
Softmax
QK₀
QK₁
PV₀
QK₂
PV₁
QK₃
PV₂
PV₃
E₀
E₁
E₂
E₃
图 4：沿用图 2 的相对时长，这次画四轮。S 槽 A 在 PV₀ 完成后才交给 QK₂；槽 B 在 PV₁ 完成后才交给 QK₃。省略输出校正阶段的耗时，不据此推算实测加速比。

按图走一遍：

  1. QK₀ 写槽 A，完成后 softmax 开始处理。
  2. softmax 处理槽 A 时，QK₁ 写槽 B。
  3. P₀ 就绪，做 PV₀；softmax 转去处理槽 B。
  4. PV₀ 用完槽 A，QK₂ 才能写回槽 A。此后不断交替。

重叠发生在 E₀ 与 QK₁ 之间，以及 E₁ 与 PV₀、QK₂ 之间。同一轮的 QK、E、PV 仍然要守数据依赖。

分数可以早算，输出必须按顺序校正

QK₁ 能早算，但如果它带来更大的行最大值，之前累积的结果就要重新缩放：

mj=max⁡(mj−1,rowmax⁡(Sj)),αj=exp⁡(mj−1−mj),Pj=exp⁡(Sj−mj),ℓj=αjℓj−1+rowsum⁡(Pj),Uj=αjUj−1+PjVj.\begin{aligned} m_j&=\max\left(m_{j-1},\operatorname{rowmax}(S_j)\right),\\ \alpha_j&=\exp(m_{j-1}-m_j),\\ P_j&=\exp(S_j-m_j),\\ \ell_j&=\alpha_j\ell_{j-1}+\operatorname{rowsum}(P_j),\\ U_j&=\alpha_jU_{j-1}+P_jV_j. \end{aligned}

这些运算逐行进行，α 也按行广播。举一行的例子：旧最大值是 2，新最大值变成 4，旧分母和旧输出累加器都要乘 exp(-2)，才能和本轮贡献处在同一尺度上。

分数可以提前算；同一 U 的更新必须协调
上一轮 PV 完成得到完整 Uⱼ₋₁
本轮 max 已知得到 αⱼ
两项均就绪 ↓
先校正旧 U，再加入本轮 PVUⱼ = αⱼ Uⱼ₋₁ + Pⱼ Vⱼ
图 5：采用原地累加时，不能一边让上一轮 PV 写 U，一边无序地重缩放同一份 U。实现可以安排独立的校正（correction）角色，或使用其他经过验证的累加方案。

所以双 S 放宽的只是分数缓冲的读写冲突，max、sum 和输出累加器各自保留依赖链。普通有限分数可以从第一块初始化；整行都被 mask 的情况要单独处理，避免 -∞-(-∞) 产生 NaN。

槽什么时候能还回去

实现里需要能表达三件事：QK 已写完 S，softmax 已写好 P，PV 已用完 P。如果 S 和 P 共用同一区域，让 softmax 读完 S 就释放整个槽，下一次 QK 会覆盖 PV 还没用完的 P。

协议可以概括成三句话：写之前等空闲，读之前等就绪，用完之后再归还。循环使用时还要区分槽里的数据属于第几代，常见做法是维护 phase。初始值和屏障怎么用，取决于所用的流水线 API，不能只照抄 j % 2。

Blackwell 上还要分清 MMA 的提交与完成。tcgen05.commit 配合 barrier 跟踪异步操作的结束，访问结果时必须满足对应的可见性规则。PTX 的 tcgen05 同步规则

四、双 Q：两条独立的计算链

单 Q 双 S 只在同一个 Q 的前后轮次之间找重叠。双 Q（dual-Q）换一个方向：让一个 CTA 同时处理 Qᴬ、Qᴮ 两个 query 块，各自保留 m、ℓ、U。

例如 Qᴬ 对应第 0～127 行 query，Qᴮ 对应第 128～255 行。dense 非因果 attention 里，它们遍历同一段 KV，但两组分数按各自的行归一化，最后分别写出 Oᴬ 和 Oᴮ。

一个 CTA，两个 query 块
Qᴬ：第 0～127 行Sᴬ / Pᴬmᴬ、ℓᴬ、Uᴬ → Oᴬ
Qᴮ：第 128～255 行Sᴮ / Pᴮmᴮ、ℓᴮ、Uᴮ → Oᴮ
同一个 KV 块可以服务两路；两路的 max、sum、输出分别维护。
时间 →
0
1
2
3
4
5
6
7
8
9
Tensor Core
Softmax
QKᴬ₀
QKᴮ₀
PVᴬ₀
QKᴬ₁
PVᴮ₀
QKᴮ₁
PVᴬ₁
PVᴮ₁
Eᴬ₀
Eᴮ₀
Eᴬ₁
Eᴮ₁
图 6：上标 A／B 表示 query 块，下标 0／1 表示 KV 轮次。共完成四个 Q-KV 配对，与图 4 的四个配对工作量相同。时间图仍省略搬运和输出校正；相似的交错节奏来自不同的数据组织。

Eᴬ₀ 运行时，Tensor Core 可以算 QKᴮ₀；Eᴮ₀ 运行时，又可以做 PVᴬ₀ 和 QKᴬ₁。A 路的输出缩放不依赖 B 路的最大值，B 路也不必等 A 路走完全部 KV 才开始。

与单 Q 双 S 相比，双 Q 多出一套独立的输出状态，给调度留了更多选择。dense 情况下，同一块 K/V 如果留在片上供两路共用，还可能省掉一次重复加载；能不能真的复用，要看搬运顺序、布局和缓冲的生命周期。

FlashAttention-4 公开的前向设计就采用每 CTA 两个 Q tile 的 ping-pong，并配一个独立的校正阶段处理输出重缩放。它还会协调两组 softmax 的 exp 阶段，减少对同一个指数单元的争抢。这里利用的是 exp 与 MMA 的互补；反过来，两路 exp 同时跑，也会竞争同一份吞吐。FlashAttention-4 的前向流水线

两种设计可以叠加，但要付出资源代价

对比项 单 Q 双 S 本文图示的双 Q
一个 CTA 负责的 query 块 1 个 2 个
可交错的工作来自哪里 同一 Q 的不同 KV 轮次 两个 Q 的工作交错
逻辑分数槽 同一 Q 使用 A／B 两槽 每个 Q 各一槽
在线 max／sum／输出状态 一套，按 query 行维护 两套，各自按行维护
主要资源增量 额外的 S／P 在途数据 额外的 Q、输出和状态

“几个 Q”和“每个 Q 有几级 S 缓冲”是两个独立的参数。双 Q 也可以再给每个 Q 加分数缓冲，代价是更多片上存储和同步。所以只凭名字判断不了某个 kernel 的流水线深度。

除了计算，还有搬运

TMA（Tensor Memory Accelerator，张量存储加速器）能把后续的 K/V 块异步送进片上，搬运本身也可以和计算重叠。把这几类放在一起，各自的作用就清楚了：

可重叠的阶段 为什么有机会并行 仍要满足的条件
本轮 softmax 与下一轮 QK 下一轮 QK 不依赖本轮 P K 已到达，且有空闲 S 槽
本轮 softmax 与上一轮 PV 上一轮 P 已就绪 上一轮输出校正已完成，输入未被覆盖
TMA 预取后续 K/V 与当前计算 搬运走异步数据通路 有空缓冲，计算前确认传输完成
一路输出校正与另一路 MMA，或后续 QK 操作的是不同数据区域 校正与逐元素运算仍可能争资源

前两行是前面几条时间线的主角。后两行同样影响实际吞吐，但 TMA 预取属于搬运与计算的重叠，单 Q 双 S 和双 Q 改变的只是计算之间可交错的部分。输出校正和 PV 如果访问同一份 U，仍要按图 5 的依赖关系排队。

这里的双 Q 都在一个 CTA 内，和前文的 2CTA 矩阵协作 是两类不同的组织方式。Warp 与 SM 的对应关系可以参看 CUDA 执行模型。

五、多一条计算链的代价

双 Q 多留了一份 Q、一份输出暂存和一套在线状态，同时把两路 query 绑在同一个 CTA 里。两件事都会反过来影响吞吐。

片上空间会压缩预取深度

寄存器保存线程正在操作的分数片段、max、sum 和地址；SMEM（共享内存）常用来放 Q/K/V 和输出暂存；Blackwell 的 TMEM（Tensor Memory）存放矩阵累加结果，也可以承接 S／P。这三处各有容量上限，不能互相顶替。

先算一笔简化的 SMEM 布局预算。取 128×128 的 FP8 输入块和 BF16 输出暂存，假设可用于这部分布局的空间是 224 KiB，每个 query 块占 16 KiB，每份输出暂存占 32 KiB，每个 K 或 V 缓冲槽占 16 KiB：

相同 224 KiB 预算，容纳的数据不同
单 Q：16 KiB Q + 32 KiB 输出暂存 + 11 × 16 KiB KV 槽
Q
O
11 个 KV 槽
双 Q：32 KiB Q + 64 KiB 输出暂存 + 8 × 16 KiB KV 槽
Q × 2
O × 2
8 个 KV 槽
图 7：条形长度按字节数绘制。KV 槽按单份 K 或 V 计，11 槽不表示 11 对 K/V。数字仅为布局算例，省略对齐、屏障等开销；输出暂存与 TMEM 中的 FP32 累加器也属于不同存储。

双 Q 多了可交错的计算，却可能把 KV 预取距离从 11 个槽压到 8 个。如果 MMA 此后更频繁地等数据，计算重叠省下的时间就会被抵消。

TMEM 也要算。按每块 128 列的逻辑布局，单 Q 双 S 的 S₀ + S₁ + U 占 384 列，双 Q 的 Sᴬ + Sᴮ + Uᴬ + Uᴮ 占 512 列，都假设 P 复用 S 的区域。常规、非 exclusive 的 TMEM 分配按 2 的幂申请列数，384 列的布局也可能申请到 512 列。把剩余的逻辑空间利用起来当然可以，但那会带来额外的寄存器、SMEM 或同步开销。PTX 的 TMEM 分配规则

并行度集中到更少的 CTA

一个 CTA 处理两个 Q，覆盖同样一批 query 需要的 CTA 数大约减半。batch 小、query 总块数本来就不多的时候，GPU 可能凑不出足够的任务铺满 SM。寄存器、SMEM 或 TMEM 用量上升，还会压低同时驻留的 CTA 数。

资源预算变小，不等于一定发生寄存器溢出。活跃值会不会被写到 local memory，要看编译产物和性能计数，不能因为用了双 Q 就下结论。

两路长度不同，会多出空轮次

不加 mask 的时候，两路遍历的 KV 块数相同。一旦有 causal mask、变长序列或稀疏清单，长度就可能不一样。

比如 A 路有 3 个有效 KV 块，B 路只有 1 个。实现如果为了配对而让两路都跑 3 轮，4 个有效配对就变成 6 个实际配对。被 mask 掉的额外计算数学贡献为零，却照样花掉加载和 MMA 时间。能跳过多少，取决于实现。

所以双 Q 应该和已经优化过的单 Q 双 S 比，并为两者分别选合适的 tile、stage 和资源配置。只赢过完全串行的基线，不足以说明它优于现有的单 Q 流水线。

六、重叠最多能省多少

把一轮里两次 MMA 占 Tensor Core 的时间记作 Tₘ，softmax 占相关运算单元的时间记作 Tₑ，其他成本暂时忽略。完全串行需要：

Tserial=Tm+TeT_{\mathrm{serial}}=T_m+T_e

如果两类工作能在不同 tile 之间充分重叠，稳态下的理想处理间隔只能接近：

Toverlap≳max⁡(Tm,Te),speedupideal≤Tm+Temax⁡(Tm,Te)≤2T_{\mathrm{overlap}}\gtrsim\max(T_m,T_e),\qquad \mathrm{speedup}_{\mathrm{ideal}}\leq \frac{T_m+T_e}{\max(T_m,T_e)}\leq2

两边各需 1 个时间单位时，理想间隔从 2 降到 1；MMA 需要 1、softmax 需要 3 时，充分重叠也只能把 4 降到 3，剩下的要靠减少 softmax 的工作量或提高它的有效吞吐。

图 4 和图 6 画的工作量相同，都是四个 Q-KV 配对、合计 16 格有效工作，各自走完用了 10 格。离 2 倍还差的那部分，是模型里的启动、排空和数据依赖。

这个“最多 2 倍”只属于“固定工作量、两类阶段、与完全串行相比”这一个模型，没有算 KV 复用之类的其他优化，也不是从单 Q 双 S 换到双 Q 的加速承诺。真实 kernel 还要付出数据移动、校正、屏障等成本，资源变化又会改变这些成本。

要判断一个实现有没有真的受益，可以依次看三件事。先看结果对不对：和参考输出比对，覆盖多个 KV 轮次、尾块、mask 和极端分数，确认每个 Q 的 max、sum 和输出都正确。再看等待有没有减少：把 Tensor Core、指数路径、访存和屏障的等待分开统计，确认省下的时间没有转移到别的瓶颈上。最后看整个算子：固定 GPU、dtype、形状和精度要求，用设备计时重复测量，比较中要包含加载、归一化和最终写出。

单 Q 双 S 让一个 query 块的下一轮分数提前就位，双 Q 再多开一条独立的 query 计算链。两者起作用的前提相同：一类硬件卡在当前依赖上时，手边恰好有另一份输入齐备的工作。收益多少，取决于这些工作能不能及时就绪，以及把它们留在片上要花多少空间。

技术资料核对日期：2026 年 10 月 28 日。图示均为本文绘制；吞吐参数取自文中链接的公开设计说明，其余数字都是在文中假设下的估算。