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

Yupeng Sun; Yanzhao Li; Zhiqiang Zou; Bai Du; Zhiyuan Zhang; Hui Dong; Gaoyige Fan; Hui Wang

cs.LGApr 14, 2026

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

Authors: Yupeng Sun, Yanzhao Li, Zhiqiang Zou, Bai Du, Zhiyuan Zhang, Hui Dong, Gaoyige Fan, Hui Wang

Organizations: Huawei Technologies

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.

Figures & tables

Operation	FA Formula	VFA Formula
MUL	$S_{ij}=S_{ij}\times\mathrm{scale}$	$S_{ij}=S_{ij}\times\mathrm{scale}$
MAX	$\tilde{m}_{i}^{(j)}=\mathrm{rowmax}(S_{ij})$	-
MAX	$m_{i}^{(j)}=\max\!\bigl(m_{i}^{(j-1)},\ \tilde{m}_{i}^{(j)}\bigr)$	-
SUB	$\bar{S}_{ij}=S_{ij}-m_{i}^{(j)}$	$\bar{S}_{ij}=S_{ij}-m_{i}^{(j)}$
EXP	$\tilde{P}_{ij}=\exp\!\bigl(\bar{S}_{ij}\bigr)$	$\tilde{P}_{ij}=\exp\!\bigl(\bar{S}_{ij}\bigr)$
SUM	$\hat{l}_{i}^{(j)}=\mathrm{rowsum}(\tilde{P}_{ij})$	$\hat{l}_{i}^{(j)}=\mathrm{rowsum}(\tilde{P}_{ij})$

Table 1: Comparison of FA and VFA Computation Procedures

Panel of Figure 1 — Figure 1: Latency ratio normalized to Tensor on C16V32.

Tasks	Qwen3-30B		Qwen3-8B		Llama3-8B
	Baseline	VFA	Baseline	VFA	Baseline	VFA
MATH500	$0.424$	$0.42$	$0.25$	$0.262$	$0.134$	$0.146$
MMLU_abstract_algebra	$0.75$	$0.76$	$0.61$	$0.61$	$0.36$	$0.35$
MMLU_college_computer_science	$0.75$	$0.75$	$0.72$	$0.70$	$0.45$	$0.46$
MMLU_computer_security	$0.81$	$0.83$	$0.80$	$0.82$	$0.82$	$0.82$
HumanEval	$0.7683$	$0.7439$	$0.6402$	$0.628$	$0.372$	$0.372$

Table 2: Accuracy results of VFA on different models.

Methods	MMLU_abstract_algebra
Baseline (FA2)	$0.75\pm 0.0435$
VFA (full, $\mathrm{sabsmax}$ )	$0.76\pm 0.0429$
VFA (repr = $K_{\max}$ )	$0.23\pm 0.0423$
VFA (repr = $K_{\mathrm{mean}}$ )	$0.21\pm 0.0409$
VFA (repr = $\|K\|_{\max}$ , unsigned)	$0.22\pm 0.0416$
VFA (w/o $m$ -init)	$0.22\pm 0.0416$

Table 3: Ablation on MMLU_abstract_algebra : block representation

k^{\mathrm{repr}}_{j}

Panel of Figure 2 — Figure 2: Empirical evidence supporting the sink+local reordering strategy.

Panel of Figure 3 — Figure 3: Intra-block similarity for $Q$ and $K$ blocks using the SpargeAttention metric.

Panel of Figure 4 — Figure 4: Magnitude variation within $Q$ and $K$ blocks visualized by $\ell_{2}$ -norm statistics.

Panel of Figure 5 — Figure 5: Representative cases of block-maximum location along the key-block index $j$ , motivating the need for $m$ -initialization.

Method	MMLU_computer_security
Baseline (FA2)	$0.81\pm 0.0394$
Blasst ( $\lambda=1\mathrm{e}{-1}$ )	$0.8\pm 0.0402(32.3\%)$
Blasst ( $\lambda=3\mathrm{e}{-1}$ )	$0.84\pm 0.0368(44.0\%)$
Blasst ( $\lambda=5\mathrm{e}{-1}$ )	$0.83\pm 0.0378(49.5\%)$
Blasst ( $\lambda=9\mathrm{e}{-1}$ )	$0.83\pm 0.0378(56.1\%)$
Blasst_SWA ( $\lambda=1\mathrm{e}{-1}$ )	$0.8\pm 0.0402(37.4\%)$

Table 4: Accuracy results on MMLU_computer_security . Each cell is reported as

\mathrm{acc}\pm\mathrm{std}

followed by (parentheses) sparsity; for Blasst_FA4 we additionally report the rescale-skip rate as (sparsity / rescale-skip). For Blasst_Rowskip , sparsity is measured row-wise.

Efficient Attention

Efficient attention reduces the compute, memory traffic, and cache costs of Transformer attention while preserving useful long-range interactions.

See all 138 papers

FlashAttention for Scalable Vector Architectures
Aug 19, 2026Sonia Rani Gupta, Nikela Papadopoulou, Miquel PericàsEfficient Transformer Inference Efficient Attention
Chalmers University of Technology Gothenburg, Sweden · University of Glasgow Glasgow, UK
QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention
Apr 28, 2026Sehyeon Oh, Yongin Kwon, Jemin LeeSoftmax Attention GPU Acceleration
University of Science and Technology, Daejeon, Republic of Korea · Electronics and Telecommunications Research Institute, Daejeon, Republic of Korea · Pusan National University, Busan, Republic of Korea +1
Block Sparse Flash Attention
Dec 7, 2025Daniel Ohayon, Itay Lamprecht, Itay Hubara +3GPU Acceleration Self-Attention
Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel

Aug 19, 2026J/K move · Enter open · S save

FlashAttention for Scalable Vector Architectures

Sonia Rani Gupta, Nikela Papadopoulou, Miquel Pericàs

Chalmers University of Technology Gothenburg, Sweden · University of Glasgow Glasgow, UK

Efficient Transformer Inference Efficient Attention Transformer Attention LLM Inference Acceleration

Open paper AskPDF arXiv

Inference with transformer models on CPUs is increasingly important, especially for Small Language Models (SLMs), where vector architectures are emerging as a promising execution substrate. The attention module is a major bottleneck due to high memory bandwidth requirements; FlashAttention mitigates this by fusing operations to improve data locality and reduce intermediate memory traffic. In this paper, we present FlashAttention-V, a blocked FlashAttention for scalable vector architectures that adapts efficiently from short to very long vectors by exploiting parallelism across attention heads, inter-head packing to enable efficient utilization of vector lengths beyond the head dimension, and improving vector register utilization and memory access locality. We integrate FlashAttention-V into ggml within llama.cpp and evaluate it on TinyLlama, Llama 3.2, Qwen2.5, and Pythia-410M using gem5 and a Banana Pi BPI-F3. On the Banana Pi BPI-F3, we confirm that loop reordering and loop unrolling across attention heads are effective optimization principles, scaling performance gains with larger models and most pronounced with short contexts and during decoding. Simulation-based analysis shows that FlashAttention-V achieves 22x-42x speedup over scalar FlashAttention at 512-bit VL in prefill, with an additional 2x-2.5x gain scaling to 64 lanes and 4096-bit VL. During decode, FlashAttention-V achieves 8x-11x speedup using 512-bit vector lengths over scalar FlashAttention, with performance showing diminishing sensitivity to vector width and lane count due to single-token, memory-bound execution. We further identify structural bottlenecks in Q8_0 quantized linear layers that limit arithmetic amortization under long-vector execution, consistent across RVV and Arm SVE, indicating that current quantization formats pose a fundamental challenge to long-vector scalability.

Abstract

Efficient Attention

Explore similar work

FlashAttention for Scalable Vector Architectures

QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention

Block Sparse Flash Attention