AdaSplash: Adaptive Sparse Flash Attention
Authors: Nuno Gonçalves
Marcos V Treviso
Andre Martins
Venue: ICML 2025 oral (2025)
Keywords: Sparse Attention
Flash Attention
Adaptive Sparsity
Long Context Transformers
Primary area: deep_learning->attention_mechanisms
TL;DR: An efficient flash attention implementation for adaptive sparsity.

Abstract: The computational cost of softmax-based attention in transformers limits their applicability to long-context tasks. Adaptive sparsity, of which $\alpha$-entmax attention is an example, offers a flexible data-dependent alternative, but existing implementations are inefficient and do not leverage the sparsity to obtain runtime and memory gains. In this work, we propose AdaSplash, which combines the efficiency of GPU-optimized algorithms with the sparsity benefits of $\alpha$-entmax. We first introduce a hybrid Halley-bisection algorithm, resulting in a 7-fold reduction in the number of iterations needed to compute the $\alpha$-entmax transformation. Then, we implement custom Triton kernels to efficiently handle adaptive sparsity. Experiments with RoBERTa and ModernBERT for text classification and single-vector retrieval, along with GPT-2 for language modeling, show that our method achieves substantial improvements in runtime and memory efficiency compared to existing $\alpha$-entmax implementations. It approaches---and in some cases surpasses---the efficiency of highly optimized softmax implementations like FlashAttention-2, enabling long-context training while maintaining strong task performance.

ADASPLASH: Adaptive Sparse Flash Attention
Nuno Gonçalves' Marcos Treviso
André F. T. Martins 123
Abstract
FlashAAleeniinn22 (CUDA)
The computational cost of softmax- -bannd atten- tion in transformers limits their applicability to long-context tasks. Adaptive sparsity, of which o-entmax attention is example, offers flexi- ble data-dependent alternative, but cxisting imple- mentations are inefficient and do not leverage the sparsity to obtain runtime and memory gains. In this work, we propose ADASPLASH, which com- bines the efficiency of GPU-optimized algorithms with the sparsity benefits of a-entmax. We first introduce hybrid Halley-bisection algorithm, sulting in a 7-fold reduction in the number of itera- tions needed to compute the a-entmax transforma- tion. Then, we implement custom Triton kernels to efficiently handle adaptive sparsity. Experi- ments with RoBERTa and ModernBERT for text classification and single-vector retrieval, along with GPT-2 for language modeling, show that method achieves substantial improvements in runtime and memory efficiency compared to ex- isting d-entmax implementations. It approaches- and in some cases surpasses-thh the efficiency of highly optimized softmax implementations like F enabling long-context training while maintaining strong task performance.!
Relalive Speed (vs, FlastVlltertion 2)
1.2
FlashAttentionn22 (Triton)
AdaSplash (Triton)
1.0
0.8
0.6
0.4
50
60
70
80
90
100
16x18 Block Sparsity (%)
Figure I. Runtime (Fwd+Bwd) as function of input sparsity for non-causal attention. While the highly-optimized FlastAttention- maintains a constant runtime across varying levels of sparsity, ADASPLASH effectively leverages sparsity to obtain speed-ups, eventually outperforming FlashAttemtionn22 as sparsity grows.
In fact, previous research shows that attention probabilities tend to peak around a small number of tokens (Voita et al., 2019; Treviso et al., 2022), which suggests that model per- formance and computational efficiency can be increased by leveraging attention sparsity. This has motivated methods that predefine sparse masks (Beltagy et al., 2020; Zaheer et al., 2020b), rely on clustering based strategies (Kitaev et al., 2020), or low-rank approximate attention (Choroman- ski et al., 2021; Peng et al., 2021; Xiong et al., 2021; Chen et al., 2021). Some of these techniques show the potential of sparsity mitigate memory and computation bottlenecks, but they often require architectural modifications crude approximations, limiting their flexibility and generality.
1. Introduction
Central to the success of transformers (Vaswani et al., 2017) lies the attention mechanism, where each token in a se- quence attends directly to every other token. Attention prob- abilitics are computed through the softmax transformation, which always assigns a nonzero probability to every token. However, for long context inputs, the accumulation of small probabilities can lead dispersion (Veličković et al., 2025).
A related line of research explores adaptive and differen- tiable sparse activations as surrogates of softmax, such as sparsemax (Martins Astudillo, 2016) and, more broadly, the o-entmax family (Peters et al., 2019; Correia al., 2019). By assigning zero probability to irrelevant tokens, these activations eliminate their residual influence, reduc- ing the dilution of attention scores and potentially improv- ing both performance and interpretability. Unfortunately, existing algorithms and implementations for these adap- tive sparse activations do not exploit the sparsity, being slower than sofumax based attention and struggling to scale
IInttituto Superior Técnico, Universidade de Lisboa, Portu- gal "Instituto de Telecomunicaçães, Lisbon, Portugal "Unbabel, Lisbon, Portugal. Correspondence to: Nuno Gonçalves <nuno. wamno...oncclleseeeee www.eeeo..inciicaaaaa...
Proceedings of the 42nd International Conference on Machine I Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Code: https: /tgpp::/gittee //ggt//bbbttt/tt com wwtp///peeee/////eeeeeeeeeeeeeee adasplash

---

ADASPLASH: Adaptive Sparse Flash Attention
effectively with context length, primarily due the lack of hardware-optimizze implementations like FlashAttention- 2 (Dao, 2024) support from programming models like FlexAttention (Dong et al., 2024).
in the following way (Vaswani et al., 2017):
QK vd
0
V Roxd
(I)
This paper addresses this problem by providing new algo- rithms and implementations to improve the computational efficiency of the family of a-entmax activations. Our main contributions include faster and GPU-friendly algorithm for calculating a-entmax, alongside Triton kernel (Tillet et al., 2019) for computing entmax-based -ased attention, which we call ADASPLASH. In particular, ADASPLASH advances the goal of supporting training of adaptively sparse mod- els with longer context lengths, as shown in Figure 1. We demonstrate the potential and scalability of our approach through experiments with synthetic data and with several natural language processing benchmarks for encoder-only and decoder-only models, achieving substantial improve- ments over previous a-entmax implementations and ap- proaching (sometimes surpassing) the efficiency of softmax- based attention with FlashAttention-2. 2, with strong perfor- mance on downstream tasks.
SERexa
The π transformation usually maps rows to distributions, with (T()); softmax (s:), being a common choice. For decoder-only models, S is masked in order to ignore the contribution from future tokens. Notably, a naive implemen- tation of Equation I leads to a O (n²) time and memory complexity for training.
2.3. FlashAttention
To address the costs of naive attention implementations, Dao al. (2022) introduced FlashAttention, an algorithm that avoids the materialization of quadratic matrices via a GPU-aware implementation of online softmax (Milakov & Gimelshein, 2018), bringing the overall memory complexity to O (n). Subsequent versions of Flash further improved GPU usage by reordering the loops, reducing the number of non-GEMM (general matrix multiply) operations (Dao, 2024), and exploiting the asynchronicity and support for FPS low-precision on the new Hopper GPUs (Shah et al., 2024). The key idea of FlashAttention is to split the imputs Q, K, V into blocks, load them from slow GPU high band- width memory (HBM) to the fast GPU on-chip SRAM, then compute the attention output regarding those blocks and, at the end, scale the output by the right normalization factor.
2. Background
2.1. Hardware Performance
Modern GPUs, such as the Nvidia H100, are designed for efficient parallel computation using a hierarchical memory architecture, with high-bandwidth memory (HBM) provid- ing large capacity but slower access compared to the smaller, faster on-chip SRAM. Efficient use of SRAM is critical to minimize the memory bottlenecks caused by frequent HBM accesses. GPUs execute operations (kernels) via thousands of threads organized into thread blocks, where data is loaded from HBM into SRAM for computation before being writ- ten back. Kernel fusion is a key optimization strategy that combines multiple operations into a single kernel, reducing intermediate HBM accesses by directly computing and stor- ing final results. While compilers like torch, compile can automate fusion for simple operations (Ansel al., 2024), complex tasks such as attention mechanisms require cus- tom strategies to reorder operations and optimize memory usage effectively. Our method leverages this GPU mem- ory organization by implementing block-wise computations, recomputation strategies, and kernel fusion specifically tai- lored for sparse attention, as detailed in $ and $3.2.2.
2.4. Sparse Attention
The original softmax-based attention is dense, i.e., it puts some probability mass on all tokens-not only a compu- tational disadvantage, but also making interpretation and generalization harder (Voita et al., 2019; Treviso et al., 2022: Veličković et al., 2025). An alternative softmax is the o-entmax transformation (Peters et al., 2019), which is differentiable and leads to sparse outputs:
o-eetmax (s) = [(a 1)s - -1]/-
(2)
where H+ is the ReLU function, and E R is a normalizing constant to ensure the output is valid probability distribu- tion. Importantly, entries with score 8; S -= get exactly zero probability. In the limit a + 1, a-entmax recovers the softmax function, while for any value of a > this trans- formation returns increasingly sparser probability vectors. When 2, we recover the sparsemax transformation (Martins & Astudillo, 2016). However, in contrast to fixed sparse patterns, such as windowed sparse attention (Child et al., 2019; Beltagy et al., 2020) and block-sparse vari- ants (Zaheer et al., 2020b; Dao et al., 2022), a-entmax's 's sparsity patterns are dynamic and hence difficult exploit
2.2. Standard Attention
Given a set of matrices Q, K, V E Radd containing d- dimensional representations for n queries, keys and values, the dot- product self-attention single head is computed
2

---

ADASPLASH: Adaptive Sparse Flash Attention
in order reduce the quadratic burden of self-attention because we still need to materialize 5 QKT before applying the transformation.
Algorithm Halley-bisection algorithm for a-entmax.
1: Input: logits a E 12", param, a E IR, iterations T
2: Define f(r) := E, - 7] (a==) -1
3: Set 8 f- (a - 1)
In the next section ($3), we outline ADASPLASH, our new method for computing a-entmax attention, along with a novel custom Triton kernel (Tillet et al., 2019) that enables efficient training of transformers for extremely long context lengths. As shown in $4, our implementation maintains com- petitiveness with state-of-the-ae of-the-art algorithms such FlashAt- tention by leveraging the sparsity given by a-entmax, ef- fectively exploiting the advantages of sparse attention at scale.
4: Initialize To = mmx( (s) -
Initialize Thi = max(8) 111-a
Initialize T (No + 7h)/2
repeat
8:
Compute To Th (r) (Equation 4)
9:
Compute TH = Hy (Equation 5)
10:
if E [o, Thil then
11:
T T TH
> (Halley's 's Update)
12:
else
13:
T T = (7 + Tm)
(Bisection Update)
3. ADASPLASH
14: end if
We start by revisiting the computation of a-entmax for gen- eral values of a in $3.1, and proposing new algorithm that has fast empirical convergence. We design an efficient Triton kernel in $3.2, dubbed ADASPLASII, that effectively leverages adaptive sparsity patterns in both the forward and backward passes of a-entmax in order to minimize runtime.
15: until T iterations
completed
16: Output: [s [--1 -1] 1((a11
ation. Achieving high precision often requires many itera- tions, resulting in frequent memory accesses, As a result, in memory- scenarios where the time taken is mostly determined by the number of memory accesses- such as in attention the number of iterations can significantly impact the runtime cost.
3.1. c-enimax Computation
In order to compute Equation 2 for given s E R", we need to find the threshold r E such that the resulting output sums to Mathematically, this is equivalent to finding the root of the following equation:
Halley- algorithm. In order to obtain a faster runtime, we propose hybrid algorithm for solving Equa- tion for any a > 1 that combines the convergence guar- antee of bisection with the faster convergence of Halley's method (Scavo & Thoo, 1995). As we show in $4. 1, this approach achieves significant wall-clock speed-ups while requiring fewer iterations to attain the same precision.
f(T) = Σ [(a - 1)s; -| -1/(1-1)) 11
(3)
Exact algorithms for a E 2). In particular, for 2. the computation is reduced to an Euclidean projection onto the probability simplex, for which efficient algorithms have been extensively studied (Held et al., 1974; Duchi et al., 2008; Condat, 2016). Similarly, for = 1.5, Peters al. (2019) introduced an exact sort-based algorithm. However, these methods either require complex data structures that are not efficiently handled in GPUs, or oorting-based algorithms, which require the materialization of the entire input.
The function defined in Equation 2 enjoys a cheap com- putation of its derivatives. Thus, methods that incorporate second-order information, such as Halley's method, can be leveraged to improve the approximation of at each itera- tion. Halley's method, which uses both the first and second derivatives, updates the solution using the following rule:
2f(r)f((r)
H₂(T) =T- -()))116))"
(5)
Bisection algorithm for a > 1. For general o, Blondel et al. (2019) introduced a bisection update rule to approxi- mate by iteratively refining its lower (7a) and higher (na) bounds:
where the derivatives are given as follows:
By(T)= = (To,7) 7) f() <0, (T, Thi) otherwise,
['(r)= = a-
Σ [(a 1)8; - -] 11(--11-
(6)
(4)
(" (a - 1)2
2-a
Σ [(a - 1)8; - 7 1/(a-1)-2
(7)
obtaining T = (no + Thi) after the last iteration. While the bisection algorithm is simple and effective, it converges at a linear rate (Kaufman & Lenker, 1986), meaning the absolute error decreases by approximately half each iter-
While Halley's 's method offers faster convergence under ideal conditions, it does not always converge, particularly when
3

---

ADASPLASH: Adaptive Sparse Flash Attention
Algorithm 2 ADASPLASH forward pass (w/o masking)
Staction Staction|/Srauk
Absollte Ertor Magriude
10
10
Staction (Chadert)
Hollyy doctkan (Dutas)
Hillyy bbcciknn (Staskan)
1: Require: Matrices Q, K, V E Rnxd in HBM, block sizes B,, Br, param. a E R
2: Divide Q into T, = [n/B,] blocks Q1 Qr, of
size B, x d
3: Divide K, V into T, = [m/ Be] blocks K KT.
10
Vi...., Vr. of size Bg x d
10"
Divide O Rand into T, blocks Is...., Or, of size
10 10
B, x d
10
15
20
25
5: Divide 7 into T, blocks T,,..., Tr, of size B,
Number herations
6: for i = 1 to T, do
7: Load Q; from HBM to on-chip SRAM
Figure 2. Comparison of mean absolute error magnitudes between Halley-bisectinn and Torch's bisection methods across iterations, measured against the exact solution for 1.5.
8: On chip, initialize O;
9: On chip, compute T; using Hybrid Halley's 's with pre- defined a, using a block version of Algorithm 1.
10: for j 1 to Te do
the initial guess is far from the solution. To ensure conver- gence, we introduce fail-safe mechanism that integrates the convergence guarantee of bisection: whenever Halley's 's method produces an update that moves the solution out of the bisection bounds, the algorithm reverts to bisection update (T). This ensures that the algorithm converges, even in the worst cases, while leveraging the cubic conver-
11: Load K,, V, from HBM to on-chip SRAM
Compute s = K E R B, xB
12:
Compute p!ù) = (a 1)s
1/4-1
13:
14: Accumulate O, O, pIO) V3
15: end for
16: Write O, and t; to HBM
gence of Halley's 's method wherever possible. We outline our hybrid algorithm in Algorithm 1.
17: end for
18: Return: Output O and r
Efficiency Benchmark. We compare the runtime of Halley-bisection against existing algorithms for computing o-entmax implemented in Torch. Specifically, we gener- ate random tensors from standard Gaussian distribution (u = 0, = 1) with a fixed sequence length of n 8192. For each configuration, we measure the average runtime over 1000 runs. Overall, we observe that Halley-bisectinn significantly more efficient than the standard bisection algorithm implemented in Torch. Halley-bisection achieves runtime of 2.38 ms, compared to 36.67 ms for the standard bisection algorithm, making it approximately 15x faster. In addition, Halley-bisection reduces memory usage by 1.75x, requiring only 512 MB compared to 896.15 MB for bisec- tion. Furthermore, in Figure 2 we show that Halley- bisection (a = 1.5) requires only 3 iterations to converge to machine precision for both the output and the gradient. On the other hand, the standard bisection algorithm takes 23 iterations to achieve the same precision for both cases.
O = However, since most of these operations are memory-hound the excessive number of HBM accesses leads to slow wall-clock times. Moreover, having to mate- rialize S and P in memory poses a major bottleneck, as their sizes quickly exceed GPU memory capacity when the sequence length n increases. To address these issues and speed up o-entmax attention on hardware accelerators like GPUs, we propose an algorithm that reduces HBM reads and writes while producing the same outputs as the naive implementation.
3.2.1. FORWARD PASS
We outline the forward pass in Algorithm 2 (without mask- ing full-zero blocks, which we introduce later on this sec- tion). Concretely, given the inputs Q, K, V E Rxd stored in HBM, the goal is to compute the attention out- put O Rexd efficiently and write it back to HBM. Akin to the approach taken in FlashAttention (Dao et al., 2022), we employ two well-known techniques ceaiiiquuss-lliing and recom- putation to address the challenge of materializing the matrices S E R and P E Rexn.
3.2. Flash o-entmax Attention
Given an algorithm compute the entmax mapping that requires T' iteration steps, a naive implementation of entmax attention proceeds as follows: (1) multiply S = QK E gax= and write the result to slow HBM on the GPU; (2) load S from HBM T times to compute т: (3) load S from HBM again, and write the result P = a-entmax (S) to HBM; (4) perform a matrix multiplication to get the output
Tiling. The key idea involves splitting the inputs K, V into smaller blocks, and then computing attention block by block. We start by loading only Q and K from the slower HBM to the faster SRAM to compute T E R using the
4

---

ADASPLASH: Adaptive Sparse Flash Attention
Valley-bisection algorithm (Alg. 1). In order to use the aforementioned algorithm, we need accumulate three val- ues: f(T), f(T), (T) Since f. well as its derivatives, is additive over its inputs, their computation can also be computed in blocks. Let B, and Be be the row and col- umn block sizes, respectively, and define T, = B, and To = Tn/ B.]. Divide Q into Q1, QI,..., QT. blocks, and K into K1. i..... Kr. blocks. Then, f(T) can be computed as:
set of all indices j' such that L///To j. We construct block mask matrix M {0, 1 as follows:
Mij- = {1 if 3rezzn : Sa > Tps otherwise,
Importantly, M is created dynamically after small prede- fined number of Halley- sellly-biseetiion iterations.
While the introduction of M breaks the linear memory com- plexity of dense fused-attention by requiring T, x Te extra memory, the overhead is still manageable as only con- tains binary values and is substantially smaller than the full P E R.** matrix. Furthermore, M needs to be materi- alized only once and its memory can be shared across all attention layers. To leverage M in practice, we propose to create two pointer-incremen lookup tables:
Tc f(T,) = Σ S() 1=1
(8)
where 5/01 = K E B.XB. and represents the zth sliced block of T with size T,.. Thus, these quantities do not need to ever be materialized and can be accumulated directly in fast memory. Afterwards, we load V to compute the attention output O for those blocks. In contrast to FlashAttention, our approach requires loading K to com- pute S at least two additional times. Therefore, the forward pass is bound to always be slower than FlashAttention's due to the extra HBM reads and computation.
1. K2 = {i I My = 1}: A table containing the row indices i of M that lead non-null blocks in p())
2. Q = {j My = 1}} A table containing the column indices j of M that lead to non-null blocks in p(O)
Recomputation. In order to avoid the materialization of the matrices S and P. we recompute them again in Algo- rithm 1, which is used to compute т, and also recompute them for obtaining the gradients for the backward pass. By doing this we are increasing the required FLOPs to re- duce the maximum amount of memory required. While this might suggest an increase in runtime, the opposite is observed (Dao et al., 2022). Despite the need for additional matrix multiplications, the reduction in total HBM reads and writes more than offsets the extra FLOPs, leading to improved performance overall.
These tables enable efficient skipping of K and V blocks that do not contribute to the final attention output O, sig- nificantly reducing unnecessary computations. Moreover, the same mechanism can be extended to accelerate the back- ward pass, where gradients with respect to Q, K, and V are computed, which we describe next.
3.2.2. BACKWARD PASS
In FlashAttention (Dao al., 2022), the backward pass is executed using single kernel that parallelizes computation across batch, head, and sequence dimensions. However, fol- lowing Triton's 's official implementation of FlashAttention," we separate the backward pass into two kernels: one for dQ (the gradient w.r.t. Q) and another for d.K and dV (the gradients w.r.t. K and V).
Sparsity-aware implementation. The key challenge of a-entmax attention lies in finding the threshold т, which requires multiple evaluations of the function fT))
which, in turn, depends on the score matrix S. While our proposed Halley-bisection algorithm alleviates the number of iterations needed recompute s(") by providing faster empirical convergence, our current implementation still iterates over all blocks of S, including null blocks-blocks where the corresponding entries of the sparse attention matrix P are zero
Sparse Jacobian of a-entmax. The sparsity in the Jaco- bian of a-entmax plays crucial role in the backward pass. For p a-entmax(s), the Jacobian is (Peters et al., 2019)
da-entmax(s* (s) Diag(u) |ulli"
uu
(10)
Ös
Furthermore, empirical evidence from Jiang et al. (2024) and (Xiao et al., 2024) suggests that for long inputs (e.g., 128k tokens in LLaMa-3-8b), approximately 3% of the en- tries in P suffice to capture over 96% of the total attention, which motivates an approach to leverage the adaptive and unstructured sparsity of a-entmax attention weights. To this end, we propose to only compute necessary blocks of P by skipping the null blocks. Concretely, let Z(i) denote the set of all indices i' such that [f//T,] = i, and J (j) denote the
where u, = (p,) Importantly, this Jacobian is sparse and only depends on p. which, in turn, function of computed during the forward pass. We denote by U Raxa the matrix defined element-wise as Unx p2-a and by U, E BB..B. its (i. block. Using this information,
5

---

ADASPLASH: Adaptive Sparse Flash Attention
Bisection (Torch)
AdaSplash wio masking (Trilon)
FlashAttention22 (Trito))
AdaSplash (Triton)
10.0
Time (s)
FlashhAteention2 (CUDA)
0.1
OOM
MOOM
OOOM
JOOM
1k
2k
4k
8k
16k
32k
64k
Figure 3. Efficiency of algorithms for computing non-causal attention in terms of the average training step time for increasingly longer sequence lengths. We 1.5 for a-entmax based methods (Bisection and ADASPLASH).
the gradient w.r.t. the score matrix s:0) E R. can efficiently computed as:
Generalization to architectures: How does ADAS- PLASH perform when integrated with encoder-only and decoder-only models?
as = U. @ dd Diag( U.)
(11)
Effectiveness finetuning: Can ADASPLASH- pretrained models outperform or match their dense counterparts in short and long-context tasks?
where d.P())) V E BBB. with dO, B.
and V, R Re and o; R.. denotes the 1h block of the vector 8 T defined element-wise òj (xx Un Un))
4.1. Efficiency Benchmark
Efficient gradient computation. In ADASPLASH, instead of storing P, we store the lookup tables K. and Q computed during the forward pass, allowing us to to efficiently skip the computations of null blocks during backpropagation. Given dS,, the gradients for Q. K,, V, B. are computed follows using the pointer-increment lookup tables:
We compare the efficiency of ADASPLASH against Fltenintt and naive implementations of o-entmax. For a fair comparison, we also include a variant of FlashAttention-2 implemented Triton that follows closely our kernel implementation of ADASPLASH. We the number of iterations of ADASPLASH to 3 and Bisection to 10. As input, we generate random tensors from a Gaussian distribution (μ1 0), simulating attention scores with high level of sparsity by setting the Gaussian variance to o² 6
dQ: Σ ds(
(12)
JEQ,
dK,
Σ (
Qi
(13) of query vectors. Sequence lengths range from 1 to 64k, with a fixed head size of d = 64.
dV, Σ P()) dO,-
(14)
Runtime. We show the average training step time for each method in Figure 3. ADASPLASH demonstrates superior scalability, efficiently handling sequences up to 64k, unlike the Bisection method implemented in Torch, which runs out of memory beyond 4k context length. We also note that, as context length increases, the amount of block sparsity naturally increases as well, leading to an advantage for our method over both implementations of FtastAntet
iEK,
Hence, by splitting the backward pass into separate ker- nels and exploiting the sparsity of o-entmax through the Jacobian structure, can achieve efficient gradient com- putation. Overall, ADASPLASH allows users to choose between memory efficiency (without block masking) and computational speed (with block masking) depending on task requirements and hardware constraints. We provide detailed derivation of a-entmax attention's 's backward pass and its implementation in Appendix A.2.
4.2. Performance on Real Tasks
Encoder-only models, such as RoBERTa (Liu et al., 2019) and ModernBERT (Warner et al., 2024), exhibit higher at- tention sparsity than decoder-only models. making them well-suited for adaptive sparse attention mechanisms like ADASPLASH. Following ModemBERRT's evaluation setup, we opt to evaluate these models on standard NI.P tasks, such as text classification. natural language inference, textual similarity, and information retrieval. Moreover, following FlashAttention's 's evaluation setup (Dao et al., 2022), we also benchmark ADASPLASH with GPT-2, a decoder-only
4. Experiments
We evaluate ADASPLASH across various scenarios to show its computational efficiency and impact on downstream tasks. Our experiments address the following questions:
Performance efficiency: How does ADASPLASH com- pare with baseline methods in terms of runtime as sequence length and sparsity vary?
6

---

ADASPLASH: Adaptive Sparse Flash Attention
Table Results for single-vector retrieval models on different
Table 3. Runtime per epoch (hh:mm:ss) and peak memory usage (GB) for long document classification with different sequence lengths. In cases where the full batch could not fit memory. gradient accumulation was used. Memory values represent the effective peak memory required to process batch of samples.
tasks from the BEIR benchmark in terms of nDCG@ 10.
Model
Seq.
SciFact
NFC
FIQA
TREC-C
RoBERTa
512
51.7
23.1
27.8
60.1
RoBERTa (a 1.5)
512
50.8
24.2
27.6
71.0
RoBERTa (a 2.0)
512
52.2
23.8
25.7
65.5
Runtime (hh:mm:ss)
Sequence Length
ModernBERT
8192
57.7
22.4
25.7
67.6
Model
512
1024
2048
4096
8192
ModernBERT (a= = 1.5)
8192
58.4
25.7
29.6
75.2
ModemBERR (a- = 2.0)
8192
58.0
25.4
29.3
71.1
RoBERTa
2:39
5:00
9:35
18:36
35:51
RoBERTa (a = 1.5)
2:45
5:20
10:24
19:54
38:08
w/ Torch Bisect
4:51
8:44
22:48
1:11:53
4:12:34
Table 2. Long document classification performance (F) micro) with softmax and a-entmax atlention.
Memory (GB)
Sequence Length
512
1024
2048
4096
8192
Sequence Length
RoBERTa
6.75
11.43
20.35
37.49
75.00
RoBERTa (a = 1.5)
6.75
11,45
20.38
39.17
79.88
Model
512
1024
2048
4096
8192
w/ Torch Bisect
7.75
16.92
44.06
142.76
508.16
RoBERTa
71.5
74.4
75.1
77.9
79.2
RoBERTa (a- 1.5)
71.8
75.5
76.4
78.0
78.6
Long document classification. We fine-tuned pre- trained RoBERTa model (Liu et al., 2019) on the ECIHR (Chalkidis et al., 2019; 2021) dataset while progressively increasing the sequence length up to 8192 tokens. Positional embeddings were extended by repetition, following the ap- proach of Beltagy et al. (2020). As a baseline, we fine-tuned the model using standard entmax attention, we linearly increased the from 1.0 to 1.5 during training to ensure smooth convergence. The results, summarized in Table 2, show consistent improvement in model performance with longer context lengths. Notably, despite the base model being pretrained with standard atten- tion, a-entmax attention was capable of effectively learning the task, achieving a slightly higher miero F1 score than the model fine-tuned with standard attention up to a sequence length of 4096 tokens.
model, to assess its efficiency in autoregressive settings where attention patterns are denser. This ensures a compre- -based methods
hensive comparison with optimized while validating the benefits of sparsity across different ar- chitectures. We provide more training and evaluation details for each task in Appendix B.
attention. For a-
Continuous pretraining. We conducted continuous pre- training of RoBERTa-base and Modern on 2B tokens of the English subset of Fineweb-edu (Lozhkov et al., 2024) using ADASPLASH for E {1.5, 2). and PyTorch's 's loolledddot_prrduccc prroduuc.aateettion for a = 1.0. To en- sure a smooth transition from dense to sparse attention, we linearly increased a from a = 1.0 to the target values E 2.03 over the first IB tokens and kept it fixed afterwards. We provide more details the continuous pre- training phase in Appendix 1. including efficiency results.
Table compares the runtime per epoch and peak memory usage for different sequence lengths on the long document classification task. We report results for RoBERTa with Fttetinttn (a = 1), RoBERTa with ADASPLASH (a = 1.5), and RoBERTa using Torch's bisection-based implementation. ADASPLASH enables scalable training with a-entmax attention. Prior to this, implementations had to resort to Torch's 's bisection, which leads to both extremely slow runtimes or even out-of-memory problems, rendering it infeasible for most realistic training setups. In contrast, our method brings the cost of o-entmax attention down to the level of existing dense attention implementations, as both runtime and memory usage with ADASPLASH remain well aligned with those of FlashAttention--2.
Single-vector retrieval. We evaluate our pretrained models on single-vector retrieval performance using the BEIR benchmark (SciFact, NFCorpus, FiQA2018, TREC- COVID), following the setup in (Warner et al., 2024). Ta- ble highlights the performance of RoBERTa and Mod- eMBERT models using a-entmax attention terms of the standard nDCG@ 10 metric. ModernBERT with a = 1.5 consistently outperformed its dense counterpart, achieving the highest scores on all tasks, demonstrating its ability to focus on relevant signals effectively. While ModemnBERT with 2.0 remained competitive, its higher sparsity might have excluded relevant information, affecting task per- formance. Finally, sparse versions of ModernBERT achieve better results than the sparse versions of RoBERTa on all tasks, highlighting the benefit of modeling long contexts.
Language understanding. We also evaluate RoBERTa and ModemBERR models with a-entmax attention on the GLUE benchmark (Wang et al., 2018) in Appendix B.2. Overall, the results indicate that models with sparse attention
7

---

ADASPLASH: Adaptive Sparse Flash Attention
Table Results language modeling with GPT-2 in terms of final validation loss and accuracy the HellaSwag task (Zellers al., 2019), along with the average runtime per training step (in seconds) and peak memory usage (GB) per GPU.
38
0.6
Model
Val. Loss
HS Acc.
Runtime Memory
GPT-2
3.283
30.4
0.98
52.5
GPT-2 (a = 1.5)
3.263
30.6
1.03
52.5
w/ Torch sorting
3.61
73.8
02
w/ Torch bisection
7.78
77.6
8
9
10
11
12
-0.0
Layer
achieve comparable performance to their dense counterparts, which underscores the ability to efficiently train a-entmmx models without sacrificing accuracy.
Figure 4. Ratio of non-zero attention scores for GPT-2 (a = 1.5).
Language modeling. Following (Dao et al., 2022), we trained a small 124M GPT-2 model (Radford et al., 2019) from scratch on 10B tokens of the FineWeb dataset (Penedo et al., 2024) with a context length of 1024 tokens. For consistent evaluation between softmax and a-entmax at- tention, we also trained a softmax- GPT-2 to serve as baseline. After training, we evaluated both models on the HellaSwag task (Zellers et al., 2019). Table 4 presents side-by-side comparison of the final validation loss and accuracy on HellaSwag, along with runtime and memory usage numbers. Sparse GPT-2 achieves slight improve- ment in validation loss (3.263 vs. 3,283) and final accuracy (30.6% vs. 30.4%) compared to its softmax counterpart, while obtaining comparable runtime and memory efforts. Furthermore, our approach achieves a runtime comparable to the GPT-2 using the highly optimized FA2 (1.03 s/step vs. 0.98 s/step) and matches its memory footprint (52.5 GB), while outperforming the sorting and bisection variants by large margins in both speed (1.03 s/step vs. 3.61 and 7.78 s/step) and memory usage (52.5 GB vs. 73.8 and 77.6 GB). In Appendix B.4, we report all training and evaluation details, including the validation loss curves of each method.
rameter. For a = 2, the problem simplifies to a projection onto the probability simplex, a well-established optimization problem. Its solution forms the base of sparsemax (Martins & Astudillo, 2016), which can be efficiently computed using sorting and root-finding methods (Held et al., 1974; Condat, 2016; Liu & Ye, 2009). Moreover, for intermediate values such as a = 1.5. Peters et al. (2019) proposed an exact sorting-basee bbann algorithm along with an implementation of a bisection algorithm applicable to any a. However, these approaches remain suboptimal for long contexts due to slow convergence or reliance on complex data structures and sort- ing operations, which are difficult to optimize for hardware.
Sparse Attention Mechanisms. Efficient sparse atten- tion mechanisms have been widely studied reduce the quadratic cost of transformers. The Sparse Transformer (Child et al., 2019) introduces fixed windowed atten- tion that can be efficiently computed using CUDA ker- nels, strategy also adopted by Longformer (Beltagy et al., 2020), and BigBird (Zaheer et al., 2020a). However, data- dependent sparse attention methods, such as Reformer (Ki- taev et al., 2020) and Routing Transformer (Roy et al., 2021), aimi to approximate softmax in return for efficiency, not leveraging the sparsity of attention weights. Other methods, such as Top-k attention (Gupta et al., 2021) and NSA (Yuan et al., 2025), provide sparsity but require a fixed, non- adaptable budget. In contrast, a-entmax attention provides natural, input-dependent sparsity patterns with an exact and differentiable transformation that generalizes softmax, mak- ing it more flexible for modeling attention distributions. Adaptively sparse transformers (Correia al., 2019) uses a-entmax attention where attention heads can learn dy- namically, improving interpretability but without leveraging sparsity for efficiency. SparseFinder (Treviso et al., 2022) aims to address efficiency issues by predicting the sparsity pattern of entmax attention a priori; however, it does not scale efficiently for long sequences.
Sparsity in attention heads. Figure 4 presents the spar- sity observed in attention heads for all layers for an input of 1024 tokens for our sparse GPT-2 model (a 1.5). Except for the first layer, all subsequent layers exhibit a high degree of sparsity, highlighting the potential efficiency gains from leveraging this property. Moreover, in Fig- ure (Appendix B.1), we illustrate the sparsity patterns in ModernBERT-bau base attention heads for a 2.0). reinforcing similar conclusions.
5. Related Works
Sparse Probability Transformations. The sparsity in- herent to the a-entmax transformation, as demonstrated by Blondel et al. (2019), is directly controlled by the a pa-
8

---

ADASPLASH: Adaptive Sparse Flash Attention
Hardware-Aware Attention. Recent works have ex- plored optimizing attention mechanisms with hardware- aware implementations. Flex Attention (Dong et al., 2024) provides an API for efficient attention computation, though they remain tied to softmax-based transformations and do not support more complex operations such as those con- sidered in our work. Closely related to our approach, FlashAttention-11 and 2 (Dao et al., 2022; Dao, 2024) op- timize softmax-based attention using tiling and recompu- tation techniques implemented in CUDA. While FlashAt- tention includes a sparse block variant, its sparsity pattern must be predefined, limiting adaptability. In this work, we compare our method, ADASPLASH, with FlashAttention-2 and demonstrate that our approach can outperform both its CUDA and Triton implementations at high input sparsity levels. Similarly, Sparse Flash Attention (Pagliardini al., 2023) extends FlashAttentionn-1 with sparse variant that reduces computational cost by either dropping queries and keys per head or grouping them using hash-based bucket- ing approach. However, despite its efficiency improvements, it relies on slow sorting operations and is constrained to causal attention, making its sparsity by-product of bucket- ing rather than an inherently adaptive feature, as in our case.
Impact Statement
Efficient attention mechanisms are crucial for scaling trans- formers to long-context tasks. Our work provides a prac- tical implementation by making adaptive sparse attention efficient, overcoming previous computational limitations of o-entmax. Therefore, the improved efficiency of ADAS- PLASH has potential applications in large-scale NLP, where sparsity can be leveraged to reduce computational costs. We do not foresee direct societal consequences from sparsity itself, but its integration into decision-making models may still reflect biases in training data. As such, we encourage careful evaluation when deploying sparse attention mecha- nisms in high stakes applications, ensuring that efficiency gains do not come at the cost of fairness or transparency.
Acknowledgments
We thank Vlad Niculae for his insightful and construc- tive comments throughout this work. We also thank the SARDINE Lab members for reviewing this paper and pro- viding helpful feedback. This work was supported by the Portuguese Recovery and Resilience Plan through project C645008882-00000005 (Center for Responsible by the EU's Horizon Europe Research and Innovation Actions (UT- TER, contract 101070631), by the project DECOLLAGE (ERC-2022-Cod 101088763), and by FCT/MECI through national funds and when applicable co-funded EU funds under UID/50008: Instituto de Telecomunicações.
Efficiency at Inference Time. Another line of work fo- cuses on optimizing transformers at inference time. Meth- ods such as Paged Attention (Kwon et al., 2023) and KV cache sparsification (Devoto et al., 2024; Luohe et al., 2024) aim to alleviate the linear complexity of inference by mod- ifying key-value caching strategies. While our approach does not directly provide KV cache compression benefits, these methods are orthogonal and can be combined with our work to further improve inference efficiency.
References
Ansel, J., Yang, E., He, H., Gimelshein, N., Jain, A., Voz- nesensky, M., Bao, B., Bell, P., Berard, D., Burovski, E., Chauhan, G., Chourdia, A., Constable, W., Desmaison,
6. Conclusion
A., Vito, Z., Ellison, E.,, Feng, W., Gong. J., Gschwind, M., Hirsh, B., Huang, S., Kalambarkar, K., Kirsch, L., Lazos, M., Lezcano, M., Liang, Y., Liang, J., Lu, Y.. Luk, C. K., Maher, B., Pan, Y., Puhrsch, C., Reso, M., Saroufim, M., Siraichi, M. Y., Suk, H., Zhang, S., Suo, M., Tillet, P., Zhao, X., Wang, E., Zhou, K., Zou, R., Wang, X., Mathews, A., Wen, W., Chanan, G., Wu, P., and Chintala, S. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Vol- ume 2, ASPLOS '24, pp. 929-947, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703850. doi: 10.114553620665.36410
In this work, we introduced ADASPLASH, a hardware-aware and efficient implementation of a-entmax attention, bridg- ing the gap between adaptive sparse activations and efficient long-context modeling. Our approach leverages a hybrid Halley- billee--bseetion algorithm for faster empirical convergence and custom Triton kernels to exploit the inherent sparsity of a-entmax. Our experiments show that ADASPLASH not only achieves substantial computational improvements over existing a-entmax implementations, but can often match or even surpass the efficiency of highly optimized softmax- based attention algorithms like FlashAttention-2. Moreover, ADASPLASH enables long-context training while maintain- ing strong task performance across diverse benchmarks, such as language understanding, information retrieval, doc- ument classification, and language modeling. Overall, our work unlocks the viability of dynamically sparse attention mechanisms in large-scale training, which was previously hindered by computational inefficiencies.
URL.
https://doi. org/ 10.
Bajaj, P., Campos, D., Craswell, N., Deng, L., Gao, J., Liu, X., Majumder, R., McNamara, A., Mitra, B., Nguyen. T., et al. Ms marco: A human generated machine reading
9

---

ADASPLASH: Adaptive Sparse Flash Attention
comprehension dataset. arXiv preprint arXiv:1611.09222 2016.
In International Conference on Learning Representa- tions, 2021. URL httppp://opearrreee.eeeeeeeeeeee id=Ua6zukOWRH.
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The Long-Document Transformer. arXiv:2004.05150 [cs], April 2020. URL http://arxiv. org 05150. arxiv: 2004.05150.
Condat, L. Fast projection onto the simplex and the & ball. Math. Program., (58(1-2):575
July 2016. ISSN
0025-5610. doi: 10. 10:1007/400000--05544446 https://doi. http:///oooorrg/////// 10. 1007 w0t.100000/0100000000000090
URL
Blondel, M., Martins, A., and Niculae, V. Learning classi- fiers with fenchel-ye yeuchel-young young losses: Generalized entropies, margins, and algorithms. In Chaudhuri. K. and Sugiyama, M. (eds.), Proceedings of the Twenty-Second Interna- tional Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 606-615. PMLR, 16-18 Apr 2019. URL https:// proceedings.co mlr ptwps/////boooeellee... www.dellneeeeee html.
Correia, G. M., Niculae, V., and Martins, A. F. T. Adap- tively sparse transformers. In Inui, K., Jiang. J., Ng, V., and Wan, X. (eds.), Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 2174-2184 Hong Kong, China, November 2019. Association for Computational Linguistics. doi: (0.18653/1//11--122 1223. URL httpps://cclleeee htwps//thllaoc....o...... Bttp:/ooooooooo/1---11101 org/D19-122///
Chalkidis, I., Androutsopoulos, I., and Aletras, N. Neural legal judgment prediction in English. Korhonen, A., Traum, D., and Mårquez, L. (eds.). Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pp. 4317-4323, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 121166///1/11--1222 URL erg/P19-10 1124//
Dai, X., Chalkidis, I., Darkner, S., and Elliott, D. Revis- iting transformer-based models for long document clas- sification. In Goldberg, Y., Kozareva, Z., and Zhang, Y. (eds.), Findings of the Association for Computational Linguisties: EMNLP 2022, pp. 7212-7230, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 10.18653/11/022. mudingg-eemnlp.. URL https: https:/ httpp:///aclanthology org/ innding-eemml..5555 534/.
Chalkidis, I., Fergadiotis, M., Tsarapatsanis, D., Aletras, N., Androutsopoulos, I., and Malakasiotis, P. Paragraph-level rationale extraction through regularization: A case study on European court of human rights cases. In Toutanova, K., Rumshisky, A., Zettlemoyer, L., Hakkani-Tuu Tur, D., Beltagy, I., Bethard, S., Cotterell, R., Chakraborty, T., and Zhou, Y. (eds.), Proceedings of the 2021 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, pp. 221,241, Online, June 2021. Association for Computational Linguistics. doi: 10. (0.18853/v1/02) 1/2021. naacl- mainl-m URL https://aclanthologg/ org/ 2021 nwa.nnacllmma.m 22/.
Dao, T. Flash Faster attention with better paral- lelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAt- tention: Fast and memory-efficien exact attention with I0-awareness. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
Devoto, A., Zhao, Y., Scardapane, S., and Minervini, P. A simple and effective 1_2 norm-based strategy for KV cache compression. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Process- ing, pp. 18476-18499, Miami, Florida, USA, Novem- ber 2024. Association for Computational Linguistics. doi: 10. 00.18653/v1/2024 http:////22222..mmmaaaaaama www.p-mainnnn....m. 1027. URL https: //aclanthology ogy....gg..... org/2024. emnle- main. 1027/.
Chen, B., Dao, T., Winsor, E., Song, Z., Rudra, A., and Ré, C. Scatterbrain: Unifying sparse and low-rank attention. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.). Advances in Neural Information Process- ing Systems, 2021. URL https://openrew// /tt.//eeerrreeeeeeeeeeeee forum?id=Seh
Child, R. Gray, S., Radford, A., and Sutskever, I. Generating Long Sequences with Sparse Transformers. 0rXiv:1904.105550 10509 [cs, stat]. April 2019. URL http: 10509. arXiv: 1904.
Dong, J. Feng, B., Guessous, D., Liang, Y., and He, H. Flex attention: A programming model for gener- ating optimized attention kernels, 2024. URL. https: //arxiv. www.xiv.rrg/bb////// otg//bss/412..555
//arxiv. faaxii.org//a
version: 1.
Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mo- hiuddin, A., Kaiser, L., Belanger, D. B., Colwell, L. J., and Weller, A. Rethinking attention with performers.
Duchi, J., Shalev-Shwartz, S., Singer, Y., and Chandra, T. Efficient projections onto the 11-ball for learning in high dimensions. In Proceedings of the 25th Inter- national Conference on Machine Learning, ICML '08,
10

---

ADASPLASH: Adaptive Sparse Flash Attention
pp. 272-279, New York, NY, USA, 2008. Association for Computing Machinery. ISBN 9781605582054. doi: 10. 10.1145/1390156 1390191. URL httpss//doi.oon 10. 1141/13991166666 1390191.
Luohe, S., Zhang, H., Yao, Y., Li, Z., and hai zhao. Keep the cost down: A review methods optimize LLM's KV-cache consumption. In First Conference on Lan- guage Modeling, 2024. URL https://openreet det/forumiid=8kke jog 4MM...
Gupta, A., Dar, G., Goodman, S., Ciprut, D., and Be- rant, J. Memory-efficient transformers via top-k atten- tion. In Moosavi, N. S., Gurevych, I., Fan, A., Wolf, T., Hou, Y., Marasović, A., and Ravi, S. (eds.), Pro- ceedings of the Second Workshop on Simple and Effi- cient Natural Language Processing, pp. 39-52, Virtual, November 2021. Association for Computational Linguis- URL https:
Martins, A. and Astudillo, R. From softmax to sparsemax: A sparse model attention and multi-label classifica- tion. In Balcan, M. F. and Weinberger, K. Q. (eds.), International Conference on Machine Learning (ICML), volume 48 of Proceedings of Machine Learning Research, pp. 1614-1623, New York, New York, USA, 20-22 Jun 2016. PMLR. URL attp://proceedine mlr. press/v48/martinngg html.
tics. doi: 10. 112.18663/112022.sssaaa //aclanthology... org/ 3rgg/221..ustaannpp11.55
Held, M., Wolfe, P., and Crowder, H. P. Validation of subgradient optimization. Mathematical Programming, (1):62-88, December 1974.
Milakov, M. and Gimelshein, N. Online normalizer cal- culation for softmax, arXiv preprint arXiv: 1805.020522 0286677 2018.
Pagliardini, M., Paliotta, D., Jaggi, M., and Fleuret, F. Fast attention over long sequences with dynamic sparse flash attention. In shirth-seventh Conference Neu- ral Information Processing Systems, 2023. URL https: // wwwpeeewwwweeeetoou/nnniiiiie.
Jiang, H., Li, Y., Zhang, C., Wu, Q. Luo, X., Ahn, S., Han, Z., Abdi, A. H., Li, D., Lin, C.-Y., Yang, Y., and Qiu, L. Minference 1.0: Accelerating pre-filling for long- context lims via dynamic sparse attention. arXiv preprint arXiv:2407.0m 02440, 2024.
Penedo, G., Kydlíšek, H., allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L. V., and Wolf, T. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eighl Conference on Neural Infor- mation Processing Systems Datasets and Benchmarks Track, 2024. URL htth:: https://openrreview wttpeeeviieeeee/frurnnn net/forum? i id=n6SCkn2QaG.
Kaufman, E. H. and Lenker, T. D. Linear convergence and the bisection algorithm. The American Mathematical Monthly, 93(1):48-51, 1986. ISSN 00029890, 19300972. URL http: htt.p:///ww...soorrrroo jstor. org/ httg/stalee/2322555
Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The effi- cient transformer. In International Conference on Learn- ing Representations, 2020. URL httptttt////peeeeiies net/ aett/toruumm.Grrrggl/aatttt
Peng, H., Pappas, N., Yogatama, D., Schwartz, R., Smith, N., and Kong, L.. Random feature attention. In In- ternational Conference on Learning Representations, 2021. URL
Kwon, W. Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Sym- posium on Operating Systems Principles, pp. 611-626, 2023.
QtTKTdVIFBB.
Peters, B., Niculae, V., and Martins, A. T. Sparse sequence-to-scquencc models. In Proceedings of the 57th Anmual Meeting of the Association for Computa- tional Linguistics, pp. 1504-1519, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.18653/w1/119-1144. URL https: http://www.llllw... org/ org/anthoology/19-111
Liu, J. and Ye, J. Efficient euclidean projections linear time. Proceedings of the 26th annual international conference on machine learning, pp. 657-664, 2009.
Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. 2019.
Roberta: A robustly optimized bert pretraining approach, 2019. URL hattp:://are //arxiv.../////// hts/abb//1199 19007 11692.
Roy, A., Saffar, M., Vaswani, A., and Grangier, D. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Lin- guistics, 985 2021. doi: 1111622tacLa_00335 URL httpe://aclanthologg org/ tacl-1.4
Lozhkov A., Ben Allal, L., von Werra, L., and Wolf, T. Fineweb-edu: the finest collection of educational content, 2024. URL
htwp:/giinnnaccc.ccctttateeeeeee co/datasets/ Hugg:ing(aceF///Eneeeeeeeed

---

ADASPLASH: Adaptive Sparse Flash Attention
Scavo, T. R. and Thoo, J. B. On the geometry of halley's method. The American Mathematical Monthly, 102(5): 417-426, 1995. ISSN 00029890, 19300972. URL http: //www.jston wwtg wwrgsstable/2975055
Florence, Italy, July 2019. Association for Computa- tional Linguistics. doi: 10. 10.186553/1//119-5550 URL https://acllanthollog 11tp://1000.....rr//111110010
Wang. A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Linzen, T., Chrupala, G. and Alishahi, A. (eds.), Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyz- ing and Interpreting Neural Networks for NLP, pp. 353- 355, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.18653/v1/W1885555 URL https://acllatthooog org/ 5446.
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate attention with asynchrony and low-precing oowpprccsioon, 2024. URL https: / /arxiv..vv.. werxiv.org//eie owg//b///////// abs/abb//7.... 08608.
Thakur, N., Reimers, N., Rückléé A., Srivastava, A., and Gurevych, I. Beir: A heterogeneous benchmark for zero-shot evaluation of information retrieval models. In Vanschoren, J. and Yeung, S. (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks, volume 2021. URL wcwppppecccppppprr.ereeee http:/pppppeetlleeappppp//////// abseaadeelcbbbffroc224777aaa.aa. Paperarr----PPPerrrrrrrrr roun2.. pdf.
Warner, B., Chaffin, A., Clavié, B., Weller, O., Hallström, O., Taghadouini, S., Gallagher, A., Biswas, R., Ladhak, F., Aarsen, T., et al. Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference. arXiv preprin! arXiv:2412. 13663, 2024.
Tillet, P., Kung, H. T., and Cox, D. Triton: an intermediate
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Effi- cient streaming language models with attention sinks. In The Twelfth International Conference on Learning Rep- resentations, 2024. URL https: htttps://openreveer: .et// forumidd=N77s51122
language and compiler for tiled neural network compu- tations. In Proceedings of the 3rd ACM SIGPLAN Inter- national Workshop on Machine Learning and Program- ming Languages, MAPL. 2019. pp. 10-19, New York, NY, USA, 2019, Association for Computing Machinery. ISBN 9781450367196 doi: 10. (0.11333311558833999999 https://doi. org/ 10. 3399999..8888889.73.. 3329973.
URL
Xiong, Y., Zeng. Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., and Singh, V. Nyströmformer: A nyström-based algorithm for approximating self alffaattenioon 2021.
Treviso, M., Góis, A., Fernandes, P., Fonseca, E., and Mar- tins, A. Predicting attention sparsity in transformers. In Proceedings of the Sixth Workshop Structured Predic- tion for NLP, pp. 67-81, Dublin, Ireland, May 2022. Asso- ciation for Computational Linguistics. doi: 10. 18653/v1/ 2022.spnlp-1.7. URL https: http///ccaaaatllgo/o .org/ 2022. 8022--p1
Yuan, J., Gao, H., Dai, D., Luo, J., Zhao, L., Zhang, Z., Xie, Z.,, Wei, Y., Wang, I.., Xiao, Z., et al. Native sparse attention: hurdware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502. 1089, 2025.
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q. Yang, L.. et al. Big Bird: Transformers for Longer Sequences. Advances in Neural Information Processing Systems, 33:17283-17297, 2020a. URL
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, Attention all you need. Advances in neural information processing systems, 30, 2017. URL http:///paapreenipp..ccpcapppppp 2tbe22455476e91110001114446844 Abstract. html.
https://papers..e httpss//ppprrs..ap..cccpppper https//ppcccppppp////////aaaaaaa c85 28811a11222884477255118887888810 -tt.:///aa0000bbabaaaba.t. Abstract html.
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Al- berti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q. Yang, L., et al. Big bird: Transformers for longer se- quences. Advances in Neural Information Processing Systems, 33, 2020b.
Veličković, P., Perivolaropoulos, C., Barbero, F., and Pas- canu, R. softmax is not enough (for sharp out-of- distribution), 2025. URL forum?id==t1jfgNNWu.
Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Ko- rhonen, A., Traum, D., and Mårquez, L. (eds.), Pro- ceedings of the 57th Annual Meeting of the Associ- ation for Computational Linguistics, pp. 5797-5808,
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sen- tence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
12

---

ADASPLASH: Adaptive Sparse Flash Attention
A. Algorithm Details
first derive high-level view of the forward and backward passes of the entmax attention and then present the full algorithms for both mentioned versions. For consistency and ease of comparison, follow the notation adopted by FlashAttention-I (Dao al., 2022).
A.1. o-entmax Attention Forward Pass
We recall that given the input sequences Q, K, V E Raxd we want to compute the attention output E Rnxd follows:
S = QK E R P (e-entmax E Rax O = PV E Rxadd
Therefore all we need is the E Rh that solves Equation 2, for which we can use Algorithm 1. We note that we do not need to materialize S as we only need to accumulate the derivatives of f(T), defined in Equation 3. Once T is computed, we can compute each row of O as follows:
- P Po - Σ max (a
[1QKK K, T.) vo) V
(15)
As in FlashAttention, we can compute O; without extra memory by incrementally summing the contributions of each a-entmaxx(Q) (Q, Vj term. We can then compute the forward pass with (n) extra memory as follows:
1. Compute for all 1 S i according to Algorithm 1, which takes O (n) extra memory.
2. Compute O for all < i S nu according to Equation 15 which takes O (n) extra memory.
A.2. a-entmax Attention Backward Pass
For the a-entmax attention backward pass, we need to compute the gradients with respect V. K, and Q. Let L be scalar loss function, and dO E nx denote 0C Our goal is compute the input gradients dV, dK, dQ E gaxd
Gradient of V
Using reverse-mode autodifferrntiiatio,, first compute dV:
dV = do,
(16)
where P = o--enmaxxSS) is the output of the a-entmax transformation applied row-wise to the score matrix $ = =QKK
Expressed element-wise, we obtain:
7
dV,
P. do;
(17)
which is analogous the softmax Since P.j is sparse due to the nature of a-entmax, skip Q: blocks that leads to blocks of P full of zeros using the pointer increment tables, as shown in Equation 14.
2. Gradient of P and S
The next step involves computing dPP and dSS. From O PV, we have:
dP = do, V,.
(18)
Next, let us recall the Jacobian of the a-entmax mapping (Peters et al., 2019). Defining a-entmax(s), the Jacobian is:
0a-entmax(s)
uu
Diag(u)
(19)
Øs
|u|
13

---

ADASPLASH: Adaptive Sparse Flash Attention
where u is defined element-wise as:
(px)2 0,
ps > 0
2kk
(20)
otherwise.
Let U denote stack of [u,,..., un) for each row of P. From the relationship P a-entmax((S), (S), and the Jacobian of the d-entmax function, we can propagate the gradients back to follows:
U.U.
ds,
Di,)(U,)
dP.
(21)
|
U dP:
U,O o dP.
U:
(22)
|
We can further simplify by defining a new quantity E R::
U dP.
6,
(23)
||Ulll
"
Σ
Un (do, v²))
(24)
|
(Σg-n UNj V3
d
(25)
||
In standard softmax attention, instead of the right-side term in the above product, we would simply obtain O, Since this new quantity is required for the backward pass, and avoid passing more through Q, K and V. we compute and store this quantity during the forward pass solely during training. Unlike in softmax attention, however, the backward pass for o-entmax does not require saving the output matrix O; instead, we only require this new quantity, which we label (2) Then, we can simplify the computation of dS to:
ds, U, @ (dP: -,)
(26)
Again, we can use the sparsity stored in M (see Equation 9) from the forward pass to efficiently skip the computation of null blocks of P.
3. Gradients of Q and K
Using the definition of Sij = Q, Kj, the gradients for Q and K are:
dQ
ds K₁,
(27)
f=1
dK, Σ ds,
(28)
Substituting dSij, we get:
dQ, Σ
Up (dP.j 6,) K₁
(29)
1=1
n
dK, Σ
Ug (dP,, - 6,)
(30)
14

---

ADASPLASH: Adaptive Sparse Flash Attention
Effectively, we only iterate through the blocks that will result in Psy + 0. As in FlashAttention, the backward pass can also be computed with O (n) extra memory:
1. Compute dV, for all j according to Equation 17, which takes O (d) extra memory.
2. Compute d; for all i according to Equation 23, which takes O (1) extra memory.
3. Compute 0 for all i, as defined in Equation 25, which takes O (d) extra memory.
4. Compute dQ, for all according Equation 29, which takes O (d) extra memory.
5. Compute , for all j according to Equation 30, which takes O (d) extra memory.
We note that the only extra memory requirement compared to FlashAttention is in having to additionally compute and storing ( E Raxd_ When using block masking, we also need O (T, x Te) extra memory to store the binary mask M. However, we recall that this memory can be shared across attention layers, as is merely a temporary matrix used to compute the pointer-increment tables.
A.3. ADASPLASH: Forward Pass (without block masking)
The full ADASPLASH's forward pass is presented in Algorithm 2. For completeness, we also provide Algorithm 3 the steps for approximating T without the need to materialize in block-wise manner.
Algorithm 3 Halley- bisection Halley-bisection
for computing T Block Version
Require: Matrices Q, K E R in HBM, block sizes Ba, B, and number of iterations M.
1: Divide Q into T, = B, blocks Q...... Qr. of size B, x d B, 2: Divide K into T. = [n/ Be] blocks KI...,, KT. of size Bg x
3: Divide т into T, blocks Ti...., Tr, size B,
4: for = 1 to T, do
5: Load Q, from HBM on-chip SRAM
6: On chip, initialize T;, Taa Thi, according to Algorithm 1.
b. Note: this requires pass over Kj for all
7: repeat
8: On chip, initialize f.f, - B..
9:
j = 1 to Ta do
10: Load K, V, from HBM to on-ehip SRAM
Compute 5 K B.BB.
11:
12:
Accumulate f. f. ff according to Equations 3, 6 and 7, respectively.
13:
end for
14:
Update Ts, To, Thi, according to Algorithm
until M iterations are completed
16:
Write Ty to HBM
17:
end for
18:
Return:
A.4. ADASPLASH: Backward Pass (without block masking)
As mentioned in $3.2.2, in contrast to FlashAttention, we propose to separate the kernels that compute the gradients dQ, d., dV. However, as in FlashAttention, we need to compute before being able to compute the gradients, which we do in a separate kernel following Equation 25. We present the full steps for computing dKK and dV in Algorithm 4, and for computing dQ in Algorithm 5.
A.5. ADASPLASH: Block Masked Version
In this version. as outlined in Section 3, a boolean block mask M E T- is created dynamically in the forward pass, allowing the exploitation of the sparsity in the matrix P the cost of linear memory complexity. The mask populated
15

---

ADASPLASH: Adaptive Sparse Flash Attention
Algorithm 4 ADASPLASH Backward Pass for d.K and dV
Require: Matrices Q, K, V, O, dO E Raxd in HBM, vector E R" in HBM, block sizes Be, Br, parameter 1: Divide Q into T, = [n/B,] blocks Q1....,Quu of size B, x d cach, and divide K, V into Te [n/ B. blocks VT, of size Be x deach.
K,,...,yy,
2: Divide dO into T, blocks dO1,... dO2, of size B, x d each.
3: Divide into T, blocks Ti,.... Tr. of size B, each.
4: Initialize and divide dK, dV E Raxd into Te blocks dKl,.... dKr. and dV dVi..... dVr. of size Be x each.
5: Divide 5 into T, blocks 61..... δr, of size B, each.
6: for S j S Te do
7: Load Kj, Vj from HBM to on-chip SRAM.
8: Initialize dK, = OB on SRAM.
9: Initialize dV, = -B,xd on SRAM.
10: for 1 / < T, do
11: Load Q, dO;, Ti 6, from HBM to on-chip SRAM.
12:
B.B. On chip, compute ) = K.Kj E
On chip, compute p(1) = max(0, (a - 1)5(() -T;; -/ E RBB.B.
13:
On chip, compute dV; dV, + (p()) T d0; E BBdd
14:
15:
On chip, compute dP,- dO,V," BBB....
On chip, compute (() p.())= ERB....
16:
On chip, compute ds( U,6) (dpli) -.)) E B.
On chip, compute dK; + dK; + (ds) )TQ. E ERa.
18:
19:
end for
20: Write dK3,, dV, to HBM.
21:
end for
Return: Gradients dK, dV.
during the final iteration of the Halley-bisectin algorithm (Algorithm 3) by evaluating the condition any(s > T,) and storing the result as a boolean value. Thus, the mask M indicates whether a specific Q, K block pair contributes to the output. This process enables the creation of a lookup table that associates each query block with the set of key blocks that contribute non-zero values, thereby allowing to skip unnecessary computations for future computations. Similarly, a reverse lookup table can be created for each key block. Both tables can be used the backward pass (Line 10 in Algorithm 4 and Line 9 in Algorithm 5) to avoid looping over unnecessary query/key blocks.
In practice, to create the lookup tables, we use the torch. arguhere function to extract the (i,j) indices of entries where My = 1. Combined with row-wise summation of non-zero entries, this approach efficiently skips computations for irrelevant blocks within the remaining kernels. Consequently, during the forward pass, only the K, V pairs identified in the lookup table are loaded, avoiding redundant memory and computational overhead. As mentioned, for the backward pass, given that we separated the computation of dQ and dK, dV, we can further use both tables (Q and XC) speedup the gradient computation.
B. Experimental Setup
B.1. Continuous Pre-training
We conducted continuous pretraining of RoBERTa-basee and ModernBERT-baas models with our custom sparse attention Triton kernel, ADASPLASH. The pretraining process was carried on 2B tokens of the Fine FineWeb-Edu dataset, due to its high-quality, diverse and large-scale content. We used the HuggingFace Transformers library for model training and implementation and the Datasets library for data handling. Concretely, we used batch size of 32 and a learning rate of 5 x 10-5, optimized with the AdamW optimizer. Training was conducted for 100,000 steps using mixed mixed-precision (fp16).
16

---

ADASPLASH: Adaptive Sparse Flash Attention
Algorithm 5 ADASPLASH Backward Pass for dQ
Require: Matrices Q, K, V, O, dO E Roxd in HBM, vector E R" in HBM, block sizes Be, Br, parameter a. 1: Divide Q into T, = [n/B,] blocks Q1....,Q1 of size B, x d cach, and divide K, V into Te [n/BL B. blocks K VI...., .T. of size B. x deach.
2: Divide dO into T, blocks 1O1,.... dO2, of size B, x each.
3: Divide into T, blocks TI...., Tr. of size B, each.
4: Initialize dQ in HBM and divide into T, blocks
,QQ, of size B, d each.
5: Divide 6 into T, blocks ó1..... 8z, of size B, each.
6: for = 1 to T. do
7: Load Q dO;, 6;, T.. from HBM to on-chip SRAM
8: Initialize dQ; OB.xd on SRAM.
9: for j = to T. do
10: On chip, compute ( Q E IR B.
)
On chip, compute p(j) max (a 1)5( - ERB.BB.
11:
12:
On chip, compute dP. - dO,V, E RBB.B..
On chip, compute U(I) p(]) B.B..
) On chip, compute ds (1) U.() @ (dP(1) -d)) EBBxB.
14:
15: On chip, compute dQ, t dQ; + +d K E BBxd
end for
Write dQ: to HBM
18:
end for
19: Return: Gradient dQ
Table 5. Runtime (s) of ModernBERT-baan (a L.5) for varying context lengths.
Sequence Length
Algorithm
512
1024
2048
4096
8192
Sorting (Torch)
0.09
0.11
0.26
0.76
OOM
Bisection (Torch)
0.11
0.15
0.42
1,35
4.99
fallly-bisecctio (Triton)
0.10
0.11
0.26
0.46
1.61
ADASPLASH (Triton)
0.10
0.12
0.21
0.48
1.53
The sparsity parameter (a) was initialized at 1.01 and annealed linearly to final value of 1.5 or 2.0 over 50,000 steps. We kept ModemBBERT's window attention layers untouched, only replacing the full softmax layers by a-entmax. Finally, we also performed continuous pretraining of RoBERTa and ModemBBERT with standard softmax attention with a fixed a = 1.0.
As shown in Figure 5, the attention mechanisms our sparse ModemBERR model 1.5) obtain high sparsity levels in practice, with overall sparsity of 95% for 1.5 and 99% for a 2.0. For this reason, used the version of ADASPLASH that leverages the pointer increment tables for training ModernBERT, which has maximum sequence length of 8, 192. For RoBERTa, which has a sequence length of 512, we opted to use the Halley-bisectinn algorithm implemented in Triton. In Table 5 we report efficiency results in terms of runtime and memory usage for different attention algorithms with ModernBERT-bare Overall, we observe that the sorting approach slower than bisection, which slower than our Hallcy-bisectiin and ADASPLASH, in that order.
B.2. GLUE and BIER tasks
For GLUE tasks, we used the checkpoints continuous pre-trained models for both RoBERTa-base and ModernBERT Then, we fine-tuned them on each GLUE task with the default hyperparameters from the Transformer library. 6 Importantly, we capped the maximum sequence length at 128 tokens to reduce computational cost while preserving task-relevant context and used fp16 for training.
17

---

ADASPLASH: Adaptive Sparse Flash Attention
1.0
1.0
0.8
0.8
0.6
0.6
0.2
30
10
-0.0
0.0
12 15
18
21
6
12
15
18
21
Figure 5. Ratio of non-zeros for non-local layers of ModemBERT-baar with 1.5 (left) and 2.0 (right).
Table Results on different tasks from GLUE benchmark (Wang al., 2018).
Single Sentence
Paraphrase and Similarity
Natural Language Inference
Model
Params
Scq.
CoLA
SST-2
MRPC
STS-B
QQP
MNLI
QNLI
RTE
Avg.
BERT
110M
512
58.6
91.9
86.9
89.0
89.3
84.0
91.0
69.3
82.5
RoBERTa
125M
512
59.8
93.7
89.5
89.6
89.8
87.7
92.3
69.3
83.9
RoBERTa (a 1.5)
125M
512
58.5
93.2
91.5
90.2
89.7
87.3
92.5
68.6
83.9
RoBERTa 2.0)
125M
512
56.8
93.0
90.9
88.8
89.0
86.7
91.9
67.2
83,0
ModemBERT
149M
8192
63.2
95.0
88.2
90.3
90.4
87.9
93.0
61.7
83.7
ModemBERR (a 1.5)
149M
8192
62.2
96.1
87.7
89.4
90.2
87.9
92.6
61.7
83.5
ModemBERT (a 2.0)
149M
8192
62.2
94.8
89.0
89.9
90.5
87.8
93.1
62.5
83.7
To evaluate the generalization of ADASPLASH in retrieval tasks, we fine-tuned ModernBFRRT-haee and RoBERTa-base models on the MS MARCO dataset (Bajaj et al., 2016) and evaluated them on the BEIR benchmark (Thakur et al., 2021). This benchmark suite assesses performance across diverse information retrieval tasks, including SciFact, NFCorpus, FiQA- 2018, and TREC-COVID. The fine-tuning and evaluation process closely follows the approach proposed in the BRR paper (Warner al., 2024). Fine-tuning was performed using the CentenceTransfommeen library." The models were evaluated on BEIR tasks using the MTEB benchmark toolkit." The evaluation metric for each task nDCG @ 10 (Normalized Discounted Cumulative Gain), following standard information retrieval practices.
B.3. Long Document Classification
The European Court of Human Rights (ECIHR) dataset comprises legal cases from the European Court of Human Rights, each associated with specific articles of the Convention Human Rights allegedly violated. For this task, fine-tuned the RoBERTa base model (Liu et al., 2019) with classification head. Since this is multi-label classification task, we used binary cross-entropy loss. To accommodate longer contexts, we followed the approach proposed by (Beltagy et al., 2020), repeating the 512 position embeddings until the target context size was reached. We used the Adam optimizer for training. For hyperparameters, we follow the recipe Dai et al. (2022). For the attention mechanism, bfloat16 precision was used.
B.4. Language Modeling
We trained both the standard GPT-2 model and sparse GPT-2 (a = 1.5) using the configuration provided in the 11m. repository," 9 following their training recipe. Specifically, we trained a GPT-2 (124M parameters) from scratch on 10B tokens of the Web dataset, with maximum sequence length of 1024 tokens. Training was conducted using bfloat16 precision. We use an effective batch size of 512, and use gradient accumulation fit into available GPU memory. We
"nttps://sbert.nee "nttps://github. http://hhhh...oeeebbediiiaaeaaaa "nttpe://github hacper//gihhbb.co///rrpttty www/karpathy/lln......
18

---

ADASPLASH: Adaptive Sparse Flash Attention
10 and weight decay of 0.1. The learning rate followed warm-up phase, linearly ramping from zero maximum of 6 10-1 over the first 700 iterations, equivalent 350 million tokens. Subsequently, the learning rate decayed to zero across the remaining training steps. We show the validation loss curves for both softmax and a-entmax (a 1.5) in Figure 6.
the AdamW optimizer, with learning
Given that, for this task, the context size was not high enough, for sparse attention we opted to use the algorithm that does not take advantage the pointer increment tables. For the benchmarking of the time spent per step, we averaged across 50 steps after the model had trained for at least 100 steps.
GPT-2
4.4
Sparse GPT-2
4.2
Validation Loss
3.6
3.4
3.2
2500
5000
7500
10000
12500
15000
17500
20000
Validation Step
Figure 6. FineWeb withheld validation loss comparison between GPT-2 and Sparse GPT-2 during training.
C. Computational Details
Experiments on masked language modeling, text classification, GLUE tasks and BIER tasks were carried on Nvidia RTX A6000 GPUs with 48GB VRAM. Experiments with GPT-2 and the efficiency benchmark in Figures and 3 were carried on single Nvidia H100 GPU (80GB). The runtime experiments with ModemBERT were carried single A6000 GPU.
19


Reviews:
[{"review_id": "5oWarXzgJp", "reviewer": "Reviewer_CoRN", "summary": "The paper introduces ADASPLASH, a novel method that combines the computational efficiency of GPU-optimized algorithms with the sparsity benefits of the α-entmax attention family. The goal is to enable efficient and scalable training of transformers for long-context tasks while maintaining or improving task performance.\n\n## update after rebuttal\n\nThe author's rebuttal addressed most of my concerns; I will keep my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments and suggestions. \n\n**“The experiments are primarily conducted on BERT and GPT-2 level models, without testing on larger-scale models such as 7B-scale LLMs. Additionally, evaluations on long-text benchmarks, such as NIAH and RULER, are missing...”**\n\nPlease note that our primary goal is to introduce AdaSplash, an efficient implementation of entmax attention [1], rather than to benchmark sparse vs. dense attention on long-text tasks or large-scale models. While we don’t include 7B-scale evaluations or long-text benchmarks, we address a clear gap: entmax lacked an efficient implementation. As a bonus, we show that AdaSplash can even outperform FlashAttention in wall-clock speed (Fig. 1 & 3). Like FlashAttention, which also omitted large-scale retrieval experiments, we believe AdaSplash is a fundamental tool to enable future research in these directions.\n\n\n**“The experiments report evaluation results under different alpha values but do not provide corresponding computational efficiency metrics, such as sparsity levels and computation time....”**\n\nWe **do** provide key runtime and memory metrics in Tables 3, 4, and 6 in the Appendix. In particular, Table 6 specifically shows runtime for a 8192-token context using RoBERTa. We'll integrate and highlight these results more clearly in the final version.\n\n\n**“What is the comparison of computational efficiency between softmax and α-entmax, including both theoretical time complexity and actual computation time?”**\n\nIn theory, α-entmax requires slightly more work than softmax in the forward pass – about $(2 + k)N$ vs. $2N$ – but we considerably reduce k to just 2 or 3 iterations using a hybrid Halley-Bisection method, greatly improving runtime compared to previous methods (Fig. 2). When used in attention mechanisms, the main efficiency gains of α-entmax come in: (1) avoiding the multiplication by $V$ in the forward pass and (2) skipping zero blocks in the backward pass, where α-entmax has sublinear complexity (see Eq. 12-14). Note that AdaSplash exploits both points for faster computation. \n\n\n**“Does Figure 1 include the full computation time, or only the time after applying the sparse mask?”**\n\nFigure 1 reports the **full** computation time of forward + backward passes, and thus, it already includes both the computation for finding $\\tau$ and producing the output.\n\n\n**“What is the relationship between the α parameter and sparsity? Does a larger α always result in greater sparsity? Is there a theoretical basis or empirical evidence supporting this relationship?”** \n\nIndeed, as discussed in line 92, there is a direct theoretical relationship between α and sparsity: for the same input vector, increasing α results in an equal or greater sparsity of the output attention distribution. This theoretical result, along with empirical evidence can be found in the original entmax paper [1, Table 3]. \n\n\n**“Additionally, does the block size setting in the Triton kernel affect the sparsity of the sparse mask? If so, what is the nature of this effect?”**\n\nIndeed, the block size affects the sparsity of the block mask. Since masking is done at the block level (a block is skipped only if all its elements are zero), larger blocks reduce the chance of being fully masked. Conversely, smaller blocks increase the likelihood of full sparsity and improve computational efficiency but might induce a slightly higher memory overhead to store the bit mask (see response to R2 for more details on the memory cost). \n\nThe table below illustrates how block-wise sparsity varies with different block sizes (row: $B_r$, column: $B_c$) for a fixed input of 32k context length:\n| | 16   | 32   | 64   | 128  |\n|--|--|--|--|--|\n| **16**  | 96.8 | 93.6 | 87.2 | 76.1 |\n| **32**  | 93.6 | 87.7 | 76.4 | 58.4 |\n| **64**  | 87.7 | 77.0 | 59.3 | 35.3 |\n| **128** | 78.1 | 61.0 | 37.4 | 14.2 |\n\nFinally, while block size impacts practical efficiency, it does not affect the mathematical sparsity produced by the α-entmax function, which depends solely on α and the inputs.\n\n\n\n\n**“Compared to traditional fixed sparsity patterns, such as selecting tokens with a stride, what are the main advantages of adaptive sparsity?”** \n\nAdaptive sparsity offers more flexibility than fixed patterns by dynamically selecting the most relevant tokens based on the input, avoiding the risk of overlooking important ones. Unlike fixed sparsity, which approximates softmax for efficiency, α-entmax is an exact, smooth transformation that naturally produces sparse outputs, allowing computational savings without sacrificing accuracy. Moreover, prior work (Treviso et al., 2021) has shown empirically that α-entmax outperforms fixed-pattern attention methods (e.g., BigBird, Longformer, Reformer) in tasks like machine translation and masked language modeling, especially under high sparsity conditions.\n\n\nReferences:\n\n[1] Peters, B., Niculae, V., & Martins, A. F. “Sparse sequence-to-sequence models.” ACL 2019.", "claims_and_evidence": "Yes, the authors first propose a hybrid Halley-bisection algorithm for faster computation of the α-entmax transformation, reducing the number of iterations by 7-fold compared to standard methods. Then, to utilize this sparse pattern, the authors propose adaptive sparsity-aware Triton kernels to exploit sparsity in both the forward and backward passes of attention computations.", "methods_and_evaluation_criteria": "The experiments are primarily conducted on BERT and GPT-2 level models, without testing on larger-scale models such as 7B-scale LLMs. Additionally, evaluations on long-text benchmarks, such as Needle in a Haystack and RULER, are missing. Therefore, it is unclear whether this sparse attention mechanism can effectively handle mainstream long-text tasks.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experiments report evaluation results under different alpha values but do not provide corresponding computational efficiency metrics, such as sparsity levels and computation time. For instance, in Table 1, the 8192-length setting, which is the longest in the paper, could effectively showcase the efficiency advantages of sparsity but lacks such analysis.", "relation_to_broader_scientific_literature": "The paper's key contributions build on prior work in sparse attention mechanisms and efficient transformer architectures:\n\n- Relation to α-entmax: It extends the α-entmax family by introducing a more efficient Halley-Bisection algorithm, significantly improving its computational feasibility for large-scale tasks.\n- Relation to Efficient Attention: It complements methods like Longformer and FlashAttention by leveraging adaptive sparsity, offering dynamic attention patterns instead of fixed sparsity structures.\n- Hardware-Aware Optimization: The introduction of sparsity-aware Triton kernels advances hardware-optimized transformer implementations, bridging theoretical improvements with practical GPU efficiency.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "- What is the comparison of computational efficiency between softmax and α-entmax, including both theoretical time complexity and actual computation time? Does Figure 1 include the full computation time, or only the time after applying the sparse mask?\n- What is the relationship between the α parameter and sparsity? Does a larger α always result in greater sparsity? Is there a theoretical basis or empirical evidence supporting this relationship? Additionally, does the block size setting in the Triton kernel affect the sparsity of the sparse mask? If so, what is the nature of this effect?\n- Compared to traditional fixed sparsity patterns, such as selecting tokens with a stride, what are the main advantages of adaptive sparsity? Are there any empirical comparisons to illustrate these advantages?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pkqQJfnMLl", "reviewer": "Reviewer_agna", "summary": "The paper introduces ADASPLASH, an efficient implementation of α-entmax attention that leverages sparsity to improve computational performance while maintaining model quality. The key contributions are:\n\n1. A hybrid Halley-bisection algorithm that reduces the number of iterations needed to compute α-entmax transformation\n2. Custom Triton kernels that efficiently handle adaptive sparsity in both forward and backward passes", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive comments. We address your points below. \n\n**“Can the proposed method reuse the memory of KV with zero attention scores? If not, can the method actually save memory?”**\n\nSimilarly to FlashAttention, AdaSplash is primarily optimized for efficient training. For inference settings, for example, other methods such as Flash Decoding (an extension of FlashAttention) are needed in order to explicitly optimize inference by splitting the KV-cache memory and at the end aggregating the results. Because our work exclusively focuses on the training scenario, we have not addressed KV-cache optimizations. Nevertheless, investigating the potential benefits of leveraging the sparsity from α-entmax to optimize KV-cache utilization is an orthogonal and promising direction for future work.\n\n\nIn the training regime, compared to naive/eager attention implementation, AdaSplash significantly reduces memory usage which can be seen by Figure 3 as well as Tables 3 and 6 where the eager method goes out-of-memory for bigger contexts. We do this by avoiding the explicit materialization of large intermediate matrices such as the attention score matrix $S$ and the attention probability matrix $P$. We show that, just like in FlashAttention, one can avoid storing these matrices which in turn both reduces the memory usage and speeds up the computation.\n\n\n**“The improvements are most significant at high sparsity levels (>90%), which may not be reached in all applications. The efficiency gains for decoder-only models (GPT-2) are more modest than for encoder models, suggesting varying utility across architectures.”**\n\nInterestingly, in the specific case of GPT-2, our analysis (Figure 4) reveals a clear pattern where the first attention layer prefers denser attention distributions, while subsequent layers become significantly sparse. This observation naturally suggests the possibility of using an hybrid approach, combining dense attention at earlier layers with AdaSplash attention at later layers in order to potentially obtain the best balance between computation efficiency and model performance. We plan to incorporate this analysis with GPT-2 in the final version.\n\n**“The approach still requires materializing the mask matrix M with O(Tr × Tc) complexity, which could become a bottleneck for extremely long sequences.”**\n\nYou are right to note the memory overhead involved in materializing the block mask. Thus, we provide two implementations of AdaSplash: one that explicitly stores a binary block-sparsity mask (adding some memory overhead), and another without explicit mask storage (matching the memory requirements of FlashAttention-2) in subsection 3.2 (Line 299). For our experiments, the block mask memory overhead was not a bottleneck (Table 6). Moreover, we stress tested our implementation up to 64K context length and did not observe any memory issues, mainly because the matrix $M$ stores single bytes – extremely memory efficient in practice.", "claims_and_evidence": "The primary claims are well-supported by evidence", "methods_and_evaluation_criteria": "The methods are appropriate and well-justified.", "theoretical_claims": "The mathematical derivations for both the Halley-bisection algorithm and the attention mechanism appear sound. The derivation of α-entmax's Jacobian and its sparse properties (Section 3.2.2) is correctly presented and consistent with prior work. The block-wise computation approach in Algorithm 1 and the memory complexity analysis in Sections 3.2.1 and 3.2.2 are mathematically sound.", "experimental_designs_or_analyses": "The experimental design is thorough and well-executed:\n\n1. Comprehensive comparison against multiple baselines (FlashAttention-2 CUDA, FlashAttention-2 Triton, Torch bisection)\n\n2. Testing across multiple sequence lengths (512 to 64k tokens)\n\n3. Evaluation on diverse tasks and architectures\n\n4. Appropriate ablation studies (masked vs. non-masked variants)", "relation_to_broader_scientific_literature": "1. Adaptive sparse attention mechanisms\n2. Efficient attention implementations", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths:\n\nThe approach is architecture-agnostic, demonstrating benefits on both encoder and decoder models.\n\nWeaknesses:\n\nThe improvements are most significant at high sparsity levels (>90%), which may not be reached in all applications.\nThe efficiency gains for decoder-only models (GPT-2) are more modest than for encoder models, suggesting varying utility across architectures.\nThe approach still requires materializing the mask matrix M with O(Tr × Tc) complexity, which could become a bottleneck for extremely long sequences.", "other_comments_or_suggestions": "None.", "questions_for_authors": "Can the proposed method reuse the memory of KV with zero attention scores? If not, can the method actually save memory?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q9WKAqZJQy", "reviewer": "Reviewer_bvB9", "summary": "This paper aims to improve the efficiency of alpha-entmax attention, which is a kind of adaptive sparse attention. Its contributions are two folds:\n\n1. Alpha-entmax attention requires a threshold to determine which tokens are masked. This paper proposes an algorithm that achieves faster estimation of this threshold.\n2. After the threshold is estimated, this paper follows the spirit of Flash Attention and develops a hardware-optimized implementation of alpha-entmax attention called Sparse Flash Attention. The key idea is to avoid materializing the quadratic similarity matrices.\n\nExperiments demonstrate significant speedups compared with previous implementations of alpha-entmax attention.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments. We are glad that you found our paper provides clear and convincing evidence about the computational improvements of our implementations, and that our proposed methods and evaluation criteria make sense. We address your main concerns below. \n\n\n**“They should also evaluate the running speed on real dataset. Moreover, the author should report the running time of Halley-bisection alone.”**\n\nRegarding the runtime of the Halley-bisection algorithm, please note that we **do** report these results in Section 3.1 under the _Efficiency Benchmark_ paragraph (Line 182).\n\nFor the AdaSplash implementation, we followed FlashAttention-2 and chose to use synthetic data to evaluate our kernels, which allows us to control sparsity levels and systematically explore performance across varying context lengths (see Figures 1 and 3). Note, however, that we **do** report runtime on real datasets in our end-to-end training benchmarks (Table 3, 4, and 6) comparing the overall speed of models using AdaSplash with those using FlashAttention-2. \n \n\n\n**“This paper improves alpha-entmax attention, which is a special case of sparse attention. The authors should discuss the pros and cons of alpha-entmax attention compared with other types of sparse attention, such as topk attention [1] and block sparse attention.”**\n\nThis is a good suggestion, which we will follow. As also noted in our response to Reviewer CoRN, methods like top‑k and block-sparse attention rely on fixed, input-agnostic sparsity patterns, which may miss important tokens. Top‑k requires setting k in advance and block-sparse methods depend on a pre-defined mask and cannot adapt to specific patterns for each query vector. In contrast, α-entmax offers adaptive sparsity, selecting relevant tokens based on the input without requiring hard thresholds or fixed mask patterns. Its sparsity emerges naturally from the transformation, which is **exact**, **differentiable**,  and it generalizes softmax (when $\\alpha \\to 1$) without any approximations.\n\nWe will further clarify these distinctions in the final version.\n\n\n\n**“When compared with Flash-attention, we would expect that the proposed method is much faster than Flash-attention due to the sparsity. [..] why would one use the proposed method instead of Flash-attention?”**\n\nWe appreciate the reviewer’s concern regarding efficiency. As stated above, AdaSplash introduces adaptive, unstructured sparsity, which contrasts with fixed sparsity patterns that intentionally trade accuracy for speed. This makes implementation and optimization more challenging, but also more expressive and flexible.\n\nOur primary contribution is accelerating α-entmax attention, rather than introducing a generic sparse attention mechanism that merely approximates softmax attention. AdaSplash is the first efficient implementation of α-entmax, and already achieves speeds comparable to FlashAttention, despite being at an earlier stage of hardware/software optimization.\nIn practice, we have seen in Figures 4-5 that sparsity emerges in real attention patterns. Beyond speed, our method may also improve task performance and generalization, particularly for long-context scenarios where it may help to mitigate attention dispersion, a problem known to affect softmax-based dense attention [1, 2]. This may also explain why AdaSplash performs strongly across multiple benchmarks (Tables 1–3), especially in retrieval tasks.\n\n\n\n\n\n**“My biggest concern is that the estimation of threshold requires multiple access of QK memory. Thus it is difficult to be faster than FlashAttention. Online estimation might be better in term of speed.”**\n\nIt is true that, unlike FlashAttention which reads $Q, K, V$ tensors once in the forward pass, AdaSplash requires at least an additional read on $k$ and an additional $QK^\\top$ block computation. However, thanks to potential block-wise sparsity in α-entmax, we can skip entire blocks of $V$ (sidestepping a slow memory read) and avoid computing their contribution during the $PV$ step in the forward pass. Similarly, during the backward pass, masked blocks are skipped entirely, which makes, in practice, AdaSplash’s backward pass faster than FlashAttention’s. However, we do emphasize that this ultimately depends on the sparsity (as per Figure 1). \n\nAs for online estimation, we agree it's an interesting direction for future work. Unlike softmax, estimating the threshold $\\tau$ for α-entmax requires full access to the input, making online computation more challenging.  Thanks for the suggestion!\n\n\n\nReferences:\n\n\n[1] Veličković, Petar, et al. \"softmax is not enough (for sharp out-of-distribution).\" arXiv preprint arXiv:2410.01104 (2024).\n\n\n[2] Ye, Tianzhu, et al. \"Differential transformer.\" arXiv preprint arXiv:2410.05258 (2024).", "claims_and_evidence": "Claim 1: This method achieves substantial computational improvements over existing alpha-entmax implementations.\nThis is supported by clear and convincing evidence.\n\nClaim 2: This method can often match or even surpass the efficiency of highly optimized softmax-based attention algorithms like FlashAttention-2.\nThis is supported by the experiments. When the block-sparsity higher than about 95%, the proposed method is faster than FlashAttention-2.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense.", "theoretical_claims": "I carefully check the correctness of algorithm 1 and algorithm 2.\nI don't check the correctness of the backward pass algorithm.", "experimental_designs_or_analyses": "I check the speed comparison part (figure 1,2 and 3). \nThe authors only evaluate the running speed on random data (figure 3). They should also evaluate the running speed on real dataset.\nMoreover, the author should report the running time of Halley-bisection alone.\n\nI check the accuracy comparison part (table 1, 2 and 3).\nThe results support the authors claims. But I am not familiar with the benchmark in table 1.", "relation_to_broader_scientific_literature": "Flash attention /\nSparse attention /\nalpha-entmax /\nHalley’s method for faster bisection convergence", "essential_references_not_discussed": "This paper improves alpha-entmax attention, which is a special case of sparse attention.\nThe authors should discuss the pros and cons of alpha-entmax attention compared with other types of sparse attention, such as topk attention [1] and block sparse attention.\n\n[1] Memory-efficient Transformers via Top-k Attention, ACL2021", "other_strengths_and_weaknesses": "Weaknesses:\nThis paper focus on the efficiency of sparse attention, while Flash-attention is a dense attention. When compared with Flash-attention, we would expect that the proposed method is much faster than Flash-attention due to the sparsity. However, their speeds are very close. This limits practical applications of the proposed method—why would one use the proposed method instead of Flash-attention?\n\nMy biggest concern is that the estimation of threshold requires multiple access of QK memory. Thus it is difficult to be faster than FlashAttention. Online estimation might be better in term of speed.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Overall, I think this is a solid contribution with the following strengths:
- It provides an efficient algorithm and hardward-aware implementation of an adaptive input-dependent sparse attention.
- This sparse attention mechanism is thoroughly evaluated across various applications and context lengths, highlighting that the proposed adasplash for sparse attention can match the efficiency of FlashAttention, while providing minor predictive performance improvements. The evaluation appears quite thorough in terms of evaluating the attention mechanism across various different scenarios.

The main limitation appears to be that this mechanism has an additional memory overhead for really long context lengths, and can only at best match FlashAttention, though the one of promises of sparsity is the ability to be faster than the quadratic cost of dense attention. However, the authors mention that the memory overhead can be quite negligible even at 64k context length as the mask matrix is a binary matrix. The speedup over dense attention is only expected when the sparsity levels would be significantly large, which is a problem dependent property as $\alpha$-entmax sparse attention sets the attention sparsity in a problem dependent manner instead of enforcing a certain level of sparsity.
Decision: Accept (oral)