VFA: Relieving Vector Operations in Flash Attention with Global Maximum Pre-computation Yupeng Sun1,* , Yanzhao Li1,* , Zhiqiang Zou1 , Bai Du1 , Zhiyuan Zhang1 , Hui Dong1 , Gaoyige Fan1 , Hui Wang1 1 Huawei Technologies
arXiv:2604.12798v1 [cs.LG] 14 Apr 2026
* sunyupeng7, [email protected]
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/normalizer [6, 8]. However, as attention kernels approach peak tensor-core/cube-core throughput on modern accelerators, the non-matmul components of online softmax—in particular, per-tile rowmax/rowsum reductions and the rescale chain—can become vector/SIMD limited and dominate end-to-end latency. This paper revisits FlashAttention and proposes Vector Relieved Flash Attention (VFA), a hardware-friendly modification that reduces the frequency of rowmax-driven updates of the running maximum while preserving the online-softmax accumulation structure. VFA initializes the running maximum using a cheap approximation based on key-block representations, reorders the key-block traversal to prioritize empirically high-impact regions (sink and local blocks), and freezes the running maximum for the remaining blocks to avoid repeated reductions and rescale operations. We further show that VFA composes naturally with online block-sparse skipping methods such as BLASST [26], yielding Vector Relieved Sparse Attention (VSA) that combines fewer processed blocks with cheaper per-block load overhead. It’s noteworthy that VFA/VSA entirely bypasses the rescale operation in the update stage, an operation which is conditionally executed in FA4.0. Extensive evaluations on representative benchmarks (e.g., MMLU, MATH500) and numerical analyses of attention statistics validate the design choices: (i) sink+local reordering is supported by early stabilization of the running maximum, (ii) naive block summaries of Q/K are insufficient due to intra-block heterogeneity, and (iii) m-initialization is necessary to cover cases where maxima occur in middle blocks. Overall, VFA/VSA provide a practical path to improving attention efficiency in regimes where online-softmax reductions become a primary kernel bottleneck, without degrading the model performance. Compared to the baseline C16V32, the C8V32, C4V32 and C4V16 versions can achieve about 2x speedup based on the latest architecture and they have all reached the vector bottleneck. As the evolution of new architecture, C4V16 will deliver a 6x speedup by increasing the exponent capacity.
1
Introduction
Transformer attention is a core primitive for modern foundation models across language, vision, and multimodal settings [9, 19, 25]. Yet, scaled dot-product attention is notoriously expensive for long sequences due to its quadratic complexity and the associated memory traffic for softmax and P V [25]. This has motivated substantial work on efficient attention mechanisms, encompassing two main directions: efficient kernel design and low-complexity algorithm design. Low complexity algorithm design primarily includes linear approximations and sparse computation patterns [23]. FlashAttention is a typical representative of efficient kernel design methods and has emerged as a widely adopted solution for exact attention: it avoids materializing the full attention matrix in HBM by tiling Q/K/V into on-chip memory and updating softmax statistics online, achieving linear memory in sequence length while retaining numerical stability [8]. The FlashAttention line of work further demonstrates that efficiency is not purely an algorithmic issue but also a kernel-mapping problem. FlashAttention-2 improves performance via better parallelism and work partitioning, increasing occupancy and reducing shared-memory communication [6]. FlashAttention-3 exploits hardware features such as asynchronous data movement and low-precision pathways (e.g., FP8) to further improve pipeline efficiency on modern GPUs [21]. FlashAttention-4 further reduces softmax-side overhead by 1/19
approximating exp2 with low-latency linear instructions and conditionally eliding the output rescale on Õ when it is numerically safe, thereby shrinking the non-matmul critical path [7]. These advances underline a key trend: once tiled matmuls are highly optimized, attention performance increasingly hinges on the remaining non-matmul components of online softmax, including per-tile reductions (rowmax/rowsum) and the rescale chain used to maintain numerical stability. Motivation: when online-softmax becomes vector limited. To achieve higher throughput and lower latency, models are evolving toward lower precision such as MXFP4 [24], NVFP4 [1] and HIF4 [16]. However, these lower precisions are only applied in GEMM and vector operations are still in high precision, which bring higher load imbalance for fused kernels. In many practical deployments, we observe that the cost of online-softmax statistic updates can become vector/SIMD limited and disproportionately impacts end-to-end latency, especially when matmul pipelines already run near peak utilization. This motivates a complementary research question to prior FlashAttention work: can we reduce the frequency of per-tile reductions and rescale operations, while keeping the same online-softmax accumulation structure and kernel compatibility? Our approach. We propose VFA, a hardware-friendly variant of FlashAttention-style online softmax. VFA modifies the standard per-tile update schedule in three steps. First, it performs a fast m-initialization that approximates the per-row running maximum using cheap dot products between Qi and precomputed key-block representations. Second, it reorders the traversal of key blocks to prioritize an initial “sink” block and the local block aligned with the query position, motivated by empirical evidence that attention mass and the eventual running-max stabilization concentrate in these regions. Third, it freezes the running maximum for the remaining blocks, skipping per-block rowmax and the rescale chain, while still accumulating exp(Sij − mi ) and P V using the same state variables and stabilization form. Importantly, VFA targets statistic-computation overhead rather than skipping attention blocks outright. Composition with FlashAttention-4. VFA naturally skip all the rescale operations in attention computation, avoiding online conditional branch in FA4 [27]. Composition with sparse attention. VFA is orthogonal to sparse attention methods that reduce the number of processed blocks. In particular, BLASST introduces an online-statistics-driven rule that skips blocks whose local maximum is sufficiently below the running maximum, requiring minimal changes to FlashAttention-style kernels [26]. We show that VFA composes naturally with BLASST, yielding VSA that combines two multiplicative gains: fewer blocks processed (sparsification) and cheaper per-block load overhead (vector-relieved online softmax). Empirical analyses supporting the design. Beyond end-to-end accuracy evaluations, we provide numerical analyses of attention statistics to justify each component. We show that the running maximum often stabilizes early at sink/local blocks, explaining why reordering is effective. We also demonstrate that naive block representations of Q/K are ineffective due to moderate intra-block directional similarity and large norm variability, and that m-initialization is necessary because the true block-maximum can occur in middle blocks that would otherwise be missed under max-freezing. Contributions.
In summary, this work makes the following contributions:
• We identify and study a kernel-level bottleneck in FlashAttention-style online softmax: per-tile reductions and rescale operations can become vector/SIMD limited in regimes where matmul is already efficient. • We propose VFA, a hardware-friendly online-softmax variant that uses m-initialization, sink+local reordering, and max-freezing to reduce the frequency of rowmax-driven updates while preserving the standard accumulation structure. • We present VSA by composing VFA with BLASST-style block skipping [26], illustrating that vectorrelieved load overhead and sparse skipping are complementary, and suggesting that VFA can serve as a reusable building block in future attention kernel evolution. • We provide numerical analyses of attention statistics (running-max stabilization, block similarity, norm variability, and block-max peak locations) that explain when and why each design choice is effective. 2/19
• We evaluate the performance of VFA based on the compute capability of the latest architecture.
2
Related Works
2.1
Vanilla Attention
Transformer has become the de facto backbone for modern foundation models across language, vision, and multimodal applications, making attention efficiency a first-order concern in both √ training and inference [9,19]. Transformer adopts scaled dot-product attention, which computes S = QK ⊤ / d, applies a row-wise softmax P = softmax(S), and produces the output O = P V . Although conceptually simple, vanilla attention is both compute- and memory-intensive for long sequences: forming S ∈ RNq ×Nk incurs O(Nq Nk d) FLOPs, while explicitly materializing S and P requires O(Nq Nk ) memory and substantial HBM traffic due to multiple reads/writes across the softmax and P V stages [25]. In practice, the softmax stabilization (e.g., subtracting the row maximum) and the normalization/reduction steps (rowmax, rowsum, and rescaling) further introduce bandwidth-heavy operations and synchronization, making the end-to-end kernel increasingly memory-bound as sequence length grows. These characteristics motivate IO-aware attention formulations that avoid storing the full attention matrix and instead stream blocks through on-chip memory while updating softmax statistics online.
2.2
Efficient Attention Mechanisms
Flash Attention FlashAttention (FA) establishes an IO-aware formulation of exact attention by explicitly accounting for the GPU memory hierarchy [8]. Instead of materializing the full attention matrix in HBM, FA tiles the computation so that blocks of Q, K, and V are staged in on-chip SRAM, and the softmax statistics are updated online using a running maximum and normalizer [17]. This design reduces HBM reads/writes and achieves linear memory in sequence length while preserving exactness. While FA significantly improves end-to-end throughput for long sequences, it still underutilizes the compute peak of modern accelerators. FlashAttention-2 (FA2) attributes this gap largely to kernel-level bottlenecks—suboptimal work partitioning across thread blocks/warps, low occupancy for single-head cases, and excessive communication via shared memory [6]. FA2 refines both algorithm and kernel mapping: (i) it reduces non-matmul work in the online-softmax pipeline, (ii) parallelizes attention for a single head across multiple thread blocks to raise occupancy, and (iii) improves intra-block warp specialization to cut shared-memory traffic. These changes bring utilization closer to GEMM-like efficiency on contemporary GPUs. FlashAttention-3 further demonstrates that attention kernels remain sensitive to evolving hardware features. Targeting Hopper GPUs, FA3 leverages asynchronous data movement (e.g., TMA) and warp specialization to overlap memory transfers with Tensor Core computation, and interleaves block-wise matmul with softmax updates to increase pipeline efficiency [21]. In addition, FA3 introduces low-precision pathways (notably FP8 with block-wise quantization) to accelerate compute while maintaining numerical accuracy, showing that correctness in low precision hinges on careful kernel scheduling and stabilization. Overall, these works optimize exact attention by co-designing IO-aware formulations with kernel-level scheduling, parallelization, and numerical handling in the online-softmax update path. FlashDecoding++ [10] is designed as a decode-time LLM inference engine, where the main optimizations include asynchronized softmax with a unified maximum, flat-GEMM optimization, and heuristic dataflow selection to better utilize GPU resources under shape-dependent workloads. By contrast, VFA is a lightweight modification to the online-softmax statistic update in FlashAttention-style exact attention kernels. As a result, FlashDecoding++ mainly improves decode-system efficiency through scheduling and synchronization design, whereas VFA improves kernel-level efficiency by reducing vector-limited statistic computation within the online-softmax loop. KV-cache-efficient attention head variants (MQA/GQA/MLA). Beyond improving the asymptotic complexity of attention, a complementary line of work targets the key–value (KV) cache cost in autoregressive decoding, where storing and reading per-head K/V often becomes a dominant memory and bandwidth bottleneck. Multi-Query Attention (MQA) reduces KV-cache size by sharing a single set of keys/values across all query heads, enabling faster decoding with minimal architectural change [22]. Grouped-Query
3/19
Attention (GQA) generalizes this idea by partitioning heads into groups that share K/V within each group, offering a smoother accuracy–efficiency trade-off between standard multi-head attention and MQA [2]. More recently, Multi-head Latent Attention (MLA) compresses the KV cache into a latent representation, aiming to significantly shrink KV memory footprint while maintaining strong model quality in long-context settings [13]. Recent work on efficient attention for long-context LLMs is commonly organized into two main families: linear attention and sparse attention [23]. Both aim to reduce the quadratic time/memory complexity of full softmax attention, but they differ fundamentally in approximation strategy and kernel implications. Linear attention. Linear attention methods replace the softmax kernel with a form that enables associative re-ordering of computations (or kernel feature maps), reducing complexity to linear in sequence length [5, 12]. Representative approaches include kernelized attention (e.g., random feature maps and positive mappings) and recurrent/fast-weight formulations that permit streaming updates. While linear attention provides attractive asymptotic complexity, it typically introduces modeling approximations and may require careful numerical treatment and kernel engineering to achieve competitive accuracy and throughput in practice. Sparse attention. Sparse attention, in contrast, retains the exact softmax form on selected entries, but restricts the set of key/value tokens (or blocks) that each query attends to. This yields efficiency by reducing the number of score computations and the amount of V -matrix traffic, and often aligns naturally with block/tile-based FlashAttention kernels. 2.2.1
Sparse Attention
Following the survey literature, sparse attention methods can be categorized by how the sparsity pattern is determined [18]: (1) Static / pattern-based sparsity. These methods employ pre-defined sparsity patterns that are independent of the input content, such as sliding-window (local) attention with a small number of global tokens. Longformer combines local windowed attention with task-motivated global attention, providing a drop-in sparse replacement for full attention [4]. BigBird further augments local attention with random and global connections, providing theoretical guarantees (e.g., expressivity) while enabling longer contexts [28]. More recently, inference-oriented pattern designs further refine static sparsity to better match the attention mass distribution and kernel tiling behavior. For example, MInference proposes an “A-shape” sparse pattern that allocates dense compute to early “sink” regions and a local band while sparsifying the long-range tail, enabling efficient million-token prompt inference without requiring content-dependent routing [11]. Such static patterns are kernel-friendly and predictable, but may under-utilize compute on easy cases and lack adaptivity on hard cases. (2) Dynamic / content-based sparsity. Dynamic sparse attention selects attended tokens/blocks based on the input content, typically via routing, clustering, or learned selection modules. Routing Transformer introduces content-based routing via online clustering to produce sparse attention patterns that adapt to the sequence [20]. Beyond clustering-style routing, recent methods increasingly adopt lightweight scoring/indexing to retrieve a sparse set of keys for each query. SpargeAttention, for instance, uses sparse retrieval to identify a compact set of relevant tokens/blocks and then performs exact attention over the selected set, improving the accuracy–efficiency trade-off for long contexts [29]. Similarly, DeepSeek Sparse Attention (DSA) introduces a “lightning indexer” that computes query-to-token index scores and retrieves top-k entries for each query token, followed by attention over the selected key-value set [14]. Finally, MoBA (Mixture of Block Attention) explores a mixture-style block selection mechanism to activate different sparse block-attention experts, offering another route to content-adaptive sparsity at block granularity [15]. Dynamic sparsity can improve accuracy–efficiency trade-offs by focusing computation on relevant regions, but often complicates implementation and may incur additional overhead for selection/routing. (3) Online-statistics-driven sparsity within FlashAttention. A particularly kernel-compatible subclass of dynamic sparsity uses online softmax statistics already computed inside FlashAttention-style kernels (e.g., running max/normalizer) to make skipping decisions without auxiliary proxy scoring passes [3]. This direction is attractive for long-context inference because it preserves FlashAttention’s numerical stability
4/19
Algorithm 1 Vector Relieved Flash Attention c c r , initialization block count Tc1 , Value blocks {Vj }Tj=1 , Key blocks {Kj }Tj=1 Require: Query blocks {Qi }Ti=1 Tr Ensure: Output blocks {Oi }i=1 1: /* Precompute key-block representations */ 2: for j = 1 to Tc1 do 3: kjrepr ← sabsmax(Kj ) 4: end for 5: for i = 1 to Tr do (0) (0) (0) 6: Initialize mi = −∞, Oi = 0, li = 0 7: Initialize mi,init = −∞ ▷ vector, one value per row of Qi 8: for j = 1 to Tc1 do ▷ vector over rows of Qi ← Qi (kjrepr )⊤ 9: scoreapprox ij ) ▷ elementwise max 10: mi,init ← max(mi,init , scoreapprox ij 11: end for (0) 12: mi ← mi,init 13: for j ∈ ⟨1, i, 2, 3, . . . , Tc ⟩ with j ̸= i in the tail do 14: Compute Sij = Qi Kj⊤ ▷ Attention scores 15: if j = 1 or j = i then (j) 16: m̃i = rowmax(Sij ) ▷ Local maximum (j) (] j −1) (j) 17: mi = max mi , m̃i ▷ j] − 1 denotes the previous iteration index in the j-loop (j) 18: P̃ij = exp Sij − mi ▷ Attention weights
19: 20: 21: 22: 23: 24:
(j)
li
(] j−1)
= e mi
(j) Oi = e
(j)
−mi
(] j −1)
li
(] j−1) (j) mi −mi
+ rowsum(P̃ij )
(] j −1) Oi + P̃ij Vj
else (j) (] j −1) mi = mi (j) P̃ij = exp Sij − mi (j)
li
(] j −1)
= li
▷ Freeze m when j ∈ / {1, i}; j] − 1 is the previous j-loop iteration ▷ Attention weights
+ rowsum(P̃ij )
(j) (] j −1) Oi = Oi + P̃ij Vj
25: 26: end if 27: end for (T ) (T ) 28: Oi = Oi c /li c 29: end for Tr 30: return {Oi }i=1
▷ No rescale factor ▷ No rescale factor
▷ Final normalization
machinery while directly reducing the number of processed tiles. However, due to the online running mechanism, it usually can’t achieve the optimal sparsity without performance degradation.BLASST proposes a drop-in sparse attention mechanism that prunes attention blocks dynamically using only information already available in FlashAttention-style online softmax [26]. These methods primarily reduce how many tiles are processed, whereas our work focuses on reducing the per-tile statistic-update overhead, making the two directions complementary.
3
Methods
3.1
Algorithm Design: Vector Relieved Flash Attention
Algorithm 1 presents VFA, a hardware-friendly variant of FlashAttention-style online softmax. VFA is motivated by a practical observation: in modern attention kernels, the non-matmul components of online softmax—most notably the per-block rowmax reduction and the subsequent rescale chain—can become vector/SIMD limited and dominate latency when tiled matmuls already approach high efficiency. The goal of VFA is therefore to reduce the frequency of rowmax-driven updates of the running maximum mi while keeping the same online-softmax state variables and kernel interface as FlashAttention.
5/19
Online-softmax state. For each query block Qi , we maintain the same running statistics as FlashAttention: (j) (j) (j) a per-row running maximum mi , a running normalizer li , and an output accumulator Oi . The final (T ) (T ) output is obtained by the standard normalization Oi = Oi c /li c . This alignment ensures VFA remains compatible with FlashAttention-style kernels and numerical stabilization. Step 1: fast initialization of the running maximum. A key difference from standard FlashAttention (0) is that VFA does not start from mi = −∞. Instead, we initialize mi using a lightweight approximation that avoids forming the full score block Sij . Concretely, for each key block Kj in a candidate set j ∈ {1, . . . , Tc } , we construct a block representation kjrepr = sabsmax(Kj ), where sabsmax selects the element with maximum absolute value per dimension while preserving its sign. We then compute an approximate score vector scoreapprox = Qi (kjrepr )⊤ , which yields one approximate score per row of Qi . Taking an elementwise maximum ij (0)
over candidates produces mi,init , and we set mi ← mi,init . Intuitively, this step provides a high-quality initial scale for exponentiation, making subsequent max updates less frequent. Since kjrepr = sabsmax(Kj ) c depends only on the key block, we precompute {kjrepr }Tj=1 once per attention invocation and reuse them across all query blocks Qi . For decoding with a growing KV-cache, this precomputation can be performed incrementally for newly appended key blocks. This eliminates redundant vector reductions from the inner loop and reduces the overhead of the approximation stage. The parameter Tc1 controls the block size used for partitioning K in the initialization stage, and in our implementation we set Tc1 = Tc , so that the initialization uses the same key-block granularity as the main attention loop. Step 2: reordered processing with selective max updates. VFA next processes key/value blocks in a reordered schedule j ∈ ⟨1, i, 2, 3, . . . , Tc ⟩ (excluding repeated i in the tail). This ordering prioritizes two blocks that empirically concentrate large attention scores (sink and local regions), so that if a true maximum is likely to be encountered, it is encountered early. Crucially, VFA only computes the exact per-row maximum (j) (j) m̃i = rowmax(Sij ) and updates mi when j ∈ {1, i}. For these special blocks, VFA follows the standard (j)
numerically-stable online-softmax update: mi (] j −1)
factor exp(mi
(] j −1)
= max(mi
(j)
, m̃i ) and applies the corresponding rescale
(j)
− mi ) to maintain consistency of li and Oi .
Step 3: freezing the running maximum for the remaining blocks. For all other blocks j ∈ / {1, i}, (j)
(] j −1)
VFA freezes the running maximum, mi = mi , and skips the costly rowmax computation as well as (j) (j) the rescale chain. Given a fixed mi , VFA still computes the attention weights P̃ij = exp(Sij − mi ) (j)
(] j −1)
(j)
(] j −1)
and accumulates li = li + rowsum(P̃ij ) and Oi = Oi + P̃ij Vj . This design eliminates per-block reductions and rescale-related vector operations on the majority of blocks, thereby relieving vector pressure in the kernel. Discussion and implications. Compared with FlashAttention, VFA trades frequent max updates for a fast initialization plus selective exact updates on a small set of blocks. Compared with threshold-based block skipping, VFA targets a complementary bottleneck: it reduces the overhead of computing the statistics (especially rowmax) rather than relying on these statistics to decide skipping. As a result, VFA can be combined with dynamic sparsification methods to further reduce the number of processed blocks, while keeping the online-softmax accumulation structure unchanged.
6/19
Table 1. Comparison of FA and VFA Computation Procedures Operation FA Formula VFA Formula MUL Sij = Sij × scale Sij = Sij × scale (j) MAX m̃i = rowmax(Sij ) (j) (j−1) (j) MAX mi = max mi , m̃i (j) (j) S̄ij = Sij − mi S̄ij = Sij − mi SUB EXP P̃ij = exp S̄ij P̃ij = exp S̄ij ˆl(j) = rowsum(P̃ij ) ˆl(j) = rowsum(P̃ij ) SUM i i (j) (j−1) (j−1) (j) (j) (j) (j−1) (j) −mi l + ˆl MAD l = e mi l =l + ˆl i
i
MUL update
4
(j−1)
(j−1)
Õi
(j)
i (j−1)
−mi Oi = e mi (j) (j−1) Oi = Õi + P̃ij Vj
i
i
i
(j)
Oi
(j−1)
= Oi
+ P̃ij Vj
Performance Evaluation Tensor
Exponential
Vector (FA)
Vector (VFA)
Latency ratio (%)
100 80 60 40 20 0 C16V32
C8V32
C4V32
C4V16
C4V16 (2×Exp)
Figure 1. Latency ratio normalized to Tensor on C16V32. Implementation benefits of VFA. To quantify the practical benefits of VFA over standard FA, we compare their operator-level computation procedures at the granularity of a single (Qi , Kj , Vj ) block interaction. VFA introduces only a lightweight preprocessing stage. The extraction of kjrepr = sabsmax(Kj ) can be naturally fused into earlier computation, and the subsequent construction of mi,init is implemented using pure tensor operations, contributing less than 1% of the total tensor computation of standard FA. Therefore, we focus on the differences inside the main attention accumulation loop. Table 1 summarizes the per-block procedure differences between FA and VFA. In FA, each iteration updates (j) the running maximum via a local row-wise maximum m̃i = rowmax(Sij ) and a running max reduction, (j−1) (j) −mi and then applies a rescaling factor emi to both the normalization term li and the partial output Oi . These steps introduce additional vector-linear work (MAX reductions and rescale-related MUL/MAD operations), as well as extra data dependencies and off-core traffic for rescaling the accumulated states. In (0) contrast, VFA first initializes the per-row maximum with an inexpensive approximation, mi ← mi,init , and then freezes m for the majority of blocks (i.e., j ∈ / {1, i} under our schedule). Consequently, VFA avoids repeated rowmax and skips the rescale multiplications on li and Oi in the update rules for those blocks. Importantly, the tensor-dominant primitives (QK and PV) and the exponential pipeline (SUB/EXP/SUM) remain unchanged; the speedup primarily stems from reducing vector-linear operations and their associated synchronization and bandwidth overheads. Figure 1 reports the normalized latency breakdown on the latest architecture. On this architecture, the peak Tensor throughput for FP16, FP8, and FP4 is in the ratio of 1 : 2 : 6, while the exponential pipeline delivers 1/250 of FP16 Tensor throughput, and the vector pipeline provides 8 times the exponential throughput. Compared with FA, VFA reduces the latency contribution of the vector portion from approximately 77% to 7/19
46% in the C16V32 setting, and it maintains a similar (∼ 46%) vector ratio under C8V32. Under C4V32, the vector component drops from roughly 85% to 54%, indicating that the benefit becomes more pronounced as the configuration increases the relative weight of vector-side load overhead in FA. When moving to C4V16, VFA further decreases the vector share to about 27%, after which the execution becomes increasingly bounded by the exponential/SFU pipeline whose contribution remains comparatively stable. This explains why the observed gains for C8V16 and C4V16 can become similar: once the kernel enters an exp-SFU–bound regime, further reductions in linear vector work yield diminishing end-to-end returns. This trend also helps clarify the relation to FA4-style optimization. FA4 reduces the softmax bottleneck mainly by approximating exp with low-latency linear instructions, effectively converting part of the exponential cost into vector-executable work and thereby allowing the exp and vector workloads to be pooled onto the same linear pipeline. Under such a mechanism, reducing vector latency remains beneficial even in regimes that would otherwise appear exp-bound, because the approximated exponential path itself now competes for the same linear execution resources. From this perspective, VFA is complementary to FA4: FA4 alleviates the exponential bottleneck by remapping it to linear instructions, while VFA directly reduces the linear-side overhead caused by repeated rowmax-driven updates and rescale operations. Consequently, when combined with FA4-style exp approximation, the vector-load reduction introduced by VFA can still translate into additional performance gains, including in configurations where the baseline implementation would already be limited by the exp/SFU path. Looking forward, if the architecture approximately doubles the throughput of exp16, the exp/SFU bottleneck would be substantially relaxed. In that scenario, the vector-side savings provided by VFA are expected to translate more directly into overall kernel speedups, and the achievable gain in configurations such as C4V16 would increase accordingly, with the limiting factors shifting back toward tensor compute and memory traffic.
5
Experiments
5.1
Research Questions
The experiments are designed to answer four questions aligned with the algorithmic components of VFA/VSA and the empirical analyses in Section 5.5. • RQ1 (Accuracy preservation). Does VFA preserve downstream task accuracy relative to FA2 when it freezes the running maximum on most blocks and reduces rowmax-driven updates? • RQ2 (Component necessity). Which design elements are necessary to maintain accuracy under max-freezing? In particular, how do (i) m-initialization, (ii) sink+local reordering, and (iii) the choice of key-block representation kjrepr affect accuracy? • RQ3 (Approximation fidelity). Is the approximation stage required to be row-wise (i.e., scoreapprox = ij repr ⊤ repr Qi (kj ) ), or can it be replaced by a cheaper block-wise approximation based on qi without degrading accuracy? • RQ4 (Empirical justification). Do the observed attention statistics—early stabilization of the running maximum at sink/local blocks, intra-block heterogeneity of Q/K, and the existence of middleblock maxima—explain the success/failure modes of the ablated variants and justify the design choices of VFA?
5.2
Experimental Setup
Models. We evaluate three instruction-tuned decoder-only LLMs: Qwen3-30B-A3B, Qwen3-8B, and Llama3-8B. For all models we use the default tokenizer/chat template from the corresponding checkpoints and disable any additional sampling randomness for accuracy evaluation. Baselines. We compare against FA2 as the primary baseline [6, 8]. For each model and task, the “Baseline” column in Table 2 corresponds to FA2 under the same decoding and prompt protocol as VFA.
8/19
Datasets and Tasks. We cover representative reasoning, knowledge, and coding workloads: (i) MATH500 for mathematical reasoning, (ii) MMLU subsets (abstract algebra, college computer science, comp uter security) for English multiple-choice knowledge evaluation, (iii) HumanEval for code generation, and (iv) CMMLU subsets (college mathematics, philosophy) for Chinese multiple-choice evaluation. All multiple-choice tasks are evaluated with exact-match accuracy under a fixed few-shot setting, and HumanEval is evaluated using standard pass@1-style greedy decoding. Implementation Details. VFA is implemented inside the FlashAttention-2-style kernel as a conditional update rule on the running max and normalizer. We ensure functional equivalence of non-attention components (tokenization, prompting, stopping criteria, and post-processing) across baseline and VFA.
5.3
Accuracy Studies
We evaluate whether VFA preserves task-level accuracy relative to FlashAttention-2. Table 2 reports mean accuracy (or pass@1 for HumanEval) across three model families. Overall, VFA matches the baseline on most tasks and shows small gains in several cases. These results suggest that VFA’s policy—reducing rowmax-driven online-softmax updates and freezing the running maximum on most blocks while retaining baseline updates on a small set of informative blocks—does not introduce systematic degradation in accuracy under our evaluation protocol. The few observed regressions are small in absolute magnitude for most tasks and appear dataset/model dependent, indicating that the approximation is stable across the tested tasks/models.. Table 2. Accuracy results of VFA on different models. Tasks MATH500 MMLU abstract algebra MMLU college computer science MMLU computer security HumanEval CMMLU college mathematics CMMLU philosophy Avg.
5.4
Qwen3-30B
Qwen3-8B
Llama3-8B
Baseline
VFA
Baseline
VFA
Baseline
VFA
0.424 0.75 0.75 0.81 0.7683 0.619 0.8571
0.42 0.76 0.75 0.83 0.7439 0.619 0.8476
0.25 0.61 0.72 0.80 0.6402 0.5905 0.8286
0.262 0.61 0.70 0.82 0.628 0.6 0.8190
0.134 0.36 0.45 0.82 0.372 0.3810 0.6286
0.146 0.35 0.46 0.82 0.372 0.3905 0.6190
−0.0011
+0.0000
+0.0017
Ablation Studies
We conduct ablations on MMLU abstract algebra to isolate the contribution of the approximation stage used for m-initialization and the subsequent control-flow decisions (reordering and max-freezing). Unless otherwise specified, all variants share the same configuration, numerical precision, and evaluation protocol, so that differences are attributable to the ablated component. Baseline and full method. Baseline (FA2) corresponds to the standard FlashAttention-style online softmax without our approximation stage. VFA (full, sabsmax) is the complete proposed pipeline, using the signed-absmax block representation kjrepr = sabsmax(Kj ) for m-initialization and the reordered schedule with selective max updates on special blocks. As shown in Table 3, the full method matches or slightly improves the baseline on this task, indicating that the approximation stage and the max-freezing mechanism do not harm task accuracy under the evaluated setting. Ablating the kjrepr construction.
We first ablate the choice of key-block representation kjrepr , which (0)
determines the approximation scoreapprox = Qi (kjrepr )⊤ used to initialize mi . The objective is to test ij whether preserving sign information and capturing extreme values per feature dimension are critical for producing a useful upper-bound-like scale for exponentiation.
9/19
• VFA (repr = Kmax ) replaces sabsmax with per-dimension maximum of Kj . • VFA (repr = Kmean ) uses per-dimension mean of Kj , which smooths extremes but may under-estimate the true maximum score. • VFA (repr = |K|max , unsigned) uses per-dimension absolute maximum but discards sign, potentially destroying alignment information between Q and K. All three alternatives lead to a pronounced drop in accuracy relative to VFA (full). This suggests that the approximation stage is highly sensitive to the representational choice: preserving the signed extreme (via sabsmax) appears necessary to produce a reliable initialization of mi . In contrast, mean-based representations are overly conservative, and unsigned absmax loses directional information, both of which can yield a poor scale for the subsequent exponentiation and accumulation. These variants collapse because the initialization becomes systematically miscalibrated, and the subsequent max-freezing prevents correction. Removing m-initialization. VFA (w/o m-init) disables the approximation stage and falls back to the (0) default initialization mi = −∞ (or the baseline initialization used in our implementation). This ablation verifies whether our gains stem from early acquisition of a high-quality mi . The substantial degradation indicates that m-initialization is not merely an optional optimization; rather, it is a necessary condition for max-freezing to remain accurate, because freezing m without a good initial scale can miscalibrate the exponentiation and bias the online accumulation. Ablating the reordered schedule. VFA (w/o reorder) keeps the same approximation-based initialization but removes the special-block-first schedule (i.e., it processes blocks in the default sequential order). This ablation tests whether early processing of empirically high-mass regions (sink/local) is important for correcting the initialization and absorbing rare large scores via exact max updates. The observed accuracy drop suggests that ordering is not a cosmetic change: prioritizing special blocks early likely increases the chance of capturing true maxima before max-freezing is applied broadly. Replacing row-wise approximation with block-wise approximation on Q. Finally, we ablate the left operand in the approximate score computation. The default approximation produces a row-wise score vector, scoreapprox = Qi (kjrepr )⊤ , yielding one approximate score per row of Qi and enabling per-row initialization of ij mi . We replace this with a cheaper block-wise score scoreapprox = qirepr (kjrepr )⊤ , where qirepr summarizes Qi ij (e.g., absmax(Qi )) and the resulting scalar score is broadcast to all rows. This ablation evaluates whether a coarse block-level summary of Qi is sufficient for initializing a per-row maximum. The results indicate a clear accuracy loss, implying that preserving per-row variation in Qi is important for building a reliable initialization of mi ; collapsing Qi to a single vector over-smooths row-specific maxima and reduces the fidelity of the approximation stage. Takeaway. Overall, the ablation results confirm that (i) the approximation stage is a key enabler of VFA, (ii) the signed-absmax key representation is crucial for maintaining accuracy, and (iii) both the special-block-first schedule and the row-wise (rather than block-wise) approximation contribute meaningfully to the robustness of m-initialization under max-freezing. Table 3. Ablation on MMLU abstract algebra: block representation kjrepr Methods Baseline (FA2) VFA (full, sabsmax) VFA (repr = Kmax ) VFA (repr = Kmean ) VFA (repr = |K|max , unsigned) VFA (w/o m-init) VFA (w/o reorder) VFA (reprQ = absmax(Qi ); block-wise score) VFA (reprQ = sabsmax(Qi ); block-wise score) VFA (reprQ = Qmean ; block-wise score)
MMLU abstract algebra 0.75 ± 0.0435 0.76 ± 0.0429 0.23 ± 0.0423 0.21 ± 0.0409 0.22 ± 0.0416 0.22 ± 0.0416 0.24 ± 0.0429 0.22 ± 0.0416 0.22 ± 0.0416 0.65 ± 0.0479
10/19
5.5
Data Analysis
(a) Running-max stabilization position ji⋆ .
(b) Attention score heatmap P .
Figure 2. Empirical evidence supporting the sink+local reordering strategy.
(a) Block similarity of Q measured by cosSim(X).
(b) Block similarity of K measured by cosSim(X).
Figure 3. Intra-block similarity for Q and K blocks using the SpargeAttention metric.
(a) ℓ2 -norm statistics of Q (token/row-level).
(b) ℓ2 -norm statistics of K (token/row-level).
Figure 4. Magnitude variation within Q and K blocks visualized by ℓ2 -norm statistics. 5.5.1
Where the Running Maximum Stabilizes
To understand why the proposed block reordering is effective, we analyze the evolution of the online-softmax running maximum during the key-block scan. Recall that FlashAttention-style online softmax maintains a (j) per-row running maximum mi as blocks are processed. For each query block Qi , we define the stabilization position as the first block index at which the running maximum reaches its final value: o n (j) (T ) ji⋆ ≜ min j mi = mi c , (1) where the equality is evaluated elementwise for per-row maxima. Intuitively, ji⋆ indicates how early the true maximum score is encountered in the scan order; if ji⋆ concentrates on a small subset of blocks, prioritizing these blocks can reduce the need for repeated max updates and rescale operations. Figure 2a summarizes the empirical distribution of ji⋆ . We observe a strong concentration of stabilization events in the initial blocks and the local block (i.e., the block aligned with the query position under causal/windowed attention). In other words, for the vast majority of query rows, the running maximum attains its final value either near the beginning of the key-block scan or when processing the local region. This provides direct evidence that the default sequential order spends substantial effort updating mi in blocks that are unlikely to change the maximum. 11/19
5.5.2
Attention Mass Concentration: Sink and Local Patterns
The above finding is consistent with the qualitative structure of attention weights. Figure 2b visualizes the attention weight patterns (the P matrix) as a heatmap. We observe that large attention mass concentrates on a small number of regions, most prominently the early “sink” positions and the diagonal local neighborhood. (T ) Such structure implies that large dot-product scores (and thus the eventual maxima that determine mi c ) are disproportionately likely to arise from these regions, while distant off-diagonal blocks contribute comparatively small weights. 5.5.3
Implications for Block Reordering
Together, the stabilization analysis and the heatmap patterns justify our reordering strategy that prioritizes the initial (sink) block and the local block before scanning the remaining blocks. By bringing the blocks with the highest probability of containing the final maxima to the front of the scan, the running maximum can stabilize early, which in turn reduces the frequency of costly rowmax reductions and rescale updates in the online-softmax pipeline. This empirical evidence connects the observed attention structure to the kernel-level objective of relieving vector pressure, and motivates the design choices in Algorithm 1. 5.5.4
Why Block Representations Are Ineffective
To explain the poor accuracy of block-representation-based variants (Table 3), we perform a numerical analysis on the intra-block variability of Q and K. Our hypothesis is that a single vector summary (e.g., max/mean/absmax) cannot faithfully represent the diverse token-level directions and magnitudes within a block, and thus yields unreliable approximations for m-initialization. Block similarity. We first measure the similarity among vectors within the same block using the blocksimilarity metric adopted in SpargeAttention. Given a block matrix X (either a Q-block or a K-block), we define XX ⊤ , (2) cosSim(X) ≜ mean |max(XX ⊤ )| where the mean is taken over all entries of the normalized Gram matrix [29]. Figures 3a and 3b report the resulting block similarities for Q and K, respectively. We observe that the average block similarity is only around 0.6, indicating that vectors within the same block are far from being well-aligned and cannot be accurately summarized by a single representative direction. Norm variability. Direction diversity alone is not the full story: we also examine magnitude variation by plotting the ℓ2 norms of Q and K vectors. Figures 4a and 4b show that the ℓ2 norms fluctuate substantially across tokens/rows, revealing strong intra-block scale heterogeneity. Such large norm variation implies that extreme rows can dominate dot-product maxima, while mean- or max-based block summaries may either under-estimate or misrepresent these extremes. Implication for block representations. Taken together, the moderate block similarity (≈ 0.6) and the pronounced norm fluctuations provide a direct explanation for why block representations are ineffective in our setting. Because intra-block vectors are neither directionally coherent nor scale-homogeneous, a single block-level summary cannot preserve the information needed to approximate row-wise maxima reliably. Consequently, block-representation-based approximations can miscalibrate m-initialization and lead to degraded downstream accuracy, consistent with the ablation results. 5.5.5
Why m-Initialization Is Necessary
We further investigate why m-initialization is necessary for the proposed max-freezing scheme. The key challenge is that the true per-row maximum score is not always attained in the early (sink) or local blocks. When the global maximum occurs in a middle key block, a design that relies solely on early/local exact max updates can miss the correct scale, and freezing m thereafter can prevent later correction.
12/19
(a) Block-max peaks at the sink block (j = 1).
(b) Block-max peaks at the local block (j = i).
(c) Block-max peaks at a middle block (j ∈ / {1, i}).
Figure 5. Representative cases of block-maximum location along the key-block index j, motivating the need for m-initialization. To quantify this phenomenon, Figures 5a–5c plot the block maximum as a function of the key-block index j under three representative settings. Specifically, for each query block Qi and key block Kj , we define the block maximum as (j) m̃i ≜ rowmax(Sij ), Sij = Qi Kj⊤ , (3) (j)
and visualize how m̃i varies as j increases. Across the three cases, we observe a consistent pattern: the location of the maximum block score can occur at the initial (sink) block (j = 1), at the local block (j = i), and in the middle of the sequence. While prioritizing sink and local blocks captures many maxima early, the presence of middle-block maxima implies that a strategy without m-initialization would systematically underestimate the true maximum scale for a non-trivial fraction of queries. Under max-freezing, such underestimation is particularly harmful because subsequent blocks do not update m, and the exponentiation exp(Sij − mi ) becomes miscalibrated. These results motivate m-initialization as a coverage mechanism: by providing a sufficiently high initial (0) estimate of mi before the main scan, VFA can remain robust even when the true maximum arises in a middle block. In practice, m-initialization complements the sink+local reordering by mitigating the residual cases where maxima do not lie in the prioritized regions, which aligns with the accuracy gains observed in our ablation study.
13/19
6
VFA Extension
Algorithm 2 Vector Relieved Sparse Attention r c c Require: Query blocks {Qi }Ti=1 , Key blocks {Kj }Tj=1 , Value blocks {Vj }Tj=1 , initialization block count Tc1 , threshold λ r Ensure: Output blocks {Oi }Ti=1 1: /* Precompute key-block representations */ 2: for j = 1 to Tc1 do 3: kjrepr ← sabsmax(Kj ) 4: end for 5: for i = 1 to Tr do (0) (0) (0) 6: Initialize mi = −∞, Oi = 0, li = 0 7: Initialize mi,init = −∞ ▷ vector, one value per row of Qi 8: for j = 1 to Tc1 do 9: scoreapprox ← Qi (kjrepr )⊤ ▷ vector over rows of Qi ij 10: mi,init ← max(mi,init , scoreapprox ) ▷ elementwise max ij 11: end for (0) 12: mi ← mi,init 13: for j ∈ ⟨1, i, 2, 3, . . . , Tc ⟩ with j ̸= i in the tail do 14: Compute Sij = Qi Kj⊤ ▷ Attention scores
15: 16:
(j)
m̃i
= rowmax(Sij ) (j) (] j −1) (j) mi = max mi , m̃i
17: 18: 19: 20: 21: 22: 23: 24: 25: 26: 27: 28: 29: 30:
(j)
▷ Local maximum ▷ Running maximum ▷ j] − 1 denotes the previous iteration index in the j-loop
(j)
if m̃i − mi < ln(λ) then continue end if if j = 1 or j = i then (j) P̃ij = exp Sij − mi
▷ Skip this block
▷ Attention weights
(] j−1) (j) (] (j) j −1) −mi li = e mi li + rowsum(P̃ij ) ] ( j−1) (j) (j) (] j −1) mi −mi Oi = e Oi + P̃ij Vj
else (j) (] j −1) mi = mi (j) P̃ij = exp Sij − mi (j)
li
(] j −1)
= li
+ rowsum(P̃ij )
(j) (] j −1) Oi = Oi + P̃ij Vj
end if 31: end for (T ) (T ) 32: Oi = Oi c /li c 33: end for Tr 34: return {Oi }i=1
▷ Freeze m when j ∈ / {1, i}; j] − 1 is the previous j-loop iteration ▷ Attention weights ▷ No rescale factor ▷ No rescale factor
▷ Final normalization
Algorithm 2 extends VFA by incorporating dynamic block sparsification [26], resulting in VSA. While VFA reduces the per-block overhead of online softmax by eliminating most rowmax-driven max updates and rescale operations, it still evaluates all Tc key blocks for each query block. VSA further reduces compute and memory traffic by skipping blocks whose contributions to the softmax output are provably negligible, thereby combining two orthogonal sources of speedup: (i) fewer blocks processed via sparsification, and (ii) cheaper processing per block via vector-relieved online softmax. BLASST-style skip criterion. VSA adopts the BLASST principle of using online-softmax statistics to decide whether a block can be skipped. For each query block Qi and key block Kj , we compute the (j) score block Sij = Qi Kj⊤ and its per-row block maximum m̃i = rowmax(Sij ). Given the current running 14/19
(j)
(j)
(j)
maximum mi , the quantity (m̃i − mi ) controls a worst-case upper bound on this block’s exponentiated contribution relative to the current scale. If this value falls below a threshold ln(λ), the block is skipped: (j)
(j)
m̃i − mi
< ln(λ) ⇒ skip block (i, j),
(4)
(j)
which avoids computing P̃ij = exp(Sij − mi ), bypasses loading the corresponding Vj tile, and elides the subsequent P̃ij Vj accumulation. This skip rule is particularly attractive in a FlashAttention-style kernel because it reuses quantities already computed for numerical stabilization and introduces minimal additional control flow. How sparsification interacts with max-freezing. A subtlety is that VFA intentionally freezes the running maximum for most blocks to avoid repeated rowmax reductions and rescale. In contrast, BLASST’s (j) (j) skip decision depends on m̃i and the current mi . VSA reconciles the two by (1) using m-initialization to (0) provide a strong initial scale mi , and (2) allowing exact max updates (and rescale) only on special blocks (sink/local), while still evaluating the skip condition on all visited blocks. As a result, VSA preserves VFA’s vector-relieved update path for non-special blocks, yet can skip a large fraction of blocks whose maxima are far below the initialized/updated running maximum. In practice, this design yields a “best of both worlds” behavior: if the attention mass concentrates on sink/local regions, most tail blocks are either (i) skipped outright, or (ii) accumulated with frozen m without rescale. Consequently, VSA can be viewed as a composable extension: VFA addresses the vector-limited components of online softmax, while BLASST-style sparsification reduces the number of blocks that must be processed. Preliminary experimental results for VSA are provided in the Appendix.
7
Conclusion and Limitation
Conclusion. This paper revisits FlashAttention-style online softmax with hardware efficiency in mind and identifies a practical bottleneck that emerges as matmul paths approach peak efficiency: per-tile statistic updates (rowmax/rowsum) and the rescale chain can become vector/SIMD limited and dominate end-to-end latency.. To address this bottleneck without changing the attention operator, we propose VFA, which reduces the frequency of rowmax-driven running-maximum updates via three lightweight mechanisms: (i) a fast m-initialization based on key-block representations, (ii) sink+local key-block reordering to prioritize high-impact regions, and (iii) max-freezing for the remaining blocks to bypass repeated reductions and rescale operations. We further demonstrate that VFA composes naturally with BLASST-style block skipping [26], yielding VSA that combines fewer processed blocks with cheaper per-block statistic updates. Our results confirm that VFA delivers tangible speedups in regimes where online-softmax statistic updates are a primary bottleneck, while preserving model accuracy on representative benchmarks. Relative to the C16V32 baseline, our C8V32/C4V32/C4V16 variants achieve up to ∼ 2× speedup; C4V16 is currently constrained by exponent throughput, and in a hypothetical scenario with higher exp16 capacity its speedup could increase (up to ∼ 6×). Limitations and future work. VFA leverages the empirical observation that the running maximum often stabilizes early around sink/local regions; workloads with highly non-local or adversarial attention distributions may reduce the effectiveness of reordering and max-freezing. A promising direction is to design adaptive schedules that detect stabilization online and dynamically adjust the freezing strategy. Initialization robustness. Although the proposed m-initialization mitigates cases where the true maximum occurs in middle blocks, its approximation quality can vary across layers, heads, and tasks. Future work includes exploring stronger yet still lightweight representations (e.g., multi-prototype or norm-aware summaries) and calibrating initialization policies per head/group. Evaluation scope. While our experiments indicate no observable quality degradation on the tested benchmarks, a broader evaluation covering more models, longer-context tasks, and additional generation settings would strengthen the conclusions. In particular, future work should include more extensive numerical stress tests for stability under extreme sequence lengths and diverse prompt distributions.
15/19
Acknowledgments Disclosure of AI-assisted editing. We used an LLM-based assistant for English polishing (grammar and readability) on selected passages. The technical contributions (method, performance evaluation, experiments, and conclusions) are entirely the authors’ original work, and all edits were reviewed by the authors, who take full responsibility for the manuscript.
References 1. F. Abecassis, A. Agrusa, D. Ahn, J. Alben, S. Alborghetti, M. Andersch, S. Arayandi, A. Bjorlin, A. Blakeman, E. Briones, et al. Pretraining large language models with nvfp4. arXiv preprint arXiv:2509.25149, 2025. 2. J. Ainslie, J. Lee-Thorp, M. De Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245, 2023. 3. K. Alexandridis, V. Titopoulos, and G. Dimitrakopoulos. Flash-d: Flashattention with hidden softmax division. arXiv preprint arXiv:2505.14201, 2025. 4. I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020. 5. K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, L. Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020. 6. T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023. 7. T. Dao et al. flash-attention: Fast and memory-efficient exact attention with io-awareness. GitHub repository. https://github.com/Dao-AILab/flash-attention. 8. T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022. 9. A. Hatamizadeh, G. Heinrich, H. Yin, A. Tao, J. M. Alvarez, J. Kautz, and P. Molchanov. Fastervit: Fast vision transformers with hierarchical attention. arXiv preprint arXiv:2306.06189, 2023. 10. K. Hong, G. Dai, J. Xu, Q. Mao, X. Li, J. Liu, K. Chen, Y. Dong, and Y. Wang. Flashdecoding++: Faster large language model inference on gpus. arXiv preprint arXiv:2311.01282, 2023. 11. H. Jiang, Y. Li, C. Zhang, Q. Wu, X. Luo, S. Ahn, Z. Han, A. H. Abdi, D. Li, C.-Y. Lin, et al. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481–52515, 2024. 12. A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020. 13. A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434, 2024. 14. A. Liu, A. Mei, B. Lin, B. Xue, B. Wang, B. Xu, B. Wu, B. Zhang, C. Lin, C. Dong, et al. Deepseek-v3. 2: Pushing the frontier of open large language models. arXiv preprint arXiv:2512.02556, 2025. 15. E. Lu, Z. Jiang, J. Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang, et al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502.13189, 2025.
16/19
16. Y. Luo, J. Huang, Y. Cheng, Z. Yu, K. Zhang, K. Hong, X. Ma, X. Wang, A. Tong, G. Hu, et al. Hifloat4 format for language model inference. arXiv preprint arXiv:2602.11287, 2026. 17. M. Milakov and N. Gimelshein. arXiv:1805.02867, 2018.
Online normalizer calculation for softmax.
arXiv preprint
18. P. Nawrot, R. Li, R. Huang, S. Ruder, K. Marchisio, and E. M. Ponti. The sparse frontier: Sparse attention trade-offs in transformer llms. arXiv preprint arXiv:2504.17768, 2025. 19. A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. 20. A. Roy, M. Saffar, A. Vaswani, and D. Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53–68, 2021. 21. J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, and T. Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658–68685, 2024. 22. N. Shazeer. Fast transformer decoding: One write-head is all you need, 2019. URL https://arxiv. org/abs, 1911. 23. Y. Sun, Z. Li, Y. Zhang, T. Pan, B. Dong, Y. Guo, and J. Wang. Efficient attention mechanisms for large language models: A survey. arXiv preprint arXiv:2507.19595, 2025. 24. A. Tseng, T. Yu, and Y. Park. Training llms with mxfp4. arXiv preprint arXiv:2502.20586, 2025. 25. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 26. J. Yuan, C. Shinn, K. Xu, J. Cui, G. Klimiashvili, G. Xiao, P. Zheng, B. Li, Y. Zhou, Z. Ye, et al. Blasst: Dynamic blocked attention sparsity via softmax thresholding. arXiv preprint arXiv:2512.12087, 2025. 27. T. Zadouri, M. Hoehnerbach, J. Shah, T. Liu, V. Thakkar, and T. Dao. Flashattention-4: Algorithm and kernel pipelining co-design for asymmetric hardware scaling. arXiv preprint arXiv:2603.05451, 2026. 28. M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297, 2020. 29. J. Zhang, C. Xiang, H. Huang, J. Wei, H. Xi, J. Zhu, and J. Chen. Spargeattention: Accurate and training-free sparse attention accelerating any model inference. arXiv preprint arXiv:2502.18137, 2025.
17/19
Appendix VSA Table 4. Accuracy results on MMLU computer security. Each cell is reported as acc ± 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. Method Baseline (FA2)
MMLU computer security 0.81 ± 0.0394
Blasst (λ = 1e−1) Blasst (λ = 3e−1) Blasst (λ = 5e−1) Blasst (λ = 9e−1)
0.8 ± 0.0402(32.3%) 0.84 ± 0.0368(44.0%) 0.83 ± 0.0378(49.5%) 0.83 ± 0.0378(56.1%)
Blasst SWA (λ = 1e−1) Blasst SWA (λ = 3e−1) Blasst SWA (λ = 5e−1) Blasst SWA (λ = 9e−1)
0.8 ± 0.0402(37.4%) 0.81 ± 0.0394(48.9%) 0.79 ± 0.0409(54.1%) 0.82 ± 0.0386(60.5%)
Blasst FA4 (λ = 1e−1) Blasst FA4 (λ = 3e−1) Blasst FA4 (λ = 5e−1) Blasst FA4 (λ = 9e−1)
0.82 ± 0.0386(32.1%/96.8%) 0.83 ± 0.0378(43.4%/97.0%) 0.81 ± 0.0391(48.5%/97.0%) 0.82 ± 0.0385(54.1%/97.1%)
Blasst Rowskip (λ = 1e−1) Blasst Rowskip (λ = 3e−1) Blasst Rowskip (λ = 5e−1) Blasst Rowskip (λ = 9e−1)
0.83 ± 0.0378(62.3%) 0.83 ± 0.0378(68.1%) 0.85 ± 0.0359(70.5%) 0.8 ± 0.0402(73.5%)
VSA (λ = 1e−3) VSA (λ = 3e−3) VSA (λ = 1e−2) VSA (λ = 1e−1)
0.84 ± 0.0368(16.1%) 0.84 ± 0.0368(31.1%) 0.84 ± 0.0368(52.8%) 0.78 ± 0.0368(84.6%)
Accuracy Results Table 4 reports task accuracy in the form acc ± std, followed by efficiency-related statistics in parentheses. Unless otherwise stated, the parenthesized number denotes the block sparsity, i.e., the fraction of (i, j) attention tiles that are skipped by a BLASST-style rule. For methods that additionally introduce a skip-rescale mechanism (e.g., Blasst FA4), we report two numbers as (sparsity/rescale-skip): the first is the block-skip ratio; the second is the fraction of non-skipped tiles for which the online-softmax rescale chain is bypassed (equivalently, the scale factor becomes 1). Concretely, block sparsity is the fraction of tiles skipped by the BLASST test, which bypasses softmax evaluation and P̃ij Vj accumulation. In contrast, rescale-skip measures how often a tile is still processed but the running-maximum update is suppressed (or the used max is kept unchanged), eliminating the vector/SIMD-heavy rescale updates on (Oi , li ). Blasst (λ). Rows labeled Blasst correspond to the vanilla BLASST integration. For each tile, BLASST evaluates the criterion (j) (j) m̃i − mi < ln(λ), (5) and triggers continue when the inequality holds. In this branch, the kernel bypasses P̃ij computation, avoids loading the corresponding Vj tile, and elides the P̃ij Vj accumulation. Therefore, the percentage in parentheses for Blasst is exactly the empirical block sparsity (the fraction of skipped tiles) under the given λ. As λ increases, the threshold becomes easier to satisfy, typically yielding higher sparsity but potentially larger accuracy degradation, reflecting the standard accuracy–sparsity trade-off of softmax-thresholded pruning. Blasst SWA (λ): sink+local reordering. Rows labeled Blasst SWA retain the same BLASST skip rule but changes the scan order to prioritize the sink and local blocks: j ∈ ⟨1, i, 2, 3, . . . , Tc ⟩
(j ̸= i in the tail).
(6) 18/19
(j)
Because BLASST’s decision depends on the running maximum mi , reordering can materially change sparsity: bringing high-mass regions earlier tends to increase mi sooner, making it more likely that later tiles satisfy (j) (j) m̃i − mi < ln(λ) and are skipped. Hence, the parenthesized percentage for Blasst SWA again reports block sparsity, but under a different scan schedule that can alter both sparsity and accuracy at the same λ. Blasst FA4 (λ): block skip + rescale skip. Rows labeled Blasst FA4 implement a two-stage saving mechanism. First, it applies the standard BLASST block-skip test and continues on tiles deemed negligible. Second, for tiles that are not skipped, it optionally skips the rescale chain when the running maximum changes only slightly: (j) (j−1) (j) (j−1) ≤ τ ln 2 ⇒ mi,used ← mi , (7) mi − m i (j−1)
(j)
−mi,used so that the scale factor becomes emi = 1. This removes the vector/SIMD-heavy rescaling of (Oi , li ) (j) while still computing P̃ij = exp(Sij − mi,used ) and accumulating P̃ij Vj . Accordingly, each cell for Blasst FA4 reports (sparsity/rescale-skip): the first number is the BLASST block sparsity; the second is the rescale-skip rate among the non-skipped tiles.
Blasst Rowskip (λ): row-wise thresholding. Rows labeled Blasst Rowskip refine BLASST from block-level decisions to row-level decisions. Instead of skipping an entire tile when the block maximum is small, it constructs a per-row keep mask (j)
(j)
row keep ← m̃i − mi
≥ ln(λ) ,
(8)
and sets the scores of skipped rows to −∞ (thus contributing zero mass after exponentiation), while applying rescale only to the kept rows. If no row is kept, the kernel continues and the entire tile is skipped. Because the pruning granularity differs, the parenthesized percentage for Blasst Rowskip reports row-wise sparsity (the fraction of rows suppressed by row-level thresholding), rather than tile-level sparsity. VSA (λ): composing VFA with sparsification. Rows labeled VSA correspond to Algorithm 2, which composes the vector-relieved online-softmax pipeline (VFA: m-initialization, sink+local reordering, and max-freezing) with BLASST-style thresholded skipping. At a high level, VSA targets two orthogonal sources of speedup: (i) fewer tiles processed via BLASST skipping, and (ii) cheaper per-tile load overhead via reduced rowmax/rescale frequency.
19/19