注意力机制作为 Transformer 架构的核心组件，其计算效率直接决定了大规模语言模型的训练与推理成本。FlashAttention 系列算法通过分块计算与核融合技术，已经显著降低了注意力计算中的高带宽显存访问开销。然而，在追求更高计算密度的道路上，一个长期被忽视的瓶颈逐渐浮现：向量间隔瓶颈。FlashAttention-T 的核心贡献在于识别并解决了这一瓶颈，通过将关键的 Softmax 原语卸载至 Tensor Core，实现了真正意义上的完全张量化注意力计算。本文将从张量化计算与内存布局两个维度，深入解析这一创新的工程化细节与可落地参数。
向量间隔瓶颈的本质与识别
在传统 FlashAttention 的实现中，Softmax 计算依赖于 CUDA Core 上的标量与向量操作，而大规模的矩阵乘法则运行在 Tensor Core 上。这种混合执行模式导致了一个微妙但显著的性能损失：当 Tensor Core 正在执行矩阵乘累加（MMA）指令时，CUDA Core 可能处于空闲状态，反之亦然。这种交替执行模式造成了两类计算单元之间的调度间隙，论文作者将这一现象定义为向量间隔瓶颈。具体而言，Softmax 归一化过程中的指数运算、减法与求和操作均以向量形式在 CUDA Core 上执行，而它们与矩阵乘法之间存在固有的数据依赖关系，必须等待中间结果就绪后才能启动下一阶段的计算。
这一瓶颈在序列长度与批规模较大时尤为突出。根据论文中的微基准测试，在标准的因果注意力配置下，CUDA Core 与 Tensor Core 的利用率往往存在 15%-25% 的差距。传统的优化思路倾向于进一步压缩分块尺寸以提升显存访问效率，但这种做法会增加核启动开销并降低 Tensor Core 的计算吞吐。FlashAttention-T 的突破性洞察在于：与其让 Tensor Core 等待 CUDA Core 完成 Softmax，不如将 Softmax 的核心计算也迁移至 Tensor Core 上执行，从而消除两类计算单元之间的同步等待，实现计算流水线的紧密耦合。
Tensor Core 上的 Softmax 原语卸载
NVIDIA Ampere 与 Hopper 架构引入了第八代 Tensor Core 指令集，其中一项关键特性是对稀疏矩阵运算与向量操作的增强支持。FlashAttention-T 利用这些硬件特性，将传统的 Softmax 计算重新表达为 Tensor Core 可执行的矩阵运算。具体而言，指数运算被拆解为查表与乘法操作的组合，减法与求和则通过广播机制在矩阵维度上并行化。这种重新表达并非简单的算法移植，而是针对 Tensor Core 的指令流水线特性进行的深度优化。关键 Softmax 原语被封装为独立的 MMA 指令块，与前后相邻的 QK^T 矩阵乘法和 AV 矩阵乘法形成统一的指令流。
在实际实现中，FlashAttention-T 将注意力计算划分为三个主要阶段。第一阶段计算 Query 与 Key 的分块矩阵乘法并就地计算指数值，这一阶段完全在 Tensor Core 上执行。第二阶段执行跨块的归一化合并，论文中引入了在线 Softmax 算法的张量化变体，通过维护行级别的最大值与指数和来支持增量式的块合并。第三阶段将归一化后的权重应用于 Value，同样以矩阵乘法形式在 Tensor Core 上完成。这三个阶段之间的数据传递通过共享显存中的中间缓冲区完成，避免了全局显存的回写与重新加载。
内存布局优化的工程实践
完全张量化计算对内存布局提出了新的要求。传统 FlashAttention 的分块策略主要考虑的是如何最大化片上缓存的利用效率，而 FlashAttention-T 的内存布局还需要兼顾 Tensor Core 的数据供给模式与指令调度特性。论文中详细描述了针对 Ampere 与 Hopper 架构的差异化布局策略。在 Ampere 架构上，由于 L2 缓存与 Tensor Core 之间的带宽限制，FlashAttention-T 采用了更为激进的预取策略，将下一块的数据提前加载至片上共享显存。在 Hopper 架构上，则充分利用了其异步执行引擎，将数据加载与计算重叠，以掩盖访存延迟。
具体到参数配置，当在 Ampere A100 GPU 上部署 FlashAttention-T 时，建议将分块大小设置为 64×64 或 128×64，具体取值取决于序列长度与隐藏维度的组合。对于隐藏维度为 128 或 256 的标准配置，64×64 的分块能够提供最佳的 Tensor Core 利用率。对于采用 FP8 精度的高吞吐配置，由于 Tensor Core 的执行粒度发生变化，建议将行分块调整为 128。对于 Hopper H100 GPU，由于其异步执行能力的增强，可以适当增大分块尺寸以减少核启动开销，推荐使用 128×128 的分块配置。值得注意的是，这些参数需要与具体的模型架构配合调整，例如在多查询注意力（MQA）或分组查询注意力（GQA）场景下，由于 Key 与 Value 的分块维度发生变化，相应的布局参数也需要相应缩放。
监控指标与性能调优建议
在生产环境中部署 FlashAttention-T 时，需要关注几个关键指标以验证张量化优化的实际效果。首先是 Tensor Core 利用率与 CUDA Core 利用率的比值，在理想情况下这一比值应当接近 1.0，表示两类计算单元处于均衡负载状态。如果 Tensor Core 利用率显著高于 CUDA Core，则可能存在未能完全张量化的计算路径。如果 CUDA Core 利用率偏高，则可能是向量间隔问题未得到充分缓解。其次是核函数的执行时间与内存传输时间的比值，FlashAttention-T 的设计目标是让计算时间主导整体执行时间，内存传输应当尽可能被计算所掩盖。
在调试过程中，建议使用 NVIDIA Nsight Compute 工具采集核函数的指令级吞吐量数据。重点关注的指标包括 Tensor Core 活跃周期占比、共享显存带宽利用率以及寄存器压力。如果发现共享显存带宽接近饱和，可能需要调整分块策略或启用 Hopper 架构的异步复制特性。对于长序列场景，建议启用 FlashAttention-T 的序列并行变体，该变体通过在序列维度上进行分区来突破单卡显存限制，同时保持张量化计算的高效性。在序列并行配置下，需要特别注意跨分区的归一化同步开销，建议将分区大小设置为至少 4096 序列长度以分摊同步成本。
FlashAttention-T 的出现标志着注意力计算优化进入了一个新的阶段。通过将计算密集型的 Softmax 原语完全迁移至 Tensor Core，它不仅提升了硬件利用效率，更重要的是为未来的架构优化提供了一个可扩展的范式。随着 GPU 架构的持续演进，这种软硬件协同设计的方法论将在更多场景中发挥作用。对于工程实践者而言，理解并落地这些优化细节，将是构建高效 AI 系统的关键能力。
资料来源：论文《FlashAttention-T: Towards Fully Tensorized Attention by Exploiting Tensor Core Units》收录于 ACM SIGARCH/IEEE MICRO 2025，由 Tri Dao 团队发布。