在 Transformer 架构的注意力机制中，计算效率与内存占用始终是一对核心矛盾。传统 FlashAttention 通过分块（Tiling）与在线 Softmax 两大技术创新，将注意力计算从高带宽内存（HBM）的瓶颈中解放出来，使长序列注意力计算成为可能。然而，随着 GPU 硬件演进到以 Tensor Core 为核心的计算范式，传统 FlashAttention 在向量级（Vector-Level）执行的 Softmax 原语开始暴露出新的性能瓶颈。FlashAttention-T 正是针对这一问题的系统性解决方案，它通过张量化（Tensorized）注意力机制，将 Softmax 原语重新映射到张量单元上执行，从而实现硬件资源的更充分利用。
传统 FlashAttention 的内存布局与计算范式
理解 FlashAttention-T 的创新，需要先回溯传统 FlashAttention 的设计哲学。标准注意力计算的痛点在于注意力矩阵的二次方复杂度 —— 对于序列长度 $S$，需要物化 $O (S^2)$ 的中间结果矩阵，这对 HBM 带宽的消耗是灾难性的。FlashAttention 的核心洞察是：将计算分解为适合片上 SRAM 的小块（Tiles），通过在线 Softmax 技术增量累积结果，避免一次性加载完整注意力矩阵。
在内存布局层面，传统 FlashAttention 采用了分块策略（Blocking）。具体而言，Query（Q）、Key（K）、Value（V）矩阵被划分为大小为 $B_c \times d$ 的块（Tile），其中 $B_c$ 是序列维度的分块大小，$d$ 是注意力头的维度。这些 Tile 被逐次从 HBM 加载到 SRAM，在片上完成 Q 与 K 的点积、Softmax 归一化、以及与 V 的加权聚合。这种布局的优势在于，它充分利用了 SRAM 的高带宽（相比 HBM 高出一个数量级）与低延迟特性，将数据复用最大化。
然而，在线 Softmax 的实现引入了新的复杂性。由于 Softmax 需要对整行元素求取最大值与指数和，传统的增量算法必须维护两个统计量 —— 当前行的最大值 $m$ 与归一化和 $l$。在每个 Tile 处理完毕后，需要根据新 Tile 的统计量对已累积的结果进行缩放（Rescale）。这个过程在代码层面表现为对 $O_i$（输出 Tile）的反复读写，以及对 $m_i$、$l_i$（统计量 Tile）的反复加载。根据实际测试数据，在 RTX 2070 上，FlashAttention v1 的 HBM 读取量高达 11.58 GB，写入量为 5.54 GB，这正是因为循环顺序不合理导致的反复 HBM 访问。
FlashAttention v2 对此进行了关键优化：反转循环顺序，让 Query 块在外层循环驻留，K 与 V 块流式加载。这一改动将 HBM 读取量降低了 92.98%（从 11.58 GB 降至 412.18 MB），几乎将 SRAM 作为累加器使用，只在最终写回 HBM。但即便如此，在线 Softmax 的统计量计算仍然面临着 MIO（Memory Input/Output）管线的瓶颈。
MIO 瓶颈：向量级执行的隐性限制
在 NVIDIA GPU 的微架构中，不同类型的运算被分发到不同的执行单元。矩阵乘法累加（MMA）指令由 Tensor Core 处理，拥有极高的吞吐量；而指数运算（Exp）、最大值比较（Max）等 Softmax 原语则由标量 / 向量单元（SFU/MIO）执行。在传统 FlashAttention 中，Softmax 的逐元素计算模式决定了它必须通过 MIO 管线，这导致了一个关键问题：MIO 管线的吞吐远低于 Tensor Core，成为流水线的瓶颈。
具体表现为，在 FlashAttention v2 的性能分析中，MIO Throttle Stalls 占据了 43.97% 的流水线停顿时间。这意味着，尽管 Tensor Core 处于高利用率状态，整个内核的执行效率却受限于 MIO 管线的排队延迟。此外，Shared Memory Bank Conflicts 也是常见问题 —— 由于 K 矩阵的列读取模式与 SRAM 的 bank 结构冲突，导致额外的内存事务开销。
这一问题的根源在于计算范式的不匹配：传统 FlashAttention 将 Softmax 视为一个向量级运算，但现代 GPU 的设计哲学是以矩阵 / 张量级运算为核心。向量级运算在 SM（Streaming Multiprocessor）上的执行单元数量有限，且容易受到流水线停顿的影响。要突破这一瓶颈，唯一的路径是将 Softmax 也张量化 —— 即重新组织计算，使其能够利用 Tensor Core 的 MMA 指令执行。
FlashAttention-T 的张量化设计原理
FlashAttention-T 的核心创新在于 ** 张量化在线 Softmax（Tensorized Online Softmax）算法，以及配套的张量 - 向量并行性（Tensor-Vector Parallelism）** 调度策略。这一设计将 Softmax 从向量级运算重新定义为张量级运算，从而将其卸载到 Tensor Core 上执行。
首先，在内存布局层面，FlashAttention-T 对 Softmax 的中间结果采用了 Tensor Core 友好的组织方式。传统 FlashAttention 的在线 Softmax 统计量（行最大值与行和）是一维向量，而 FlashAttention-T 将其扩展为与分块大小对齐的二维结构。这种布局使得 Softmax 的缩放运算能够通过矩阵乘法实现，而非逐元素操作。具体而言，已累积的输出 $O$ 与新 Tile 的贡献之间的加权合并，可以表示为矩阵乘法形式：$O_{\text {new}} = \alpha \cdot O_{\text {prev}} + \beta \cdot (P_{ij} \cdot V_j)$，其中 $\alpha$ 与 $\beta$ 是缩放因子，$P_{ij}$ 是注意力分数矩阵块。
其次，在计算图优化层面，FlashAttention-T 重新利用了 Tensor MMA 指令来执行 Softmax 的核心运算。传统的 MMA 指令用于执行 $C = A \times B + C$ 形式的矩阵乘法累加，而 FlashAttention-T 通过巧妙设计 $A$、$B$、$C$ 矩阵的布局，使得 Softmax 的缩放与合并运算能够嵌入到 MMA 的累加过程中。这种重利用不仅规避了 MIO 管线的瓶颈，还能够利用 Tensor Core 的高吞吐特性，将 Softmax 的计算密度提升到与矩阵乘法相当的水平。
最后，在硬件适配层面，FlashAttention-T 引入了架构感知的调度技术（Architecture-Aware Scheduling）。不同代际的 NVIDIA GPU 在 Tensor Core 的支持特性上存在差异，例如 SM 7.0（Volta）与 SM 8.0（Ampere）在 MMA 指令的精度支持与流水线深度上有所不同。FlashAttention-T 能够根据目标 GPU 架构动态调整分块策略与 MMA 指令的选择，从而在不同硬件上实现最优性能。这种自适应调度是其能够基于 FlashAttention-2 与 FlashAttention-3 的实现进行扩展的关键原因。
工程落地：监控参数与回滚策略
对于工程实践而言，理解 FlashAttention-T 的设计原理只是第一步，更重要的是掌握落地的关键参数与异常处理策略。以下是经过验证的监控指标与调优建议。
在性能监控层面，推荐使用 NVIDIA Nsight Compute（ncu）工具采集以下指标，以评估 FlashAttention-T 的实际效果。第一项关键指标是 MIO Throttle Stalls，该指标反映了 MIO 管线的饱和程度，传统 FlashAttention 的 v2 版本中该值通常在 40% 到 50% 之间，而 FlashAttention-T 应该将其降低到 10% 以下。第二项指标是 Tensor Core Utilization，可以通过 sm__throughput.avg.pct_of_peak_sustained 间接评估，该指标应接近或达到 Tensor Core 的理论吞吐上限。第三项指标是 HBM 带宽占用（dram__bytes.sum），用于验证张量化是否减少了片外内存访问。如果 MIO Throttle Stalls 未显著下降，则可能存在以下问题：Head Dim 未对齐 Tensor Core 的最优 Shape（通常为 16 的倍数）；或者调度策略未适配当前 GPU 架构。
在参数调优层面，分块大小（Block Size）的选择对性能影响显著。对于传统 FlashAttention，建议 $B_c$（K/V 维度分块）设置为 32 到 64 之间，以平衡 SRAM 占用与流水线并行度；对于 FlashAttention-T，由于需要适配 MMA 指令的 Shape 约束，建议 $B_c$ 与 Head Dim（$d$）的乘积为 8 的倍数（如 $d=64$ 时 $B_c=64$，$d=128$ 时 $B_c=64$ 或 $128$）。此外，Head Dim 的选择也需要考虑 Tensor Core 的最优计算形状 —— 在 NVIDIA Ampere 及以后架构上，$d=64$ 或 $d=128$ 通常能获得最佳性能。
在回滚策略层面，FlashAttention-T 的部署应遵循渐进式验证原则。首先，在测试环境（A10、A100 等）上运行基准测试，对比 FlashAttention-T 与 FlashAttention-3 在相同配置下的端到端延迟与吞吐。如果在特定 GPU 型号（如 T4、V100 等老旧架构）上出现性能退化或数值误差，应自动回退到 FlashAttention-3 或 FlashAttention-2。回退条件可以设定为：当 ncu 采集的 Tensor Core 相关指标低于预期阈值（如峰值吞吐的 50%），或者端到端延迟相比 FlashAttention-3 下降超过 10% 时触发。此外，需要注意 FP8 支持的兼容性 ——FlashAttention-3 引入了 FP8 量化支持，而 FlashAttention-T 在部分架构上可能仅支持 FP16/FP32，应根据实际精度需求选择合适版本。
总结与展望
FlashAttention-T 的张量化注意力机制，代表了注意力优化从「内存带宽受限」向「计算单元异构协同」范式的转变。传统 FlashAttention 通过分块与在线 Softmax 解决了 HBM 瓶颈，但将 Softmax 的计算留在了向量级执行单元上；FlashAttention-T 则通过重利用 MMA 指令，将 Softmax 也纳入张量级执行范畴，从而打通了从矩阵乘法到 Softmax 的完整张量计算流水线。这种演进路径对工程实践的启示在于：GPU 性能优化已经从单一内存或计算的优化，演进为对整个计算图与硬件微架构的协同设计。
在落地层面，建议采用渐进式部署策略，在 CI/CD 流程中集成 ncu 性能基准测试，设置合理的回滚阈值以保障稳定性。同时，关注 NVIDIA 未来架构（如 Blackwell）在 Tensor Core 与 MIO 单元上的演进，张量化注意力设计可能需要进一步适配新的硬件特性。
参考资料
- Zenodo Artifact: "FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor–Vector Parallelism" (PPoPP '26)
- A. Diro, "Reimplementing FlashAttention for performance and giggles" (2025)