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

Jianxing Xu

doi:10.1145/3774934.3786425

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

2026-01-01

0

PRE

AI

DOI：10.1145/3774934.3786425

原文链接

原文求助

摘要

En 中文

注意力机制是现代深度学习的核心，特别是在大型语言模型（LLMs）中，但其计算复杂度呈二次方增长。为在GPU上加速注意力计算，融合注意力技术（如FlashAttention）将矩阵乘法（GEMM）和Softmax计算合并为一个内核。然而，这些操作在计算上仍相互分离：GEMM利用高性能张量单元（Tensor Cores），而Softmax在较慢的向量单元（CUDA核心）上执行。这种不平衡导致严重的向量间隔周期，张量单元闲置等待向量单元完成，显著降低了张量单元的利用率。此外，硬件持续进步带来的更快速张量单元加剧了这一瓶颈。为解决向量间隔瓶颈，本文提出FlashAttention-T，一种推进全张量化融合注意力的实现。我们的关键思路是将关键Softmax原语卸载到闲置张量单元，最大化硬件利用率和吞吐量。具体而言，我们首先引入一系列操作数赋值方法，重新利用张量矩阵乘累加（MMA）指令执行Softmax原语（如逐元素缩放）。在此基础上，我们设计了一种具有数值稳定性保证的张量化在线Softmax算法，符合重新利用的张量MMA指令约束。为最大化性能，我们通过架构感知调度技术，在张量和向量单元间并行化在线Softmax计算，充分利用异构并行性。在包括NVIDIA Ampere（A100和AGX Orin）和Hopper（H100）GPU在内的多平台对各种注意力配置的广泛评估表明，FlashAttention-T在Ampere GPU上实现向量间隔比基准低1.17-2.18倍，在Hopper H100上将其降低至2.7%，相较于FlashAttention-2和FlashAttention-3平均加速达1.17倍，同时保持数值稳定性和精度。

Keyword:

Fused Attention

GPU Kernel

Tensor Core

TensorVector Parallelism

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

摘要

期刊

机构

引用论文