arXiv:2604.12798cs.LGcs.AI2026-04

优化注意力计算中的向量瓶颈,提升FlashAttention在现代硬件上的速度。

VFA: Relieving Vector Operations in Flash Attention with Global Maximum Pre-computation

  • 通过预计算全局最大值减少每块的行最大值更新次数。
  • 在不损失精度前提下,实现比基线快近两倍的加速效果。
  • 适合追求极致推理性能的AI模型部署与硬件优化研究者。

FlashAttention-style在线softmax通过流式处理得分块并维护运行最大值和归一化因子,以线性内存实现精确注意力计算。然而,随着注意力核接近现代加速器的张量核/立方核峰值吞吐,非矩阵乘法组件——尤其是每块的行最大值与行求和归约及重缩放链——可能成为向量或SIMD瓶颈,主导延迟。本文重新审视FlashAttention,提出硬件友好的向量缓解闪存注意力(VFA),通过廉价的键块表示近似初始化运行最大值,重排键块遍历顺序以优先处理高影响的源块和局部块,并冻结剩余块的最大值以避免重复归约与重缩放。进一步将VFA与块稀疏跳过方法(如BLASST)结合,形成向量缓解稀疏注意力(VSA),降低块数和每块开销。值得注意的是,VFA和VSA完全避免了FA4.0中更新阶段的条件重缩放操作。在MMLU和MATH500等基准上的广泛评估,结合注意力统计分析验证了设计:(i) 源块和局部块重排可早期稳定运行最大值;(ii) 简单的Q和K块摘要因块内异质性失效;(iii) 当最大值出现在中间块时,需进行m-初始化。总体而言,VFA和VSA在无性能损失前提下有效缓解在线softmax归约瓶颈。相比C16V32基线,C8V32、C4V32和C4V16在现代硬件上实现近两倍加速,突破向量瓶颈;随着架构改进,C4V16有望实现六倍加速。

原文摘要 · Abstract (English)

FlashAttention-style online softmax enables exact attention computation with linear memory by streaming score tiles through on-chip memory and maintaining a running maximum and normalizer. However, as attention kernels approach peak tensor-core/cube-core throughput on modern accelerators, non-matmul components of online softmax -- especially per-tile rowmax and rowsum reductions and rescale chains -- can become vector or SIMD limited and dominate latency. This paper revisits FlashAttention and proposes Vector Relieved Flash Attention (VFA), a hardware-friendly method that reduces rowmax-driven updates of the running maximum while retaining the online-softmax structure. VFA initializes the running maximum via a cheap approximation from key-block representations, reorders key-block traversal to prioritize high-impact sink and local blocks, and freezes the maximum for remaining blocks to avoid repeated reductions and rescaling. We further integrate VFA with block-sparse skipping methods such as BLASST to form Vector Relieved Sparse Attention (VSA), which reduces both block count and per-block overhead. Notably, VFA and VSA completely avoid the conditional rescale operation in the update stage used in FA4.0. Extensive evaluations on benchmarks including MMLU and MATH500, together with attention statistics, verify our design: (i) sink and local reordering stabilizes the running maximum early; (ii) simple Q and K block summaries fail due to intra-block heterogeneity; (iii) m-initialization is required when maxima appear in middle blocks. Overall, VFA and VSA efficiently alleviate online-softmax reduction bottlenecks without performance loss. Compared to the C16V32 baseline, C8V32, C4V32 and C4V16 achieve nearly two times speedup on modern hardware while hitting the vector bottleneck. With upcoming architecture improvements, C4V16 will deliver six times speedup by enhancing exponent capacity.

注意力机制硬件优化加速器

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。