VFA：通过全局最大值预计算减轻 Flash Attention 中的向量运算
机器学习
2026-04-15 v1 人工智能
摘要
FlashAttention 风格的在线 Softmax 能够通过流式传输评分块并保持运行最大值和归一化器来实现线性内存的精确注意力计算。然而，随着注意力核接近现代加速器上的峰值张量核心/立方核心吞吐量，非矩阵乘法组件——尤其是每块行最大和行求和约简以及 rescale 链条——可能成为向量或 SIMD 受限并主导延迟。本文重新审视 FlashAttention，提出一种向量减轻 Flash Attention (VFA) 的硬件友好方法，通过保留在线 Softmax 结构来减少由行最大驱动的运行最大值更新。VFA 通过对键块表示进行廉价近似来初始化运行最大值，重新排序键块遍历以优先处理高影响 sink 和局部块，并冻结剩余块的最大值以避免重复约简和缩放。我们进一步将 VFA 与块稀疏跳过方法如 BLASST 集成，形成向量减轻稀疏注意力 (VSA)，从而减少块计数和每个块的开销。值得注意的是，VFA 和 VSA完全避免了 FA4.0 中更新阶段所使用的条件 rescale 操作。对包括 MMLU 和 MATH500 在内的广泛基准测试，以及注意力统计数据进行充分评估，验证了我们的设计：(i) sink 和局部重新排序早期稳定运行最大值；(ii) 简单的 Q 和 K 块摘要因块内异质性而失败；(iii) 需要在中间块中出现最大值时进行 m 初始化。总体而言，VFA 和 VSA在不损失性能的情况下有效缓解了在线 Softmax 约简瓶颈。相对于 C16V32 基线，C8V32、C4V32 和 C4V16 在面向向量瓶颈的现代硬件上实现了近乎两倍的加速；随着即将到来的架构改进，C4V16 将实现六倍加速，通过增强指数容量。
引用
@article{arxiv.2604.12798,
title = {VFA: Relieving Vector Operations in Flash Attention with Global Maximum Pre-computation},
author = {Yupeng Sun and Yanzhao Li and Zhiqiang Zou and Bai Du and Zhiyuan Zhang and Hui Dong and Gaoyige Fan and Hui Wang},
journal= {arXiv preprint arXiv:2604.12798},
year = {2026}
}