FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor-Vector Parallelism

发表时间: 2026-01

Jianxing Xu, Yuanbo Wen, Jun Bi, Ruibai Xu, Guanglin Xu, Rui Zhang, Wei Li, Ling Li, Tianshi Chen, Qi Guo, Yunji Chen
University of Science and Technology of China; Institute of Computing Technology, Chinese Academy of Sciences; Cambricon Technologies.

速读

一句话结论 本文提出了 FlashAttention-T，通过复用 GPU 的张量单元来执行 Softmax 计算，解决了注意力机制中因等待向量单元而导致的计算空闲瓶颈，在保证精度的前提下将吞吐量最高提升了 1.19 倍。

要解决什么问题 现有的融合注意力机制（如 FlashAttention-2 和 FlashAttention-3）在计算时将矩阵乘法和 Softmax 操作解耦，矩阵乘法交给算力强悍的张量单元（Tensor Cores）执行，而 Softmax 则交给速度较慢的向量单元（CUDA Cores）执行。随着硬件架构的演进，张量单元的吞吐量翻倍增长（例如从 Ampere 架构到 Hopper 架构），但向量单元的性能提升并没有跟上。这种硬件异构性导致了一个严重的机制卡点：向量间隔（Vector Interval）。具体来说，张量单元在极快地算完矩阵乘法后，必须停下来等待向量单元慢吞吐地完成 Softmax 的一系列操作（求最大值、乘加、指数运算等），才能进入下一轮迭代。在 FlashAttention-3 的实际运行中，这种等待导致张量单元的空闲时间占到了整个迭代周期的 36.3%，大量的算力被白白浪费，严重限制了注意力机制的整体并行度和执行效率。

怎么做的 核心思路是将 Softmax 的关键计算原语卸载到处于空闲状态的张量单元上，让张量单元和向量单元真正并行起来。为了绕开张量单元原本只能做矩阵乘法的限制，作者设计了一套操作数赋值方法，通过巧妙构造输入片段（fragments）的数据排布，复用 GPU 的张量矩阵乘累加（MMA）指令来执行 Softmax 中的逐元素缩放、融合乘加和行和归约。针对张量单元内部数据映射带来的置换开销，作者将其转化为最小化交换次数的优化问题，直接在寄存器内低成本恢复正确结果；对于行和归约，则利用 MMA 输出的部分和进行两次线程内加法，彻底消除了传统向量化方法中显式的线程间同步和通信开销。然而，张量 MMA 指令强制要求整行数据必须共享同一个缩放因子，这与标准 Softmax 每行缩放因子不同的数学定义相冲突。为此，作者设计了张量化在线 Softmax 算法，引入了代理最大值（Surrogate Maximum）$\hat{m}[i]$，将其定义为包含第 $i$ 行的 $X$ 行图块内的最大值：$$ \hat{m}[i] = \max_{k=\lfloor i/X \rfloor \cdot X}^{\min((\lfloor i/X \rfloor + 1) \cdot X, n)} m[k] $$ 因为代理最大值始终大于等于该行真实的局部最大值，这在数学上严格保证了后续的指数计算不会发生数值溢出。最后，由于全张量化在当前硬件上并非绝对最优，作者设计了架构感知的混合调度策略。在 Ampere 架构上采用指令级并行，在同一个 Warp（线程束）内将张量指令和向量指令交错发射，让向量指令填补张量指令的发射气泡；在 Hopper 架构上采用线程级并行，受限于异步指令的寄存器依赖，仅将注意力分数的行和归约操作张量化，并将其调度到下一次迭代的矩阵乘法批次中，与其他线程块的向量化操作完全重叠。

效果如何 实验在服务器级 GPU（NVIDIA A100 和 H100）以及边缘端 GPU（Jetson AGX Orin）上进行，固定总 Token 数为 16K，测试了 512 到 8K 的序列长度。对比基线点名了四种主流方案：代表 Ampere 架构强基线的 FlashAttention-2、代表 Hopper 架构强基线的 FlashAttention-3、专为大语言模型推理优化的 FlashInfer，以及 OpenAI 深度学习编译器生成的 Triton Kernel。在吞吐量表现上，FlashAttention-T 对比 FlashAttention-2 和 FlashAttention-3 实现了 1.02 到 1.19 倍的加速，对比 FlashInfer 和 Triton 平均加速 1.03 到 1.39 倍。这种加速在非因果配置和较小的头维度（如 64）下尤为明显。在解决核心卡点方面，Hopper 架构上的线程级并行调度将向量间隔比率大幅压缩至 2.7%，几乎消灭了张量单元的空闲等待。该方法的局限在于，引入代理最大值后，如果遇到极端的长尾数据分布，可能会增加指数项全部下溢为零的概率，此时系统需要回退到传统的向量化计算。但在 Llama3、Mistral 和 Qwen3 等真实模型的生成式基准测试中，FlashAttention-T 的功能正确性指标与基线完全一致，未观测到任何数值失败，证明了其在实际应用中的可靠性。

A1 主要贡献

本文针对现有的融合注意力机制（Fused Attention，如 FlashAttention-2/3）中存在的向量间隔（Vector Interval）瓶颈问题，提出了一种向全张量化注意力机制演进的方案 FlashAttention-T。主要贡献归纳如下：

1. 核心问题识别与解决：现有的融合注意力实现将矩阵乘法（GEMM）和 Softmax 操作在计算上解耦，前者使用高性能张量单元（Tensor Cores），后者使用较慢的向量单元（CUDA Cores）。这种性能异构性导致了严重的向量间隔，即张量单元在等待向量单元完成 Softmax 时处于空闲状态（如图 1(a) 所示）。FlashAttention-T 通过将关键的 Softmax 原语卸载到空闲的张量单元上来解决这一瓶颈。
2. 基于张量 MMA 的 Softmax 原语：提出了一系列操作数赋值方法，通过重新利用（Repurpose）GPU 的张量矩阵乘累加（MMA）指令来执行 Softmax 的关键原语（如逐元素缩放、融合乘加和行和归约），最大限度地利用硬件并减少开销。
3. 张量化在线 Softmax 算法：设计了一种具有数值稳定性保证的张量化在线 Softmax 算法。该算法引入了代理最大值（Surrogate Maximum）的概念，以适应被重新利用的张量 MMA 指令对统一缩放因子的约束，从而使得在 FlashAttention 流程中可以张量化 Softmax 计算。
4. 张量-向量并行调度：提出了架构感知的调度技术，通过利用指令级并行（ILP，针对 Ampere 架构）和线程级并行（TLP，针对 Hopper 架构），将张量化 Softmax 算法并行映射到张量和向量单元上，充分利用异构并行性。

A3 背景知识/关键Observation/设计原则

融合注意力的算法流程
融合注意力通过将计算并行化到 GPU 内核中的 Warpgroup（或 Threadblock）来执行，每个 Warpgroup 计算输出矩阵 $O$ 的一个行块。如 Fig 2(a) 所示，算法流程包含：❶ 计算注意力 Logit 块 $S$；❷ 更新行最大值 $m$；❸❺❻ 基于 $m$ 和旧的行最大值 $m_{old}$ 对 $S$ 进行指数化生成注意力分数块 $\tilde{P}$；❹ 重新缩放注意力输出 $O$；❼ 更新注意力分数行和 $l$；❽ 将结果累加到 $O$。这种分块方案确保中间结果驻留在寄存器中，为利用张量 MMA 指令进行张量化提供了基础。

Softmax 时间占比分析
在当前的范式中，Softmax 操作（max, fma, mul, add, exp）使用向量单元，而 GEMM 使用张量单元。对于序列长度 $N$ 和头维度 $h$，Softmax 的总操作数为 $4N^2 + \frac{N-1}{N}N^2h$，而 GEMM 为 $2N^2h$ FMAs。随着硬件发展（如 Hopper H100 相比 Ampere A100 吞吐量翻倍），张量单元与向量单元的吞吐量比率 $P_T/P_V$ 增加，导致 Softmax 的时间占比显著上升，成为计算瓶颈。

向量间隔瓶颈的揭示
通过分析 FlashAttention-2（Ampere 优化）和 FlashAttention-3（Hopper 优化）的执行时间线（见 Fig 3），作者揭示了向量间隔瓶颈：
* FlashAttention-2 (Fig 3a)：采用顺序调度，GEMM 计算后必须等待向量单元完成 Softmax 才能进行下一次 GEMM，导致张量单元空闲时间 $t_{vec}$。
* FlashAttention-3 (Fig 3b)：利用 Hopper 的异步 WGMMA 实现流水线调度，部分重叠了 GEMM 和 Softmax。然而，未重叠的 Softmax 计算仍导致显著的向量间隔。
* Table 1 的 Profiling 结果显示，向量间隔在 FlashAttention-2 中平均占迭代周期的 29.8%，在 FlashAttention-3（FP8-FP32）中升至 36.3%。

A2 方法细节

重新利用张量 MMA 指令

操作数赋值原理与映射
现代 GPU 的张量 MMA 指令将逻辑操作数矩阵划分为分布在线程间的片段（fragments）。定义 $v[i, j]$ 为逻辑矩阵 $v$ 在 $i$ 行 $j$ 列的元素，$v(u, t)$ 为线程 $u$ 持有的片段中的第 $t$ 个元素。映射函数 $\phi_v: v[i, j] \to v(u, t)$ 是特定于架构和操作数类型的（如图 Fig 4 所示）。为了重新利用 MMA 指令执行 Softmax 原语，需要确定输入片段的赋值，使得 MMA 输出片段 $D(v, t)$ 等于输入片段 $A$ 的置换后的逐元素缩放。

逐元素缩放与融合乘加（Scaling & FMA）
对于缩放因子 $\alpha$，目标是确定片段 $B$ 的赋值，使得 MMA 计算出的输出片段 $D$ 满足 $D(v, t) = \alpha A(\sigma(v), t)$，其中 $\sigma$ 是由张量单元内部片段映射引起的不可避免的置换（Permutation）。
* 置换开销最小化：为了恢复非置换输出（即 $D(v, t) = \alpha A(v, t)$），需要在线程寄存器内进行交换（swap）操作，这会产生开销。作者将此问题建模为约束优化问题：最小化置换 $\sigma$ 的 Cayley 距离（即将其转换为恒等置换所需的最小交换次数）。
* 最优赋值：Fig 5(a) 展示了最优的片段 $B$ 赋值方法，对于 $\sigma = (1\ 2)$，只需一次交换即可恢复。
* 融合乘加扩展：基于此，通过利用张量单元的累加器添加偏移量 $C(v, t)$，该方法可扩展为实现 $D(v, t) = \alpha A(\sigma(v), t) + C(v, t)$ 的融合乘加操作。

行和归约（Row-Sum Reduction）
行和归约涉及对注意力分数矩阵 $\tilde{P}$ 的行进行求和。传统的向量化方法需要线程内加法和显式的线程间 All-Reduce，带来同步和通信开销。作者提出利用张量 MMA 指令更高效地执行此操作：
* 输入构造：设置输入片段 $A$ 来自矩阵 $\tilde{P}$，将片段 $B$ 设置为如 Fig 5(b) 所示的特定模式，并将累加器 $C$ 置零。
* 输出计算：MMA 输出片段 $D'$ 本质上包含了部分和。通过执行两次线程内加法 $s(0, t) = D'(0, t) + D'(2, t)$ 和 $s(1, t) = D'(1, t) + D'(3, t)$，即可得到最终的归约结果。
* 优势：这种方法避免了显式的线程同步和通信，且置换 $\sigma$ 被吸收到线程内加法中，不产生额外的交换开销。

讨论
这些方法不需要额外的数据拷贝开销，因为它们直接在寄存器中操作。主要限制是算法约束：缩放和 FMA 要求 MMA 指令的所有行共享统一的缩放因子 $\alpha$。这需要后续的张量化算法来配合。

张量化在线 Softmax 算法

统一缩放因子的挑战与代理最大值
原始融合注意力算法中的输出重缩放步骤 $O = \exp(m_{old} - m) \cdot O$ 需要每行的缩放因子 $\exp(m_{old} - m)$，这与张量 MMA 缩放指令要求的统一缩放因子冲突。
* 代理最大值 $\hat{m}$：为了解决此冲突，作者提出了代理最大值（Surrogate Maximum） $\hat{m}[i]$，定义为包含第 $i$ 行的 $X$ 行图块（tile）内的最大值：$\hat{m}[i] = \max_{k=\lfloor i/X \rfloor \cdot X}^{\min((\lfloor i/X \rfloor + 1) \cdot X, n)} m[k]$。其中 $X$ 是图块大小（如 16 或 64）。
* 动态适应性：与 FlashDecoding++ 的静态最大值不同，$\hat{m}[i]$ 根据局部行分布动态调整。

数值稳定性保证
引入代理最大值必须避免两种失败情况：
1. Any-overflow：由于 $\hat{m}[i] \ge m[i]$，保证了 $\exp(S[i, j] - \hat{m}[i]) \le \exp(S[i, j] - m[i]) \le 1$，因此严格防止了溢出（Infinity-over-Infinity）。
2. All-underflow：即某行的所有指数项都下溢为零。尽管较大的 $\hat{m}[i]$ 增加了下溢概率，但对于典型分布（如 Gaussian），所有项同时下溢的联合概率渐近地小。如果遇到极端分布，实现可选择性回退到向量化计算。

算法流程
Algorithm 1 展示了张量化在线 Softmax 算法：
* Step 1: 计算 $X$-row 图块最大值 $\hat{m} \leftarrow \text{tilemax}(S, X)$。
* Step 2-3: 使用重新利用的张量缩放指令重缩放注意力输出 $O \leftarrow \exp(\hat{m}_{old} - \hat{m}) \cdot O$（利用统一的 $\hat{m}$）。
* Step 4-5: 分配代理最大值并使用重新利用的张量 FMA 指令重缩放注意力 Logits $Z \leftarrow \log_2(e) \cdot S - (\log_2(e) \cdot m)$。
* Step 6: 使用向量单元计算指数 $\tilde{P} \leftarrow \exp_2(Z)$（因为 $\exp$ 难以张量化）。
* Step 7: 使用重新利用的张量行和归约指令更新行和 $l$。

张量-向量并行调度（Scheduling）

全张量化的局限与混合并行策略
初步结果表明，仅进行全张量化带来的收益有限，因为当前硬件（如 A100）的 MMA 吞吐量尚未达到碾压向量单元的临界点（Effective throughput 均为 16 elem/cycle）。因此，作者提出了架构感知的调度技术，利用指令级并行（ILP）和线程级并行（TLP）来并行化张量和向量单元的 Softmax 计算。

Ampere GPU 的 ILP 调度
在 Ampere GPU 上，FlashAttention-T 将 Softmax 计算拆分为张量化和向量化部分，并在每个 Warp 内进行交错调度。
* 张量-向量拆分策略：基于 16 行图块方案，首先使用 Warp All-Reduce 实现 16 行最大值代理 $\hat{m}$。然后选择以下策略之一拆分 $S$ 和 $O$ 的重缩放操作：
* 水平拆分（Horizontal split）：如 Fig 6(a) 所示，在 Warp 的矩阵图块内水平拆分，比例约为 1:1。
* 垂直拆分（Vertical split）：如 Fig 6(b) 所示，跨分配给同一 Warp 的多个矩阵图块进行拆分。

• 指令级并行（ILP）：在每个 Warp 内均匀交错执行重新利用的张量指令和向量指令，使得向量指令在张量指令的 Issue Bubble 中并行运行。此外，预生成并重用公共的 $B$ 片段以减少开销。这种调度将 Softmax 时间从 $t_{vec}$ 压缩到 $t'_{softmax}$（见 Fig 7(a)）。

Hopper GPU 的 TLP 调度
Hopper GPU 支持异步 WGMMA 指令，允许非阻塞 GEMM 执行和跨 Warpgroup 并行。
* 张量-向量拆分策略：由于 nvcc 编译器的限制（异步 WGMMA 要求累加器寄存器无依赖访问，否则会串行化），无法张量化 $S$ 或 $O$ 的重缩放（因为存在跨阶段寄存器依赖）。因此，在 Hopper 上仅张量化 $\tilde{P}$ 的行和归约（Row-sum summation）。这是一个叶子阶段操作，最大限度减少了依赖。
* 线程级并行（TLP）：如 Fig 7(b) 所示，将重新利用的 WGMMA 行和归约指令调度到 Warpgroup 下一次迭代的 $QK^T$ 和 $PV$ WGMMA 批次中。这使得张量化计算能与其他 Warpgroup 的向量化 $S$ 和 $O$ 重缩放操作以 TLP 方式并行运行，显著减少向量间隔时间。

A4 实验环境

• 硬件配置：

  • Server-class: NVIDIA A100 80GB SXM4 (Ampere, SM80); NVIDIA H100 80GB PCIe (Hopper, SM90A)。
  • Edge-class: NVIDIA Jetson AGX Orin 64GB (Ampere, SM87)。

• 软件配置与基线：

  • FlashAttention-2 (FA2)：Ampere 上的强基线。
  • FlashAttention-3 (FA3)：Hopper 上的强基线。
  • FlashInfer：专为 LLM 推理优化的库。
  • Triton：OpenAI 的深度学习编译器生成的 Kernel。
  • 实现采用 CUDA/C++，针对不同架构使用了特定的调度策略（Ampere 用 ILP，Hopper 用 TLP）。

• 数据集与参数：

  • 总 Token 数固定为 16384，总隐藏层维度 2048。
  • 变化参数：序列长度（512 ~ 8k），头维度（64, 128）。
  • 精度：FP16-FP32 和 FP8-FP32。

A4 实验结果

Attention 吞吐量性能
Fig 8 展示了不同配置下的吞吐量对比。FlashAttention-T 在绝大多数配置下均优于所有基线。

• 对比 FA2/FA3：实现了 1.02× 到 1.19× 的加速，平均加速比为 1.05-1.17×。
• 对比 FlashInfer/Triton：平均加速比达到 1.03-1.39×。
• 趋势分析：在非因果（Non-causal）配置下加速更明显；较小的头维度（如 64）加速比更大；由于 MMA 操作基于寄存器，对内存压力不敏感，因此在不同序列长度下增益保持稳定。

向量间隔比率评估
为了量化 Tensor 单元空闲时间的减少，计算了向量间隔比率（Vector Interval Ratio）。Fig 9 显示：
* Ampere GPU：FlashAttention-T 的 ILP 调度使得向量间隔比率比基线低 1.17–2.18 倍。
* Hopper GPU：TLP 调度效果更显著，将向量间隔比率降至 2.7%（相比基线大幅降低），这得益于更灵活的动态 Tensor-Vector 重叠。

消融实验
Fig 10 展示了各技术对性能的贡献（以 A100 为例）：
* FA2+Max16：仅集成 16 行代理最大值 $\hat{m}$，获得 1%-3% 的提升（归功于 Warp All-Reduce 优于 Shuffle）。
* AllTensor：全张量化变体反而导致轻微性能下降，证实了单纯张量化因 Swap 开销和吞吐量瓶颈并非最优。
* FA-T (ILP)：完整的 ILP 调度实现了最高 18.4% 的性能增益，证明了混合并行调度的有效性。

数值稳定性与精度
合成数据测试：Fig 11 显示，在不同离群值方差 $\sigma^2$ 下，FlashAttention-T 的 RMSE 均小于 $10^{-3}$，与 FA2 处于同一数量级，且未观测到数值失败（Numerical Failure）。在 $\sigma^2$ 较小时，误差主要来自 TF32 计算精度；在 $\sigma^2$ 较大时，误差主要由数值范围主导。
* 生成式基准测试：在 Llama3, Mistral, Qwen3 等模型上的测试表明（Table 2*），FlashAttention-T 的功能正确性指标（如 Pass@10, MMLU Score）与基线几乎完全一致，证实了其在真实应用中的可靠性。

A5 结论

本文提出了 FlashAttention-T，通过重新利用 GPU 张量 MMA 指令并配合张量化在线 Softmax 算法，成功将 Softmax 的关键计算转移至张量单元。结合架构感知的 ILP（Ampere）和 TLP（Hopper）调度策略，FlashAttention-T 有效解决了向量间隔瓶颈，在保证数值稳定性和精度的前提下，在多种 GPU 平台上显著提升了 Attention 的计算吞吐量。这项工作是向全张量化 Attention 执行迈出的重要一步。