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

2026-04-14

AI Summary

This work addresses the SIMD bottleneck in FlashAttention caused by non-matrix operations—such as row-wise maxima computation and reductions—in the online softmax on modern accelerators. The authors propose Vector-friendly FlashAttention (VFA), which significantly reduces vector operation overhead by precomputing global maxima, reordering key blocks, and employing a max-freezing mechanism, thereby eliminating conditional rescaling. A novel, low-cost initialization of maxima based on key blocks and a high-impact block prioritization strategy are introduced to further enhance efficiency. When integrated with sparse attention, this approach yields Vector-friendly Sparse Attention (VSA), jointly reducing both the number of blocks and per-block computation. Experiments demonstrate that configurations such as C4V16 achieve nearly 2× speedup over the C16V32 baseline on MMLU and MATH500 benchmarks, with potential for up to 6× acceleration on future architectures.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Large Vision Models
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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.

Problem

• online softmax
• vector bottleneck
• attention computation
• rowmax reduction
• SIMD limitation

Innovation

• FlashAttention
• online softmax
• vector bottleneck
• global maximum pre-computation
• block-sparse attention