注意力机制作为 Transformer 架构的核心，其计算复杂度和内存开销始终是性能优化的关键挑战。自 FlashAttention 问世以来，通过算子融合与分块计算技术，业界逐步突破了传统注意力实现的显存瓶颈。然而，FlashAttention 系列的前代版本在硬件资源利用率上仍存在进一步挖掘的空间，尤其是张量单元与向量单元之间的并行潜力尚未被完全释放。FlashAttention-T 作为面向 PPoPP '26 的最新研究成果，在 FlashAttention-2 与 FlashAttention-3 的基础上更进一步，提出了完全张量化融合注意力的实现方案。其核心思路是充分利用 GPU 上的张量 - 向量并行性，将原本依赖向量单元执行的 softmax 原语重载至张量单元执行，从而最大化硬件利用率与整体吞吐量。
传统融合注意力的局限性与张量化契机
理解 FlashAttention-T 的创新价值，首先需要回顾前代技术在硬件利用层面的取舍。传统注意力计算涉及三个主要的矩阵乘法运算与一个 softmax 操作。在标准的融合注意力实现中，矩阵乘法通常借助 Tensor Core 的张量 MMA（Matrix Multiply-Accumulate）指令高效执行，这些指令能够在单个时钟周期内完成大规模矩阵运算的乘加操作。然而，softmax 作为一种归一化指数函数，其计算逻辑与矩阵乘法截然不同，通常在向量单元上按序执行。这种计算模式的分离导致了硬件资源利用的不均衡：张量单元在执行完 GEMM 后往往处于空闲状态，而向量单元则承担着额外的计算负担。
FlashAttention-T 的核心洞察在于：现代 GPU 架构（特别是 NVIDIA Hopper 及以后版本）提供了足够的灵活性，允许开发者以非标准方式利用张量指令集。具体而言，张量 MMA 指令虽然设计用于矩阵运算，但其底层的计算单元结构理论上可以适配 softmax 所需的指数与除法操作。问题的关键在于如何改造算法流程，使其既满足张量指令的约束条件，又能保持在线 softmax 所需的数值稳定性与并行特性。这一技术挑战的解决，直接导向了 FlashAttention-T 所提出的张量化在线 softmax 算法。
重载张量 MMA 指令：软硬件协同设计
FlashAttention-T 实现中最具技术深度的部分，在于其对 Tensor Core 指令的重载策略。标准的张量 MMA 指令接受两个输入矩阵，执行矩阵乘法，并将结果累加到输出矩阵的指定位置。为了将这一机制应用于 softmax 计算，研究团队需要将 softmax 的数学运算拆解为一系列可映射到 MMA 指令的操作序列。指数运算可以通过查找表与近似算法在张量单元上实现，而归一化所需的除法操作则可以转化为乘以预计算的倒数因子。
这种重载并非简单的指令替换，而是涉及算法层面的深度改造。张量化在线 softmax 算法需要在符合重载张量 MMA 指令约束的前提下，完成数据块的尺度变换、最大值传递与归一化因子计算。传统的在线 softmax 通过两轮遍历实现归一化因子的精确计算：第一轮获取全局最大值以确保数值稳定性，第二轮计算归一化结果。FlashAttention-T 的张量化版本则需要在单次遍历或更少的迭代内完成这些操作，以充分利用张量并行带来的效率提升。这要求研究团队对 GPU 的内存访问模式与指令流水线特性有极其精细的把控。
架构感知调度：释放张量 - 向量并行潜能
除了核心算法的张量化改造，FlashAttention-T 的另一项关键创新是架构感知的调度技术。现代 GPU 并非单一计算资源的堆砌，而是由多个异构处理单元组成的复杂系统。张量单元专精于高吞吐量的矩阵运算，向量单元则更适合标量与向量形式的控制流密集型操作。FlashAttention-T 的调度器能够智能地将在线 softmax 的计算任务在张量单元与向量单元之间进行划分与并行化调度。
具体而言，对于特定规模的注意力分块，调度器会评估张量单元执行 softmax 原语与向量单元执行残余计算之间的负载平衡。当分块较大时，张量单元可以承担更多的 softmax 计算，从而释放向量单元去处理其他并行的注意力头或批次任务。这种细粒度的调度策略需要充分考虑不同 GPU 架构（Hopper、Ampere 等）的资源配比差异，因此 FlashAttention-T 的调度器是架构感知的，能够根据目标硬件自动调整任务分配比例。
工程意义与未来展望
从工程实践的角度看，FlashAttention-T 的意义在于为长上下文场景下的注意力计算提供了更优的解决方案。随着上下文窗口扩展至数十万甚至百万 token 级别，注意力计算的性能瓶颈从单纯的计算密度转向了内存带宽与指令吞吐的综合利用。完全张量化融合注意力的实现，意味着开发者可以在相同的硬件配置下运行更大的模型批次，或者在相同延迟预算内处理更长的序列。
PPoPP '26 的 artifact 评估表明，FlashAttention-T 作为一个原型实现，已经展现出超越 FlashAttention-3 的性能潜力。其基于 FlashAttention-2 与 FlashAttention-3 的工程基础，也保证了与现有生态的兼容性。对于致力于 GPU 性能优化的工程师与研究者而言，FlashAttention-T 不仅提供了一套可复用的技术方案，更展示了软硬件协同设计在深度学习编译器优化中的巨大价值。未来，随着 GPU 架构的持续演进，张量 - 向量并行的编程模型有望成为高性能深度学习应用的标准范式。
资料来源：Zenodo Artifact for PPoPP '26 Paper "FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor-Vector Parallelism"。