在 Transformer 架构主导大模型的时代，注意力机制的计算效率是核心瓶颈。FlashAttention 系列通过 IO 感知和算子融合，革命性地降低了注意力计算的显存占用与延迟。然而，随着硬件架构的演进，尤其是 NVIDIA Hopper 架构引入了强大的张量核心（Tensor Core），传统的向量化计算模式已无法充分利用这些新型计算单元。在此背景下，中科大与计算所联合发表的 PPoPP '26 论文《FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor–Vector Parallelism》提出了一个关键演进方向：完全张量化的注意力（Fully Tensorized Attention）。
FlashAttention-T 并非另起炉灶，而是基于 FlashAttention-2 和 FlashAttention-3 构建的原型实现。其核心目标是将注意力计算中最后的 “堡垒”——Softmax 原语 —— 也进行张量化，从而在张量核心上完成整个注意力计算的闭环。传统 FlashAttention 虽然融合了 GEMM（矩阵乘）和 Softmax，但其 Softmax 计算通常由向量单元（CUDA Core）执行，而 GEMM 则由张量核心执行。这种异构计算模式可能导致计算单元间的空闲与等待，形成新的性能瓶颈。
张量化注意力的核心：重定义内存布局与计算模式
FlashAttention-T 的核心创新在于 “重用张量 MMA 指令执行 Softmax 原语”。这听起来像是一种 “黑客” 行为，但实则是对硬件指令集的深度挖掘。张量核心的矩阵乘累加（MMA）指令本是为密集矩阵乘法设计的。论文作者发现，通过精心设计的数据布局和计算流程，可以将 Softmax 计算中的指数、求和、归一化等操作，映射为一系列特殊的张量 MMA 操作。
这带来了两个根本性的变化：
- 内存布局的重新对齐：为了适配张量核心的 MMA 指令，输入数据（Q、K、V）需要在 SRAM（共享内存）或寄存器中进行特殊的排布。这种排布不再仅仅是为了向量化加载，而是为了满足张量核心对数据格式（如
m16n8k16）的严格约束。这意味着内存加载模式从 “面向向量” 转变为 “面向张量块”。 - 计算模式的统一：整个注意力计算流（QK^T、Softmax、与 V 相乘）理论上可以在张量核心上以统一的、流水线化的方式执行，减少了数据在向量核心与张量核心之间来回搬运的开销。论文中提出的 “张量化在线 Softmax 算法” 是关键，它能够在进行分块计算时，利用张量核心实时维护和更新 Softmax 所需的统计量（最大值和求和值）。
与 FlashAttention 系列的硬件适应性对比
要理解 FlashAttention-T 的价值，必须将其置于 FlashAttention 系列演进的脉络中。
- FlashAttention v1/v2：核心贡献是IO 感知和工作划分。通过将注意力计算分解为块（Tile），在高速 SRAM 中进行计算，避免将巨大的中间注意力矩阵写回 HBM（高带宽内存），从而大幅降低内存读写。其计算主力是向量核心，张量核心仅用于加速矩阵乘法部分。
- FlashAttention v3：为 Hopper 架构深度优化，引入了异步拷贝和更细粒度的双缓冲等技术，旨在更好地隐藏内存延迟，并提升张量核心在矩阵乘部分的利用率。然而，Softmax 仍然主要依赖向量核心。
- FlashAttention-T：将优化焦点从 “内存带宽” 进一步推向 “计算单元利用率”。其目标是让张量核心不仅处理 GEMM，也处理 Softmax，实现计算负载的完全张量化，从而最大化 Hopper 等架构中张量核心的吞吐量。这是一种从 “内存墙” 到 “计算墙” 的攻坚。
这种差异直接体现在硬件适应性上。FlashAttention 系列（尤其是 v2/v3）的设计相对通用，经过适配可以在 NVIDIA（CUDA）和 AMD（ROCm）平台上运行。而 FlashAttention-T 的 “张量化 Softmax” 高度依赖于 NVIDIA GPU 张量核心的特定 MMA 指令集和行为，目前是一个架构特化的解决方案。这既是其性能潜力的来源，也限制了其通用性。在 AMD 或其他 AI 加速器上实现类似优化，可能需要完全不同的指令重映射方案。
工程落地：参数、监控与未来演进
对于希望尝试或理解 FlashAttention-T 的工程师而言，以下几点是关键：
- 核心依赖：实现基于 CUTLASS（CUDA Templates for Linear Algebra Subroutines）库，这是 NVIDIA 的高性能线性代数模板库。深入理解 CUTLASS 的 Gemm 与 Epilogue（后处理）概念是定制化开发的前提。
- 关键参数：
- Tile Size：决定每次加载到 SRAM 中计算的块大小。FlashAttention-T 的 Tile 尺寸选择必须同时满足张量核心 MMA 指令的数据格式对齐要求，以及在线 Softmax 的数值稳定性需求。
- Wavefront 调度：论文提出的 “架构感知调度” 技术，核心是协调张量单元和向量单元的并行工作。需要精细控制计算流水线，确保当一个 Warp 的张量核心在执行 Softmax 的某个张量化步骤时，其他 Warp 的向量核心或张量核心能执行其他不冲突的任务，最大化 SM（流多处理器）内的利用率。
- 数值精度边界：张量化 Softmax 可能引入与标准向量计算不同的数值舍入行为。在混合精度训练（如 BF16/FP16）中，需要验证其对模型收敛性的影响，尤其是对于非常长的序列（如 128K+）。
- 监控要点：性能分析应重点关注 Tensor Core Utilization（张量核心利用率）和 SM Occupancy（SM 占用率）指标。理想情况下，FlashAttention-T 应实现接近 100% 的张量核心利用率，并保持高 SM 占用率，表明其成功地将计算压力从向量核心转移至张量核心。
- 未来演进方向：
- 支持更广泛的硬件：将张量化思想迁移至 AMD MI300X 的 Matrix Core 或 Google TPU 的 MXU。
- 动态稀疏性支持：当前张量化方案针对稠密矩阵。如何与激活稀疏化（如 ReLU-based Attention）或结构化稀疏（如 Block-Sparse）结合，是下一个挑战。
- 编译器集成：将张量化注意力作为一级原语集成到 AI 编译器（如 TorchInductor、MLIR）中，实现自动调度与代码生成。
结语
FlashAttention-T 代表了注意力计算优化从 “减少 IO” 到 “榨干算力” 的范式转变。它不再满足于让张量核心和向量核心 “各司其职”，而是试图让更强大的张量核心 “承包” 整个计算流程。虽然目前其实现高度特化于 NVIDIA 硬件，但其核心思想 ——通过指令重映射和统一计算模式来最大化专用计算单元的利用率—— 为后摩尔定律时代的 AI 计算优化指明了方向。对于追求极致推理 / 训练性能的团队，深入研究 FlashAttention-T 的设计哲学，比直接使用其代码可能具有更长远的价值。它提醒我们，在硬件快速演进的背景下，算法的优化必须与指令集架构（ISA）深度结合，方能持续释放性能红利。
资料来源
- Xu, Jianxing, et al. "FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor–Vector Parallelism." Artifact on Zenodo, 2025. https://zenodo.org/records/17673796
- Dao, Tri. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." GitHub Repository, Dao-AILab/flash-attention. https://github.com/Dao-AILab/flash-attention