arXiv CSopen access

AdaSplash-2: Faster Differentiable Sparse Attention

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

machine learning, deep learning, neural networks

A DA S PLASH -2: Faster Differentiable Sparse Attention

Nuno M. T. Gonçalves * 1 2 3 Hugo Pitorro * 1 2 Vlad Niculae 4 Edoardo M. Ponti 5 Lei Li 3 André F. T. Martins 1 2 6 7 Marcos V. Treviso 1 2 7

Sparse attention has been proposed as a way to alleviate the quadratic cost of transformers, a central bottleneck in long-context training. A promising line of work is α-entmax attention, a differentiable sparse alternative to softmax that enables input-dependent sparsity yet has lagged behind softmax due to the computational overhead necessary to compute the normalizer τ . In this paper, we introduce A DA S PLASH -2, which addresses this limitation through a novel histogram-based initialization that reduces the number of iterations needed to compute τ to typically 1–2. The key idea is to compute a coarse histogram of attention scores on the fly and store it in on-chip SRAM, yielding a more accurate initialization that enables fast forward and backward computation. Combined with a sparsity-aware GPU implementation that skips zero blocks with low overhead, A DA S PLASH -2 matches or improves per-step training time relative to FlashAttention-2 when block sparsity is moderate-to-high (e.g., >60%), which often occurs at long-context lengths. On downstream tasks, models trained with our efficient α-entmax attention match softmax baselines at short-context lengths and achieve substantial gains in long-context settings.

Relative Speed (vs. FA2 Triton)

arXiv:2604.15180v1 [cs.LG] 16 Apr 2026

Abstract

FlashAttention-2 (CUDA) FlashAttention-2 (Triton)

2.25

AdaSplash AdaSplash-2

2.00 1.75 1.50 1.25 1.00 0.75 0.50 0

64x64 Block Sparsity (%)

100

Figure 1. Runtime (forward + backward) as a function of input sparsity for causal attention. A DA S PLASH -2, implemented in Triton, improves the sparsity-efficiency tradeoff, outperforming a highly-optimized CUDA version of FlashAttention-2 in moderate sparsity regimes and yielding larger gains at high block sparsity.

avoid this quadratic memory cost in the case of softmax attention. By making the computation IO-aware and by fusing operations over tiles that fit in on-chip fast memory, FlashAttention avoids storing intermediate scores and reduces memory complexity to linear in sequence length while achieving substantial speedups in training. The key operation requiring efficient implementations is normalization. For a score vector s, softmax can be written as X softmax(s) = exp(s − τ 1), τ = log exp(sj ), (1)

1. Introduction

The self-attention mechanism in transformers constitutes a major computational bottleneck since computing and materializing the score matrix S = QK ⊤ incurs quadratic time and memory complexity, limiting scalability to long sequences (Katharopoulos et al., 2020; Tay et al., 2022). FlashAttention (Dao et al., 2022; Dao, 2024; Shah et al., 2024) provides efficient GPU implementations that

where τ is an additive log-sum-exp normalizer. This form admits stable online accumulation and makes softmax particularly amenable to single-pass, tiled GPU implementations. FlashAttention and its successors (Dao, 2024; Shah et al., 2024) exploit precisely this structure. However, softmax attention is inherently dense as every token receives nonzero probability mass. For long inputs, this density can be undesirable—attention mass spreads over many irrelevant tokens (Veličković et al., 2025), and token representations become less distinguishable (Barbero et al., 2024)—motivating adaptive sparse alternatives. The α-entmax transformation (Peters et al., 2019) provides such

Equal contribution 1 Instituto Superior Técnico, Universidade de Lisboa 2 Instituto de Telecomunicações 3 Carnegie Mellon University 4 Language Technology Lab, University of Amsterdam 5 University of Edinburgh 6 TransPerfect 7 ELLIS Unit Lisbon. Correspondence to: Nuno Gonçalves <[email protected]>. Preprint. April 17, 2026.

A DA S PLASH -2: Faster Differentiable Sparse Attention

an alternative: it is differentiable, yields exact zeros in an input-dependent way, and has been shown to improve longcontext generalization relative to softmax (Vasylenko et al., 2025). Like softmax, α-entmax also requires computing a normalizer similar to τ in (1) (to be described in §2.3); however, its evaluation is not additive and requires iterative root-finding methods. This has historically prevented α-entmax attention from matching optimized softmax attention kernels in end-to-end training throughput.

moderate sparse regimes, but also matches or outperforms standard softmax attention on short- and long-context downstream tasks.

2. Background 2.1. Hardware Performance Modern GPUs are built for efficient parallel execution over a hierarchical memory system. High-bandwidth memory (HBM) provides large capacity but higher access latency than the smaller, faster on-chip SRAM (Jia et al., 2018). High performance therefore depends on efficient use of SRAM to reduce bottlenecks from frequent HBM traffic. GPUs run computation as kernels launched over thousands of threads grouped into thread blocks; data are staged from HBM into SRAM for computation and then written back. Kernel fusion is a core optimization that merges multiple operations into a single kernel, avoiding intermediate HBM reads/writes by directly producing final outputs. While compilers such as torch.compile can automate fusion for relatively simple operator chains (Ansel et al., 2024), attention mechanisms typically require custom strategies to reorder operations and optimize memory usage effectively.

While A DA S PLASH (Gonçalves et al., 2025) has recently mitigated this problem by introducing a GPU-friendly implementation of α-entmax attention, the provided algorithm still requires multiple expensive passes over the attention scores to compute τ . This limits efficiency, particularly in the forward pass and in moderate-sparsity regimes where sparsity is present but not extreme. In this work, we introduce A DA S PLASH -2, a new and more efficient hardware-aware method for α-entmax attention that reduces the computational cost of normalization. As we show in Figure 1, A DA S PLASH -2 outperforms a highlyoptimized CUDA implementation of FlashAttention-2 on moderate sparse regimes and can double the speedup at high sparsity cases. The key idea of our method is to construct a compact histogram of attention scores on the fly while streaming tiles through on-chip SRAM. This histogram yields a provable lower bound on the true normalizer τ , providing a high-quality initialization that allows the subsequent root solver to converge to the exact solution in one (typically) or two iterations, without materializing dense intermediates. Our contributions include:1

2.2. Standard Dot-Product Attention Given a set of matrices Q, K, V ∈ Rn×d containing ddimensional representations for n queries, keys and values, the dot-product self-attention at a single head is computed in the following way (Vaswani et al., 2017): S=

• Normalization via on-chip histogram. We derive a provable lower and upper bound on the α-entmax normalizer τ by computing a compact histogram of streamed scores entirely in on-chip SRAM, enabling fast α-entmax evaluation without forming dense intermediates.

O = π (S) V ∈ Rn×d .

(2)

The π transformation maps rows to normalized probability vectors, with softmax (Equation 1) being the most common choice. Crucially, the normalization τ is a stable reduction along each row of S and can be computed online in a single pass (Milakov & Gimelshein, 2018).

• One-pass refinement with a safeguarded hybrid solver. Initialized from the histogram estimate, a single additional pass typically suffices to recover the exact normalizer τ ⋆ using a safeguarded hybrid solver, yielding near-singleiteration convergence in practice for α ∈ {1.5, 2.0}.

FlashAttention. To address the costs of naive attention implementations, Dao et al. (2022) introduced FlashAttention, an algorithm that avoids the materialization of intermediate quadratic attention matrices via a GPU-aware implementation of online softmax (Milakov & Gimelshein, 2018), bringing the overall memory complexity to O (n). The key idea of FlashAttention is to split the inputs Q, K, V into blocks, load them from slow GPU high bandwidth 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. Later, FlashAttention-2 (Dao, 2024) improved the original algorithm and effectively defined the algorithmic structure adopted by subsequent variants, such as FlashAttention3 (Shah et al., 2024), which introduced hardware-specific

• Efficient exploitation of dynamic sparsity on GPU. We design an optimized attention kernel in Triton with finegrained tiling and a lightweight bit-packed encoding of nonzero blocks, enabling input-dependent sparsity to be exploited with negligible overhead. • Strong empirical results on synthetic and language modeling benchmarks. Our results show that A DA S PLASH -2 is not only faster than FlashAttention-2 on 1 Code available at: deep-spin/adasplash

QK ⊤ √ ∈ Rn×n , d

https://github.com/

A DA S PLASH -2: Faster Differentiable Sparse Attention

A DA S PLASH. Gonçalves et al. (2025) proposed a GPUoriented implementation of α-entmax attention that exploits dynamic sparsity and IO-awareness. Its main contribution is a hybrid Halley-bisection solver for the normalization threshold τ ⋆ , combining fast local convergence with bracketing guarantees. Implemented as custom Triton (Tillet et al., 2019) kernels, A DA S PLASH can skip zero blocks in forward and backward passes, improving performance at very high sparsity via a naive block masking strategy. However, refining τ ⋆ typically requires multiple iterations, each involving additional passes over the keys. In addition, its naive block masking strategy incurs extra overhead to determine which blocks can be skipped. Our method, described next, reduces the number of iterations and introduces a lightweight scheme to skip zero blocks.

optimizations for NVIDIA Hopper GPUs (e.g., TMA, WGMMA instructions, warp specialization) while preserving the same high-level algorithm. 2.3. α-entmax Transformation Softmax-based attention is inherently dense, as every key receives a strictly positive weight. A principled differentiable sparse alternative is the α-entmax transformation (Peters et al., 2019). For α > 1, α-entmax can produce sparse probability vectors, and it interpolates between softmax (α → 1) and sparsemax (α = 2, Martins & Astudillo 2016): 1

α-entmax(s) = [(α − 1)s − τ 1]+α−1 ,

(3)

where [·]+ is the ReLU function, and τ ∈ R is the (unique) normalizing constant which P ensures the output is a valid probability distribution, i α-entmax(s)i = 1. This means that coordinates with (α − 1)si ≤ τ become exactly zero. Importantly, α-entmax yields dynamic sparsity, where the pattern of zeros depends on the input s. Despite its flexibility, α-entmax does not have a direct closed-form solution, so computing it requires more involved methods, which we discuss next.

We present A DA S PLASH -2, a hardware-aware sparse attention mechanism that achieves efficient τ computation through a histogram-based approximation built entirely in on-chip SRAM. Our approach improves runtime compared to A DA S PLASH while maintaining the theoretical guarantees and sparsity benefits of α-entmax.

2.4. α-entmax Computation

3.1. Histogram Approximation

Computing α-entmax is equivalent to finding the normalizing constant τ ⋆ satisfying f (τ ⋆ ) = 0, where

Given attention scores s ∈ Rn , our goal is to compute the threshold τ ⋆ such that Equation 4 satisfies f (τ ⋆ ) = 0. Since α-entmax is invariant to adding a constant to all scores, we apply the change of variables

f (τ ) := −1 +

n X

3. Our Method

[(α − 1)sj − τ ]+α−1 .

(4)

j=1

z = (α − 1)s − (m − 1)1,

The function f is continuous and strictly decreasing in τ , hence the root τ ⋆ is unique (Blondel et al., 2019).

obtaining max(z) = 1. We henceforth work with centered scores z, which restricts both the active entries and the threshold search to the unit interval (see Equation 5).

Sorting-based solvers. For α = 1.5 and α = 2 (sparsemax), specialized solvers can be derived by sorting s and exploiting the fact that the support is given by the top-k entries (Michelot, 1986; Duchi et al., 2008; Condat, 2016; Peters et al., 2019). These approaches are less attractive on GPU because sorting and support selection are expensive and hard to fuse.

Histogram construction. With this change of variables, we discretize the interval [0, 1] into B equal-width bins with width h = 1/B. Formally, each score zj ∈ [0, 1] is mapped to a bin index as follows: b(zj ) = min(⌊Bzj ⌋, B − 1) .

Bisection (bracketing). For general α > 1, τ ⋆ can be found by a bracketing method such as bisection (Blondel et al., 2019), which only needs an interval containing the root. Let m = (α − 1) max(s). Following Peters et al. (2019), τ ⋆ satisfies m−1 ≤ τ

⋆

1−α

≤ m−n

(6)

(7)

At the same time we construct and maintain a histogram H in SRAM to store bin counts, Hk = { j : b(zj ) = k } ,

k = 0, . . . , B − 1.

(8)

Importantly, this representation only requires O(B) storage (independent of n), and entries with zj < 0 are not included in the histogram since they cannot belong to the active set.

(5)

which provides a valid initial bracket. While bisection is robust, it converges linearly and requires repeatedly evaluating f , motivating faster refinement strategies on GPU.

Histogram objective. Now in objective (4), we replace each score zj by the left endpoint of its bin. Concretely, if 3

A DA S PLASH -2: Faster Differentiable Sparse Attention

reverses the inequality to τh ≥ τ ⋆ , yielding a sparser output at the cost of potentially discarding relevant scores, while a centered approach b(zj ) = (⌊B · zj ⌋ + 0.5)/B may provide a more balanced approximation but does not induces any a one-sided bound guarantee. In A DA S PLASH -2, we adopt the left-edge binning strategy to ensure that the active support induced by τh contains the true support.

zj is assigned to bin k = b(zj ), we score it by z̃j := k/B. This allows us to rewrite a discretized objective as a sum over bins: fh (τ ) = −1 +

n X

[z̃j − τ ]+α−1

j=1

= −1 +

B−1 X

k=0 j:b(zj )=k

= −1 +

B−1 X k=0

Hk ·

1 α−1 k −τ B +

Smoothness and safe higher-order refinement. Following our histogram-based threshold initialization, we typically apply a posterior refinement step that evaluates f and its derivatives at candidate thresholds. Because each term 1/α−1 [zj − τ ]+ may be non-smooth at τ = zj , the existence and continuity of higher-order derivatives depend on α. The next proposition makes this connection explicit, and we use it to justify when higher-order updates are well-behaved.2

1 α−1

k −τ B +

(9)

Structurally, this construction replaces the original set of scores by B distinct values with associated counts Hk , yielding a reduced problem of size B ≪ n with the same monotone structure as the exact normalization in Equation 4. In turn, this reduced problem can be solved using any algorithm developed for α-entmax normalization. In particular, exact solvers apply for α = 1.5 and α = 2 (see §2.4). The key difference is that, in our setting, these solvers operate on a compact histogram representation that fits entirely in SRAM, minimizing the computation overhead compared to operating on the full set of n scores, which need to be recomputed from HBM. For completeness, we provide the corresponding solver adaptations for the histogram objective in Appendix C.3. We now formalize the approximation quality of the histogram objective and establish guarantees on the accuracy of the resulting normalization.

Proposition 2 (Continuity of f and α-entmax.). Let f (τ ) be defined as in Equation 4 for a given α. If 1 < α < t+1 t t t , then f ∈ C and, consequentially, α-entmax ∈ C . The proof is given in Appendix B.2. Proposition 2 clarifies which higher-order updates are numerically safe. When α ∈ (1, 1.5], f ′′ is bounded, so Halley updates (which use f ′′ ) are numerically stable.3 When α ∈ (1.5, 2], f ′ remains well-behaved but f ′′ becomes unbounded, making Newton updates more robust than Halley in practice. When α > 2, f ′ also becomes unbounded, and thus gradientbased methods are less stable. Based on this, we propose a new safeguarded hybrid solver: we apply Halley steps for α ≤ 1.5, Newton steps for 1.5 < α ≤ 2, and secant steps for α ≥ 2, always falling back to bisection whenever the proposed update falls outside the brackets. Combined with Proposition 1, this algorithm yields a refinement procedure that is both provably safe and fast in practice. We now outline the GPU-aware implementation of A DA S PLASH -2.

3.2. Theoretical Properties and Guarantees We first show that the histogram objective yields a conservative estimate of the true normalization threshold. We then characterize how the smoothness of f and α-entmax depends on α, motivating when safeguarded higher-order updates (e.g., Halley (Scavo & Thoo, 1995)) are appropriate. Proposition 1 (Histogram Lower Bound). Let α > 1 and h = 1/B be the bin width. Let f and fh be defined as in Equations 4 and 9. If τ ⋆ denotes the unique root of f (τ ) = 0 and τh the unique root of fh (τ ) = 0, then: τ ⋆ − h < τh ≤ τ ⋆ .

3.3. GPU-Aware Implementation Our implementation orchestrates multiple coordinated passes over the key matrix, and carefully exploits memory hierarchies for data movement between HBM and SRAM. Concretely, we perform a grid loop over Tr = ⌈n/Br ⌉ query blocks, and for a specific ith query block Qi ∈ RBr ×d c attending over key blocks {Kj }Tj=1 where Tc = ⌈n/Bc ⌉, our method proceeds as follows.

(10)

Thus, the absolute error satisfies 0 ≤ τ ⋆ − τh < h = B1 . The proof is given in Appendix B.1. This result ensures that the histogram approximation never overestimates the true threshold (τh ≤ τ ⋆ ) and provides explicit control over the maximum approximation error through the bin width h. In particular, the error satisfies 0 ≤ τ ⋆ − τh < h, and the approximation converges to the true solution as B → ∞.

Phase 1 (maximum). Stream key blocks to compute the row-wise maximum mi , which is scaled by (α − 1). We say f ∈ C t if f has continuous derivatives up to tth order. In fact from the monotonicity of f ′′ we have |f ′′ (τ )| ≤ n(2− α)/(α − 1)2 . This is an upper bound on the Lipschitz constant of f ′ , which by Kantorovich’s theorem controls the interval of convergence for Newton’s method (Kantorovich, 1949). 2 3

Other binning approaches lead to different precision-sparsity tradeoffs. Using a right-edge strategy b(zj ) = ⌈B · zj ⌉/B 4

A DA S PLASH -2: Faster Differentiable Sparse Attention Inner Loop

Q N×D

Mean Absolute Error

K N×D

Load

Bc D

Load

Bc 2

Br D

Grid Loop

Update histogram in SRAM

Zij

Bisection Hybrid Hybrid w/ hist. (B=4) Hybrid w/ hist. (B=8) Hybrid w/ hist. (B=16)

10000x closer to 𝜏* after 1 iteration in comparison to Hybrid without histogram

5 3 1

Hij 0 x 06 02 03 01 05 04 03 02

Iteration

Solve

𝜏i ≤ 𝜏*i

𝜏h

Figure 3. Comparison of mean absolute error of previous rootfinding methods and our Hybrid approach with histogram initialization, measured against the exact solution for α = 1.5.

Store

N×1 Grid Loop

P to compute Oi = j:Mij =1 α-entmax(Qi Kj⊤ , τ )Vj , enabling O(|M|) traversal complexity rather than O(Tr × Tc ).

Figure 2. Diagram of the A DA S PLASH -2 histogram kernel. For each query block Qi and key block Kj , we compute a score block Zij , update a per-row histogram H in SRAM, and use it to estimate an initial threshold τh for each query.

Histogram capacity. The bitpacking scheme naturally limits capacity since each bin can count up to 2b − 1 items before overflow. With Bc parallel accumulators per query block, the maximum capacity is

Phase 2 (histogram). Using mi , we center and scale scores into [0, 1] and build a B-bin per-row histogram entirely in SRAM while streaming over key blocks. To avoid atomics, we maintain local histogram accumulators for each position (i, j) in a query-key block and reduce them across columns after all keys are processed. Each accumulator packs B bins (representing bin counts) into a single w-bit integer (b = w/B bits per bin). A score in bin k updates its respective local accumulator via Hlocal ← Hlocal + 2kb . Final bin counts are recovered using vectorized shift-andmask operations with final reductions. Beyond fitting in SRAM, the histogram also provides an algorithmic speedup since binning implicitly orders scores by value. As a result, for α ∈ {1.5, 2}, we can apply exact solvers directly to the bin counts and recover τh without an explicit sort step. The diagram in Figure 2 illustrates this process.

Cmax = Bc × (2w/B − 1).

(11)

For our default configuration (w = 64, B = 8, Bc = 64), this yields Cmax = 16,320 keys per query block, sufficient for sequences up to 16K tokens. When sequences exceed this capacity, we employ a periodic flushing of bin counts to a Bc × B accumulator sitting in SRAM to avoid overflow. We provide more information on this in Appendix C.5. Backward pass. Crucially, gradients are nonzero only on the α-entmax support (Peters et al., 2019), we therefore reuse the block mask obtained in the forward pass to efficiently iterate only over nonzero blocks when computing the gradients with respect to queries, keys, and values. For this, we make use of native GPU instructions like fns and popc (population count) for efficient traversal. Further implementation details are presented in §C and §F, including the complete forward and backward algorithms.

Phase 3 (refinement). Starting from τh , we perform a fixed number of iterations to refine τ with a Hybrid solver that combines different root-finding methods (detailed in §3.2). In practice, we find that a single iteration is often sufficient for refinement due to the good accuracy of the histogram initialization. Simultaneously, we construct a binary block mask M ∈ {0, 1}Tr ×Tc indicating which blocks contain non-zero attention weights. The mask is stored using bitpacking, with each group of 32 consecutive column blocks encoded in a single int32, requiring only O (Tr × Tc ) bits of memory. In contrast to A DA S PLASH (Gonçalves et al., 2025), this design avoids auxiliary index buffers and incurs a negligible overhead, even for long contexts.

4. Experiments Our goals are to (i) verify that our histogram-based initialization for τ reduces convergence steps; (ii) quantify the training-time speedups of A DA S PLASH -2 relative to softmax-based baselines; and (iii) study model accuracy and long-context capabilities of LLMs trained with A DA S PLASH -2. 4.1. Normalization Solver Analysis

Phase 4 (output). Using mask M, we traverse only nonzero blocks via GPU-native find-next-set (fns) instructions

Before end-to-end GPU benchmarks, we first isolate the normalization step and study how quickly different root5

Runtime relative to FA-2 (Triton)

A DA S PLASH -2: Faster Differentiable Sparse Attention

1.50 1.25 1.00 0.75 0.50 0.25 0.00

FlashAttention-2 (Triton)

FlashAttention-2 (CUDA)

4.7ms

15ms

3.9ms

14ms 3.1ms

3.3ms

57ms

52ms 11ms

11ms

42ms

AdaSplash (Triton) 199ms

41ms

199ms

160ms

142ms

Forward Backward

AdaSplash-2 (Triton) 727ms

3.0s

645ms 583ms 461ms

2.4s

2.1s 1.5s

4096 51%

8192 58%

16384 64% 32768 70% Context Length - Block Sparsity

65536 75%

131072 85%

Figure 4. Runtime efficiency of causal self-attention implementations across context lengths of 4K-128K tokens with varying 64x64 block sparsity. Bar heights are normalized to FlashAttention-2 (Triton), with the opaque part denoting forward and the lighter part denoting backward. Numeric labels report the total forward + backward step time. Lower bars represent faster runtimes.

finding strategies recover the entmax threshold τ ⋆ . Concretely, we follow Gonçalves et al. (2025) and sample a Gaussian score vector of length n = 4096 and compute a high-accuracy reference threshold τ ⋆ using an exact method (see §2.4). We then run (i) bisection; (ii) our safeguarded hybrid solver without histogram initialization;4 and (iii) the same hybrid solver initialized with the histogram estimate τh using B ∈ {4, 8, 16} bins. Figure 3 reports the mean absolute error averaged over 10 runs for α = 1.5.

Sparsity-speed tradeoff. Figure 1 shows runtime as a function of block sparsity for random query, key and value tensors drawn from a zero-mean Gaussian. Following Gonçalves et al. (2025), we control the block sparsity of the attention probability matrix by changing the query variance, set head dimension to d = 64 and use bf16 precision. A DA S PLASH -2 consistently outperforms A DA SPLASH and, unlike FlashAttention, benefits from increased block sparsity, surpassing both Triton and CUDA implementations at moderate to high sparsity regimes. In the limit, A DA S PLASH -2 achieves more than a 2× speedup over both variants of FlashAttention-2.

After the first iteration, the histogram initialization places the solver orders of magnitude closer to the true root, making subsequent refinement into a near one-pass correction. As expected, more bins lead to better initial estimates.

Context scaling. In this experiment, we do not hand-tune the input sparsity. Instead, we use the block-sparsity ratios learned by a 1B-parameter transformer language model at different sequence lengths (we present the full language modeling results in the next subsection). We describe how these ratios are extracted from our models and provide block sparsity heatmaps across sequence lengths in §E. Figure 4 reports the average per-step training runtime (forward + backward) as a function of context length, with bar heights normalized to FlashAttention-2 Triton (lower is better) and decomposed into forward and backward contributions.

4.2. Efficiency Benchmark We implement A DA S PLASH -2, described in Section 3.1, in Triton. We compare it with FlashAttention-2 (Dao, 2024) implemented in CUDA and in Triton,5 and A DA SPLASH written in Triton (Gonçalves et al., 2025). A natural question is why compare with FlashAttention-2 rather than FlashAttention-3. Our efficiency experiments target NVIDIA Ampere GPUs, for which FlashAttention-2 is the reference implementation. FlashAttention-3 (Shah et al., 2024) introduces hardware-specific optimizations for NVIDIA Hopper GPUs (e.g., TMA, WGMMA instructions, warp specialization) while preserving the same high-level algorithm. This choice allows us to isolate algorithmic effects independent of such hardware tuning. We note that the low-level optimizations in FlashAttention-3 are largely orthogonal to A DA S PLASH -2 and can be incorporated with additional engineering effort. All runtime benchmarks are performed in a single A6000 NVIDIA GPU.

As expected, A DA S PLASH is slower than FlashAttention-2 in the forward pass due to the additional passes required to compute the normalization threshold τ (phases 1, 2, and 3 in §3.3). However, as context length increases, block sparsity naturally emerges, reducing the forward pass gap. In the backward pass, even at shorter contexts, existing sparsity already yields faster runtimes for A DA S PLASH -2, with the advantage increasing at longer contexts. Since the backward pass dominates training time, these gains are substantial. This suggests that as models scale to longer contexts, A DA S PLASH -2 can translate naturally emerging sparsity into training efficiency gains, contrasting with the dense nature of softmax-based attention.

4 For α = 1.5, our Hybrid solver empirically matches the Halley-Bisection solver of Gonçalves et al. (2025). 5 https://github.com/triton-lang/triton/blob/main/python/tutorial s/06-fused-attention.py

A DA S PLASH -2: Faster Differentiable Sparse Attention Table 1. RULER benchmark results for 1B parameter models trained up to 32K context-length. Best average results are in bold. Model

QA-H QA-S Avg.

4K Softmax 8K (RoPE) 16K 32K

80.0 78.4 67.6 41.4

62.2 45.0 10.0 8.2

38.2 30.2 14.8 0.8

47.6 36.7 22.1 17.4

43.4 31.0 22.7 20.8

100.0 100.0 100.0 96.2

100.0 99.2 99.2 61.6

98.8 98.2 99.2 75.6

5.1 0.1 0.0 0.0

24.9 24.3 19.0 6.3

8.0 2.8 0.2 0.0

24.8 22.4 17.0 16.8

35.0 20.2 18.7 10.1

51.4 45.3 37.7 27.3

4K 8K 16K 32K

75.4 71.8 53.2 33.2

26.0 19.4 8.6 4.4

23.6 8.6 1.0 0.6

29.9 22.5 12.9 7.5

52.3 42.5 22.0 12.9

100.0 100.0 100.0 97.0

100.0 99.8 81.4 40.8

98.8 94.6 79.4 64.2

36.3 8.9 0.1 0.0

32.3 23.7 18.6 15.0

3.9 3.6 1.2 0.2

25.8 23.2 18.2 18.6

36.7 19.6 20.8 9.2

49.3 41.4 32.1 23.4

4K Softmax 8K (NAPE) 16K 32K

88.2 78.4 85.6 66.6

29.8 38.6 19.6 4.2

13.0 12.6 7.4 2.0

28.9 33.2 26.8 26.4

27.9 34.1 30.6 26.5

100.0 100.0 100.0 100.0

100.0 100.0 100.0 97.0

100.0 99.0 98.8 98.0

35.3 27.0 18.0 8.3

27.4 26.3 25.6 3.3

14.8 12.3 9.8 7.7

21.8 19.2 20.0 19.0

37.9 22.9 21.3 19.5

48.1 46.4 43.3 36.8

4K Entmax 8K (NAPE) 16K 32K

81.2 73.4 79.2 63.0

12.2 4.6 1.4 1.8

27.2 17.6 9.8 4.8

51.9 62.3 38.7 25.8

55.6 31.5 24.0 16.8

100.0 100.0 100.0 100.0

100.0 100.0 100.0 92.0

79.8 84.0 74.2 76.0

48.9 33.7 22.9 11.7

58.1 48.3 51.1 45.8

31.4 38.4 34.5 27.4

30.2 29.0 27.2 25.0

34.5 26.0 27.5 22.1

54.7 49.9 45.4 39.4

Entmax (RoPE)

Len. MK-1 MK-2 MK-3 MQ MV NIAH-1 NIAH-2 NIAH-3 CWE FWE

4.3. Language Modeling Benchmarks

Table 2. Results across context lengths for the In-Context Learning tasks in the HELMET benchmark. Best average result is in bold.

To evaluate the effectiveness of α-entmax attention in language modeling tasks, we train 350M and 1B parameters versions of the LLaMA-3 architecture (Grattafiori et al., 2024) using A DA S PLASH -2 with α = 1.5. Our experimental design focuses on general capabilities and long-context performance and, thus, has two distinct settings: (i) one version that does pretraining from scratch for 50B tokens from the DCLM-Edu dataset (Allal et al., 2025), and (ii) an alternative version with a context extension phase with the ProLong (Gao et al., 2025) methodology in the final 20% of tokens, endowing our models with 32K context length. Furthermore, following Vasylenko et al. (2025), who show that RoPE can be suboptimal for entmax attention, we also evaluate their proposed positional encoding scheme, NAPE: within each layer, half of the heads use no positional encoding (NoPE; Kazemnejad et al. 2023) and the other half use ALiBi (Press et al., 2022). We compare A DA S PLASH -2 runs against softmax baselines with RoPE (Su et al., 2024) and with NAPE to ensure a fair comparison. Additional experimental details are provided in §D.

Model

Len. TREC-C TREC-F NLU B77 C150 Avg.

Softmax 8K (RoPE) 16K 32K

62.4 67.2 69.2

33.6 42.2 47.8

12.6 6.0 13.6 25.6 26.4 10.8 25.8 34.5 25.8 12.8 35.6 38.2

Entmax 8K (RoPE) 16K 32K

63.2 66.4 66.8

26.6 35.8 42.8

25.4 8.4 20.2 28.8 32.8 13.8 27.4 35.2 36.2 16.8 34.0 39.3

Softmax 8K (NAPE) 16K 32K

58.8 64.4 71.2

30.8 35.8 45.6

27.2 13.4 32.6 32.6 38.6 20.6 43.2 40.5 41.4 21.2 51.0 46.1

Entmax 8K (NAPE) 16K 32K

78.0 83.4 85.0

43.8 56.4 62.4

44.6 19.4 41.0 45.4 58.0 25.0 57.0 56.0 60.6 34.4 68.4 62.2

Next, we evaluate the models’ in-context learning (ICL) performance using the ICL subset from the HELMET benchmark (Yen et al., 2025), which measures few-shot performance across increasing context lengths. Table 2 shows that α-entmax substantially improves ICL accuracy, especially when paired with NAPE, achieving the highest average score at every tested context length and outperforming softmax baselines. Notably, while all models improve as the context grows from 8K to 32K, we observe particularly large gains with α-entmax + NAPE, suggesting that sparsity helps the model in leveraging the relevant in-context examples, reinforcing the finding that α-entmax + NAPE constitutes a strong synergy for long-context modeling.

Long-context results. We start by investigating the performance of our models on the full RULER benchmark (Hsieh et al., 2024), evaluated up to 32K context length. The results are presented in Table 1. Overall, α-entmax with NAPE achieves the strongest average performance across all sequence lengths, outperforming both softmax baselines. Among the different tasks, we note that the gains are particularly pronounced on Variable Tracking (VT)—a core subproblem underlying complex reasoning (Feng & Steinhardt, 2024; Dai et al., 2024)—as well as on common and frequent word extraction (CWE and FWE), which require precise aggregation of the relevant input spans.

Short-context results. For evaluating short-context performance, we follow the few-shot prompting strategy under the OLMES evaluation standard (Gu et al., 2025). The 7

A DA S PLASH -2: Faster Differentiable Sparse Attention Table 3. Short-context benchmark results for 350M and 1B models with 4K context length. Scores are computed with the OLMES framework using the core 9mcqa suite. Best perplexity and average results are in bold. Model

LMB (ppl) LMB ARC-E ARC-C CSQA

OBQA PIQA SocialQA WG Avg.

350M params. Softmax with RoPE Entmax with RoPE Softmax with NAPE Entmax with NAPE

23.93 22.36 19.23 18.62

40.6 39.5 41.2 42.4

61.9 63.0 61.8 61.9

32.1 34.8 33.4 33.0

50.5 50.5 47.6 51.2

41.6 40.2 40.5 41.0

38.4 39.0 38.0 39.2

66.6 64.8 64.8 66.0

45.1 44.1 43.5 44.4

48.9 53.4 52.7 53.7

47.3 47.7 47.1 48.1

1B params. Softmax with RoPE Entmax with RoPE Softmax with NAPE Entmax with NAPE

15.01 15.76 11.97 11.42

44.7 43.9 48.0 49.2

69.0 65.4 69.3 67.7

36.0 36.6 37.4 39.9

58.7 56.2 56.7 57.3

49.6 47.4 48.8 48.7

42.8 41.6 45.2 45.0

70.0 69.8 69.5 68.3

46.4 45.7 46.5 47.1

55.6 54.2 55.2 55.1

52.5 51.2 53.0 53.1

results are shown in Table 3. We first observe a substantial perplexity gap between softmax and α-entmax models on LAMBADA (LMB). Second, we note that α-entmax models are on par or better across the board on accuracybased downstream tasks at both model scales (350M and 1B params.). Taken together, these results highlight the strengths of sparse attention by not just outperforming dense models on long-context settings, but also on short-context tasks.

that exploit GPU memory hierarchies. FlexAttention (Dong et al., 2024) offers a programmable interface for custom masks/score modifications while using similar fused, tiled execution. A DA S PLASH (Gonçalves et al., 2025) offers a hardware-aware implementation of α-entmax attention by solving τ with iterative refinement, but incurs considerable overhead due to repeated scans over keys and naive block-masking. Long-context modeling. A growing body of work studies why transformer performance degrades as context length grows, attributing failures to the softmax function itself such as attention dispersion (Zhai et al., 2023; Veličković et al., 2025) and representational collapse (Barbero et al., 2024; Arroyo et al., 2025). Complementary lines improve extrapolation through positional bias design (Press et al., 2022; Jelassi et al., 2024) and through length/entropy-aware scaling (Peng et al., 2024; Nakanishi, 2025; Zhang et al., 2024). In this context, a sparse learnable-temperature alternative with α-entmax have been proposed as a direct way to mitigate these issues (Vasylenko et al., 2025).

5. Related Works α-entmax computation. Exact computation exists for special cases such as α ∈ {1.5, 2} using sorting-based solvers (Duchi et al., 2008; Condat, 2016; Peters et al., 2019). For general α, the normalization τ can be computed using root-finding procedures (Blondel et al., 2019). However, these approaches remain suboptimal due to slow convergence or reliance on complex data structures and sorting operations, which are difficult to optimize for hardware. Sparse attention mechanisms. Sparse attention has been widely studied through fixed sparsity patterns (Zaheer et al., 2020; Beltagy et al., 2020). Recent methods aim for datadependent sparsity, commonly via top-k selection followed by softmax over the selected subset (Yuan et al., 2025; Liu et al., 2025; Nawrot et al., 2025). Top-k is attractive for inference but can be unstable in training due to discontinuities. In contrast, α-entmax induces exact, differentiable, input-adaptive sparse attention. A DA S PLASH -2 focuses on making α-entmax attention efficient during training; while extending it to inference is non-trivial as the forward pass requires multiple key scans, and we leave a fully optimized inference kernel for future work.

6. Conclusion In this work, we introduced A DA S PLASH -2, a hardwareaware efficient implementation of α-entmax attention. Our key idea is to construct a compact histogram of streamed attention scores in on-chip SRAM, producing a strong initialization for the α-entmax threshold τ , and reducing refinement to a small (typically 1-2) number of steps while integrating naturally with GPU execution. To leverage sparsity, A DA S PLASH -2 uses a lightweight bitpacked block mask to skip zero blocks efficiently. Empirically, A DA S PLASH -2 delivers training speedups over prior α-entmax kernels and, in moderate-to-high block-sparsity regimes, matches or exceeds FlashAttention-2, with gains driven primarily by a substantially faster backward pass. In language modeling, α-entmax models match or improve short-context performance and deliver clear gains on long-context evaluations.

GPU-aware attention for training. FlashAttention and its successors (Dao et al., 2022; Dao, 2024; Shah et al., 2024) provide highly optimized CUDA kernels for softmax attention via fused normalization and tiled accumulation 8

A DA S PLASH -2: Faster Differentiable Sparse Attention

Impact Statement

’24, pp. 929–947, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703850. doi: 10.1145/3620665.3640366. URL https://doi. org/10.1145/3620665.3640366.

Our method, A DA S PLASH -2, provides an efficient implementation of α-entmax attention. Efficient attention mechanisms are crucial for scaling transformers to handle longcontext sequences. As a result, the improved efficiency has potential applications in large-scale NLP applications, especially in cases where sparsity can be leveraged to reduce computational costs such as in long-context modeling. We do not foresee direct societal consequences from our method itself, but its integration into decision-making models may still reflect biases in training data. As such, we encourage careful evaluation when deploying models trained with A DA S PLASH -2 in high-stakes applications, ensuring that efficiency gains do not overcome ethical concerns.

Arroyo, A., Gravina, A., Gutteridge, B., Barbero, F., Gallicchio, C., Dong, X., Bronstein, M. M., and Vandergheynst, P. On vanishing gradients, over-smoothing, and oversquashing in GNNs: Bridging recurrent and graph learning. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https: //openreview.net/forum?id=N4cyRMuLyl. Barbero, F., Banino, A., Kapturowski, S., Kumaran, D., Araújo, J. G. M., Vitvitskyi, A., Pascanu, R., and Veličković, P. Transformers need glasses! information over-squashing in language tasks. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/ forum?id=93HCE8vTye.

Acknowledgments We would like to the SARDINE lab team for the helpful discussions. This work was supported by the project DECOLLAGE (ERC-2022-CoG 101088763), by the Portuguese Recovery and Resilience Plan through project C64500888200000055 (Center for Responsible AI), and by FCT/MECI through national funds and when applicable co-funded EU funds under UID/50008: Instituto de Telecomunicações. Vlad Niculae is supported by the Dutch Research Council (NWO) via VI.Veni.212.228. Edoardo M. Ponti is supported by the ERC Starting Grant AToM-FM (101222956).

Basant, A., Khairnar, A., Paithankar, A., Khattar, A., Renduchintala, A., Malte, A., Bercovich, A., Hazare, A., Rico, A., Ficek, A., et al. Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model. arXiv preprint arXiv:2508.14444, 2025. URL https://arxiv.org/abs/2508.14444. Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv:2004.05150, 2020. URL https://arxiv.org/abs/2004.05150.

References

Blondel, M., Martins, A., and Niculae, V. Learning classifiers with fenchel-young losses: Generalized entropies, margins, and algorithms. In Chaudhuri, K. and Sugiyama, M. (eds.), Proceedings of the Twenty-Second International 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.mlr.press/v89/ blondel19a.html.

Allal, L. B., Lozhkov, A., Bakouch, E., Blázquez, G. M., Penedo, G., Tunstall, L., Marafioti, A., Kydlı́ček, H., Lajarı́n, A. P., Srivastav, V., Lochner, J., Fahlgren, C., Nguyen, X.-S., Fourrier, C., Burtenshaw, B., Larcher, H., Zhao, H., Zakka, C., Morlon, M., Raffel, C., von Werra, L., and Wolf, T. Smollm2: When smol goes big – datacentric training of a small language model, 2025. URL https://arxiv.org/abs/2502.02737. Ansel, J., Yang, E., He, H., Gimelshein, N., Jain, A., Voznesensky, M., Bao, B., Bell, P., Berard, D., Burovski, E., Chauhan, G., Chourdia, A., Constable, W., Desmaison, A., DeVito, 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, Volume 2, ASPLOS

Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019. URL https: //arxiv.org/abs/1905.10044. Condat, L. Fast projection onto the simplex and the l1 ball. Mathematical Programming, 158(1):575–585, 2016. Dai, Q., Heinzerling, B., and Inui, K. Representational analysis of binding in language models. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 17468–17493, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.emnlp-main. 9

A DA S PLASH -2: Faster Differentiable Sparse Attention

967. URL https://aclanthology.org/2024. emnlp-main.967/. Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum? id=mZn2Xyh9Ec.

Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. URL https://arxiv.org/abs/2407.21783. Gu, Y., Tafjord, O., Kuehl, B., Haddad, D., Dodge, J., and Hajishirzi, H. OLMES: A standard for language model evaluations. In Chiruzzo, L., Ritter, A., and Wang, L. (eds.), Findings of the Association for Computational Linguistics: NAACL 2025, pp. 5005– 5033, Albuquerque, New Mexico, April 2025. Association for Computational Linguistics. ISBN 979-889176-195-7. doi: 10.18653/v1/2025.findings-naacl. 282. URL https://aclanthology.org/2025. findings-naacl.282/.

Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //openreview.net/forum?id=H4DqfPSibmx. Dong, J., Feng, B., Guessous, D., Liang, Y., and He, H. Flex attention: A programming model for generating optimized attention kernels, 2024. URL https: //arxiv.org/abs/2412.05496.

Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, J. W., and Sifre, L. An empirical analysis of compute-optimal large language model training. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //openreview.net/forum?id=iBBcRUlOAPR.

Dontchev, A. L. and Rockafellar, R. T. Implicit Functions and Solution Mappings: A View from Variational Analysis. Springer Series in Operations Research and Financial Engineering. Springer Nature, New York, NY, second edition. edition, 2014. ISBN 149391037X. Duchi, J., Shalev-Shwartz, S., Singer, Y., and Chandra, T. Efficient projections onto the l1-ball for learning in high dimensions. In Proceedings of the 25th international conference on Machine learning, pp. 272–279, 2008.

Hsieh, C.-P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., and Ginsburg, B. RULER: What’s the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL https: //openreview.net/forum?id=kIoBbc76Sy.

Feng, J. and Steinhardt, J. How do language models bind entities in context? In The Twelfth International Conference on Learning Representations, 2024. URL https: //openreview.net/forum?id=zb3b6oKO77.

Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, E. Repeat after me: Transformers are better than state space models at copying. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 21502–21521. PMLR, 21–27 Jul 2024. URL https://proceedings.mlr.press/ v235/jelassi24a.html.

Gao, T., Wettig, A., Yen, H., and Chen, D. How to train long-context language models (effectively). In Che, W., Nabende, J., Shutova, E., and Pilehvar, M. T. (eds.), Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7376–7399, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176-2510. doi: 10.18653/v1/2025.acl-long.366. URL https: //aclanthology.org/2025.acl-long.366/.

Jia, Z., Maggioni, M., Staiger, B., and Scarpazza, D. P. Dissecting the nvidia volta gpu architecture via microbenchmarking. arXiv preprint arXiv:1804.06826, 2018.

Gelberg, Y., Eguchi, K., Akiba, T., and Cetin, E. Extending the Context of Pretrained LLMs by Dropping their Positional Embeddings. Technical report, Sakana AI, January 2026. Technical Report.

Kantorovich, L. V. On Newton’s Method. Trudy Mat. Inst. Steklov, 28:104–144, 1949. URL https://cs.uwaterloo.ca/˜y328yu/ classics/Kantorovich48.pdf.

Gonçalves, N., Treviso, M. V., and Martins, A. Adasplash: Adaptive sparse flash attention. In Fortysecond International Conference on Machine Learning, 2025. URL https://openreview.net/forum? id=OWIPDWhUcO.

Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers 10

A DA S PLASH -2: Faster Differentiable Sparse Attention

with linear attention. In International Conference on Machine Learning, pp. 5156–5165. PMLR, 2020.

Paperno, D., Kruszewski, G., Lazaridou, A., Pham, N. Q., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Erk, K. and Smith, N. A. (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1525–1534, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10.18653/v1/P16-1144. URL https://aclanthology.org/P16-1144/.

Kazemnejad, A., Padhi, I., Natesan, K., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum? id=Drrl2gcjzl. Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., and Idreos, S. Torchtitan: One-stop pytorch native solution for production ready LLM pretraining. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //openreview.net/forum?id=SFN6Wm7YBI.

Peng, B., Quesnelle, J., Fan, H., and Shippole, E. YaRN: Efficient context window extension of large language models. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview. net/forum?id=wHBfxhZu1u.

Liu, A., Mei, A., Lin, B., Xue, B., Wang, B., Xu, B., Wu, B., Zhang, B., Lin, C., Dong, C., et al. Deepseek-v3. 2: Pushing the frontier of open large language models. arXiv preprint arXiv:2512.02556, 2025. URL https: //arxiv.org/abs/2512.02556. Martins, A. and Astudillo, R. From softmax to sparsemax: A sparse model of attention and multi-label classification. 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 http://proceedings.mlr. press/v48/martins16.html. Michelot, C. A finite algorithm for finding the projection of a point onto the canonical simplex of n. Journal of Optimization Theory and Applications, 50(1):195–200, 1986.

Peters, B., Niculae, V., and Martins, A. F. T. Sparse sequence-to-sequence models. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1504–1519, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 18653/v1/P19-1146. URL https://www.aclweb. org/anthology/P19-1146. Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL https://openreview.net/ forum?id=R8sQPpGCv0. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. Rudin, W. Principles of Mathematical Analysis. McGraw– Hill, New York, USA, 3rd edition, 1976. ISBN 9780070856134.

Milakov, M. and Gimelshein, N. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867, 2018. URL https://arxiv.org/abs/1805.02867.

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.jstor.org/stable/2975033.

Nakanishi, K. M. Scalable-softmax is superior for attention. arXiv preprint arXiv:2501.19399, 2025. URL https: //arxiv.org/abs/2501.19399. Nawrot, P., Li, R., Huang, R., Ruder, S., Marchisio, K., and Ponti, E. M. The sparse frontier: Sparse attention trade-offs in transformer llms. arXiv preprint arXiv:2504.17768, 2025. URL https://arxiv. org/abs/2504.17768. Olmo, T., Ettinger, A., Bertsch, A., Kuehl, B., Graham, D., Heineman, D., Groeneveld, D., Brahman, F., Timbers, F., Ivison, H., et al. Olmo 3. arXiv preprint arXiv:2512.13961, 2025. URL https://arxiv. org/abs/2512.13961. 11

Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/ forum?id=tVConYid20. Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomput., 568(C), February 2024. ISSN 0925-2312. doi: 10.1016/j.neucom. 2023.127063. URL https://doi.org/10.1016/ j.neucom.2023.127063.

A DA S PLASH -2: Faster Differentiable Sparse Attention

Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. ACM Comput. Surv., 55(6), December 2022. ISSN 0360-0300. doi: 10.1145/3530811. URL https://doi.org/10.1145/3530811.

Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., and Ahmed, A. Big bird: Transformers for longer sequences. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17283–17297. Curran Associates, Inc., 2020. URL https://proceedings.neurips. cc/paper_files/paper/2020/file/ c8512d142a2d849725f31a9a7a361ab9-Paper. pdf.

Tillet, P., Kung, H. T., and Cox, D. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL 2019, pp. 10–19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367196. doi: 10. 1145/3315508.3329973. URL https://doi.org/ 10.1145/3315508.3329973.

Zhai, S., Likhomanenko, T., Littwin, E., Busbridge, D., Ramapuram, J., Zhang, Y., Gu, J., and Susskind, J. M. Stabilizing transformer training by preventing attention entropy collapse. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 40770–40803. PMLR, 23–29 Jul 2023. URL https://proceedings.mlr.press/ v202/zhai23a.html.

Tsallis, C. Possible generalization of boltzmann-gibbs statistics. Journal of statistical physics, 52(1):479–487, 1988.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017. URL https: //papers.nips.cc/paper/2017/hash/ 3f5ee243547dee91fbd053c1c4a845aa-Abstract. Zhang, Y., Li, J., and Liu, P. Extending llms’ context winhtml. dow with 100 samples. arXiv preprint arXiv:2401.07004, 2024. URL https://arxiv.org/abs/2401. Vasylenko, P., Treviso, M., and Martins, A. F. Long07004. context generalization with sparse attention. arXiv preprint arXiv:2506.16640, 2025. URL https:// arxiv.org/abs/2506.16640. Veličković, P., Perivolaropoulos, C., Barbero, F., and Pascanu, R. Softmax is not enough (for sharp size generalisation). In Forty-second International Conference on Machine Learning, 2025. URL https://openreview. net/forum?id=S4JmmpnSPy. Yen, H., Gao, T., Hou, M., Ding, K., Fleischer, D., Izsak, P., Wasserblat, M., and Chen, D. HELMET: How to evaluate long-context models effectively and thoroughly. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview. net/forum?id=293V3bJbmE. Yuan, J., Gao, H., Dai, D., Luo, J., Zhao, L., Zhang, Z., Xie, Z., Wei, Y., Wang, L., Xiao, Z., Wang, Y., Ruan, C., Zhang, M., Liang, W., and Zeng, W. Native sparse attention: Hardware-aligned and natively trainable sparse attention. In Che, W., Nabende, J., Shutova, E., and Pilehvar, M. T. (eds.), Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 23078–23097, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176-251-0. doi: 10.18653/v1/2025.acl-long.1126. URL https:// aclanthology.org/2025.acl-long.1126/. 12

A DA S PLASH -2: Faster Differentiable Sparse Attention

softmax (

entmax ( = 1.5)

top-k softmax (k = 2)

sparsemax ( = 2)

1.0

0.8

0.6

0.4

0.2

0.0

0.2 0.0

Figure 5. Visualization of α-entmax for different values of α. We also include top-k softmax with k = 2 for completeness. Each panel shows how the probability mass of p0 varies for the input z = [0, z1 , z2 ]. For softmax, p0 is always non-zero, regardless of z1 and z2 . As α increases, α-entmax increasingly assigns exactly zero probability to z0 . While α-entmax changes smoothly with the scores (yielding piecewise-smooth gradients), top-k softmax changes the “selected set” abruptly and, consequently, these non-differentiable boundaries induce discontinuous gradients that might lead to training instabilities.

A. α-entmax Transformation The α-entmax transformation of a score vector z ∈ Rn is defined as follows (Peters et al., 2019): α-entmax(z) := arg max p⊤ z + Hα (p), p∈△n

△n := {p ∈ Rn : p ≥ 0, 1⊤ p = 1},

(12)

where Hα (p) is the Tsallis(α) entropy (Tsallis, 1988). Solving the optimization problem above corresponds to finding a threshold τ so that p⋆ sums to 1. Given τ , we can easily evaluate α-entmax as per Equation 3, which we re-state here for easiness: 1/α−1 α-entmax(z) = [(α − 1)z − τ 1]+ , (13) where [·]+ is the ReLU function. Figure 5 illustrates how α-entmax(z) behaves for different choices of α. From the equation above, it is clear that coordinates with (α − 1)si ≤ τ become exactly zero. In other words, α-entmax yields dynamic sparsity, where the pattern of zeros depends on the input. However, this flexibility comes at a computational cost: unlike softmax, α-entmax does not have a direct closed-form solution but rather requires more involved methods, which we discuss next. α-entmax computation. Order-based algorithms have only been proposed for α = 2 (Michelot, 1986; Duchi et al., 2008; Condat, 2016) and α = 1.5 (Peters et al., 2019). This family of algorithms builds upon the equivalence between τ and the “active set”, denoted as S = {j : zj > τ }. At τ ⋆ , this set indexes the nonzero terms in the expression of f , or, equivalently, the nonzero elements in the α-entmax output. If the active set contains an index, it must contain all indices with greater or equal value, so there are n possible active sets S1 , . . . , Sn , where contains the indices of the k largest values of z. Moreover, since the output must be a valid probability distribution, Sk is never empty. At the true support size k ⋆ , we have f (τ ) = fk⋆ (τ ) := −1 +

(zj − τ ) /α−1 . 1

(14)

j∈Sk⋆

Note that the ReLU function is no longer necessary as all terms are strictly positive. Finding the root of fk for any k can be done efficiently in the case of α = 2 (linear equation) and α = 1.5 (quadratic equation). In contrast, root-finding algorithms apply for all values of α. Since max(z) = 1, following Peters et al. (2019), we have 0 ≤ τ ⋆ ≤ 1 − n1−α .

(15)

Moreover, since f is continuous and f (0) ≤ 0 and f (1 − n1−α ) ≥ 0, the root must be found in the interval. Blondel et al. (2019) propose a bisection or binary search approach to finding τ ⋆ . Similar in spirit, Gonçalves et al. (2025) introduces a GPU-oriented solver of α-entmax that uses a hybrid Halley-bisection method, which combines the fast local convergence of higher-order root-finding (Scavo & Thoo, 1995) with the convergence guarantees of bisection. 13

A DA S PLASH -2: Faster Differentiable Sparse Attention

B. Proofs Throughout, we let α > 1, represent scores as s ∈ Rn , and we work with centered and scaled scores z = (α−1)s−(m−1)1, where m = (α − 1) max(s), which ensures max(z) = 1. Also, we recall from Equation 15 that τ ≥ max(z) − 1 = 0 and τ ≤ max(z) − n1−α = 1 − n1−α . B.1. Proof of Proposition 1 Proof of Proposition 1. We divide the proof in four parts. Step 1: Establishing the sandwich inequality. By construction of the binning scheme, for any score zj ∈ [0, 1] assigned to bin k, we have: k+1 k ≤ zj < , (16) B B which implies: 1 bj ≤ zj < bj + h, where h = , (17) B k where bj := B denotes the left edge of the bin to which zj is assigned. For scores zj < 0, we have [zj − τ ]+ = 0 = [bj − τ ]+ for all τ > 0, so these scores do not affect the proof. Consider the 1

function ϕ(t) = [t]+α−1 for t ∈ R, which is monotone non-decreasing and continuous for α > 1. For any fixed τ ∈ [0, 1], using the monotonicity of ϕ and inequality in Eq. 17: bj − τ ≤ zj − τ < bj + h − τ

(18)

=⇒

[bj − τ ]+ ≤ [zj − τ ]+ < [bj + h − τ ]+

(19)

=⇒

ϕ(bj − τ ) ≤ ϕ(zj − τ ) < ϕ(bj + h − τ ).

(20)

The strict inequality on the right holds because: if zj − τ > 0, then zj < bj + h implies zj − τ < bj + h − τ , and ϕ is strictly increasing on (0, ∞). Summing over all j = 1, . . . , n: n X j=1

ϕ(bj − τ ) ≤

n X

ϕ(zj − τ ) <

j=1

n X

ϕ(bj + h − τ ).

(21)

j=1

Subtracting 1 from each term: fh (τ ) ≤ f (τ ) < fh (τ − h),

∀τ ∈ [0, 1],

(22)

where we used the fact that: n X j=1

Step 2: Properties of f and fh .

ϕ(bj + h − τ ) =

n X

ϕ(bj − (τ − h)) = fh (τ − h) + 1.

(23)

j=1

Both f (τ ) and fh (τ ) are:

• Continuous on [0, 1] • Strictly decreasing on [0, 1] (since ϕ is strictly increasing where positive) • Satisfy f (0) > 0 and f (1) < 0 (similarly for fh ) By the intermediate value theorem, there exist unique roots: τ⋆ :

f (τ ⋆ ) = 0,

(24)

τh :

fh (τh ) = 0.

(25)

A DA S PLASH -2: Faster Differentiable Sparse Attention

Step 3: Bounding τh relative to τ ⋆ .

Evaluating the sandwich inequality in Eq. 22 at τ = τ ⋆ : fh (τ ⋆ ) ≤ f (τ ⋆ ) < fh (τ ⋆ − h).

(26)

fh (τ ⋆ ) ≤ 0 < fh (τ ⋆ − h).

(27)

Since f (τ ⋆ ) = 0, we have: Because fh is strictly decreasing and continuous, and fh (τ ⋆ − h) > 0 while fh (τ ⋆ ) ≤ 0, there must exist a unique τh ∈ [τ ⋆ − h, τ ⋆ ] such that fh (τh ) = 0. More precisely: • If fh (τ ⋆ ) < 0, then τh ∈ (τ ⋆ − h, τ ⋆ ) by the intermediate value theorem. • If fh (τ ⋆ ) = 0, then τh = τ ⋆ . In both cases, we have: τ ⋆ − h < τh ≤ τ ⋆ . Step 4: Error bound.

(28)

The absolute error is: |τ ⋆ − τh | = τ ⋆ − τh ≤ h =

1 . B

(29)

This completes the proof. The proof shows that the histogram approximation is conservative as it never overestimates τ ⋆ , which means τh ≤ τ ⋆ . This ensures that the approximate solution produces a sparsity pattern that preserves the support of the real solution—an useful property for safe block-masking. B.2. Proof of Proposition 2 We divided the proof of Proposition 2 into two parts: (i) continuous differentiability of the threshold map f , and (ii) continuous differentiability of the α-entmax transformation (from scores into probabilities). Throughout, we assume α > 1, and work with translated and scaled scores z = (α − 1)s − (m − 1)1 where m = (α − 1) max(s). The domain of possible such scores after translation is D = {z ∈ Rn : max(z) = 1}. Scaling and centering are continuous smooth operations and thus composing with them does not affect the continuity of any order of differentiation. We denote the active set induced by a threshold τ as Sτ := {i ∈ [n] : zi > τ }. Part (i): continuity of f .

Given the active set Sτ , the derivatives of f have the form: X 1 f (τ ) = −1 + (zi − τ ) α−1

(30)

i∈Sτ 1 1 X −1 (zi − τ ) α−1 α−1 i∈Sτ X 1 1 1 −2 f ′′ (τ ) = −1 (zi − τ ) α−1 α−1 α−1 i∈Sτ "k−1 # X Y 1 1 −k (k) k f (τ ) = (−1) −i (zi − τ ) α−1 α − 1 i=0

f ′ (τ ) = −

(31) (32)

(33)

i∈Sτ

We can study the continuity of any such derivative by considering the sum X 1 −k sk (τ ) := (zi − τ ) α−1 , i∈Sτ

(34)

A DA S PLASH -2: Faster Differentiable Sparse Attention

because every term can be written as a continuous transformation of sk . For a small enough ϵ > 0, Sτ +ϵ = Sτ . As a sum of continuous functions we therefore have lim sk (τ + ϵ) = sk (τ ). (35) ϵ→0+

If τ ̸= zi for any i ∈ [n] then Sτ −ϵ = Sτ also, so the only possible discontinuities are when τ = zi for some i. In this case, zi enters the support (and thus the sum) alongside any other tied zj = zi . Let’s say the value zi appears ni ≥ 1 times in the vector. We can write X 1 −k sk (τ − ϵ) = (zj − τ + ϵ) α−1 (36) j∈Sτ −ϵ

(zj − τ + ϵ) α−1

−k

+ ni (zi − τ + ϵ) α−1

−k

(37)

j∈Sτ

Taking the limit to 0, the sum is a sum of continuous functions and so we have 1

lim sk (τ − ϵ) = sk (τ ) + lim ni (zi − τ + ϵ) α−1 −k

ϵ→0+

= sk (τ ) + ni lim ϵ α−1 −k ϵ→0+

(38) (39)

It follows that sk is continuous at τ iff. (α − 1)−1 − k > 0. In particular: • f (τ ) is continuous when 1 < α. • f ′ (τ ) is continuous when 1 < α < 2. • f ′′ (τ ) is continuous when 1 < α < 3/2, and so on. Part (ii): continuity of α-entmax.

Let Fα (z, τ ) = −1 +

(zi − τ ) α−1 .

(40)

i∈Sτ

Recall that, by continuity and monotonicity of f , for any z there is a unique τ ∈ [0, 1] such that Fα (z, τ ) = 0; denote by g(z) the function mapping z to that corresponding unique root τ .6 If α ∈ (1, 2), from the continuity of f shown above, we have that Fα is C 1 . The condition that the jacobian be nonsingular in this case is equivalent to f ′ (τ ) ̸= 0 (Eq. 31). The terms in the sum sk are all positive. For τ ∈ [0, 1] the sum contains at least one term corresponding to the maximum (scaled) score and thus f ′ (τ ) < 0, confirming nonsingularity. The conditions of the implicit function theorem (Dontchev & Rockafellar, 2014, Thm. 1B.1) are therefore satisfied, which implies the solution mapping g is C 1 for any z ∈ D. Since α-entmax can be written coordinate-wise as 1 pj (z) = [zj − g(z)]+α−1 , (41) its (almost everywhere) gradient with respect to z is ∇pj (z) =

1 1 −1 [zj − g(z)]+α−1 (ej − ∇g(z)) , α−1

(42)

and ∇pj (z) = 0 whenever zj < g(z). For 1 < α < 2 this function is a composition/product of continuous functions, and thus it is continuous (Rudin, 1976, Theorems 4.7 & 4.9). Therefore, for this range of α, the “almost everywhere” above becomes “everywhere”. Furthermore, by the higher-order extension of the implicit function theorem (Dontchev & Rockafellar, 2014, Prop. 1B.5), k we can repeat this argument for higher-order derivatives. When α < k+1 k , from the result in the first part we have Fα ∈ C k This implies g is also C . Applying the same argument coordinate-wise, it follows that α-entmax is C k whenever 1 < α < k+1 k . Indeed, as α → 1+ we recover softmax, which is C ∞ . 6 At the root the active set is nonempty: if Sτ = ∅, then Fα (z, τ ) = −1 ̸= 0. Equivalently, since max(z) = 1, any root satisfies τ < 1 (because Fα (z, τ ) = −1 for τ ≥ 1).

A DA S PLASH -2: Faster Differentiable Sparse Attention

Bin 0 0

Bin 1 4

Bin 2 1

Bin 4 0

Bin 3 2

Bin 5 0

Bin 6 1

Bin 7 0

Bit representation: 0x0001000002010400

Figure 6. Example of a bitpacked histogram with B = 8 bins and b = 8 bits per bin. Each colored segment represents a bin’s count encoded in 8 bits of a uint64 integer.

C. A DA S PLASH -2 Implementation Details This appendix provides a detailed exposition of A DA S PLASH -2’s implementation, including the bitpacking schemes that enable efficient histogram construction and block mask traversal. We begin by establishing notation (§C.1), then describe the histogram construction and capacity analysis (§C.2), block mask encoding (§C.4), and overflow handling strategies (§C.5). C.1. Notation and Problem Setup Sequence and Tiling Parameters. We work with input sequences of length n, head dimension d, batch size B, and NH attention heads. The sequences are partitioned into tiles: query tiles of size Br and key/value tiles of size Bc , yielding Tr = ⌈n/Br ⌉ query tiles and Tc = ⌈n/Bc ⌉ key/value tiles. Throughout, we use indices i ∈ [Tr ] for query tiles and j ∈ [Tc ] for key tiles. Histogram Parameters. The histogram has shape H ∈ NBr ×B . To avoid atomics or reductions, we first accumulate the counts in local histograms Hlocal private to each position, providing Br × Bc parallel accumulators. Our local histograms use B bins encoded in a w-bit unsigned integer (e.g., w = 64 for uint64), allocating b = w/B bits per bin. The bin width (resolution) is h = 1/B. Bitwise Operations.

We use standard notation for bit manipulation:

• x ≪ k: Left shift by k bits (multiply by 2k ) • x ≫ k: Right shift by k bits (divide by 2k , rounded down) • x ∧ y: Bitwise AND • x ∨ y: Bitwise OR • popc(x): Population count (number of 1-bits) • fns(x, k): Find first set bit at position ≥ k C.2. Histogram Construction and Bitpacking The histogram construction in Algorithm 1 (Phase 2) is the computational core of A DA S PLASH -2’s threshold approximation. We now describe its implementation in detail. local Bitpacked Encoding Scheme. Each local histogram accumulator Hik ∈ {0, 1, . . . , 2w − 1} encodes B bin counts by local partitioning its w bits into B equal segments of b = w/B bits each. The value Hij can be decomposed as:

local Hij =

B−1 X

ck 2kb ,

where ck ∈ [0, 2b − 1] is the local counts for bin t.

t=0

Figure 6 illustrates this structure for B = 8 bins with b = 8 bits per bin in a uint64. 17

(43)

A DA S PLASH -2: Faster Differentiable Sparse Attention

Accumulation.

For each entry Zij (see Eq. 6), we compute the corresponding bin, kij = min max(⌊B Zij ⌋, 0), B − 1 .

(44)

Each update increments the respective bit field inside the packed integer: local local Hij ← Hij + 1[Zij ≥ 0] · 1 ≪ (b ki,j ) .

(45)

In Triton, this update can be achieved with a shift plus an integer addition. Extraction and Aggregation. After accumulation, we extract the count for bin k ∈ {0, . . . , B − 1} from the packed uint64 word by shifting and masking. Let Bb = 2b − 1 be the b-bit mask (i.e., b consecutive 1-bits). Then, we can aggregate across the Bc parallel accumulators (over j) to obtain per-query bin counts X local Hik = Hij ≫ (k b) ∧ Bb . (46) j

In Triton, we vectorize the shift-and-mask over b using a compile-time range. C.3. Histogram initialization With the final histogram H obtained, we can solve the approximate problem (see Eq. 9) using a modified version of the sorting-based algorithms discussed in Section 2.4. The τ0 and bracket [τlo , τhi ] obtained from this phase will be used for posterior refinement. Here, we present the algorithms in detail. Case α ∈ {1.5, 2.0}. As the counts are conveniently sorted in the histogram H, we scan bins in descending order, k ∈ [B − 1, .., 0]. We maintain prefix sums over the active (ordered) set P = [B − 1, B − 2, . . . , k]: 2 X X X k k S0 = Hk , S1 = Hk · , S2 = Hk · . (47) B B k∈P

k∈P

Intuitively, we start with an empty set where fh (τ ) = −1. Iteratively, we add contributions from the next bin (k = B − 1, then k = B − 2, ...) while evaluating f (τ = k/B). If adding the contribution of a bin, k, cause f (τ ) to be positive, it means that the root is between the values k/B < τ < (k + 1)/B. When that set Pk⋆ is found, we can drop the ReLU, [·]+ , and solve for τ . Example for α = 2.0: X k Hk −τ =1 B ⋆ k∈Pk

X S1 − 1 k −τ Hk = 1 =⇒ τ = B S0 k∈Pk⋆ k∈Pk⋆ | {z } | {z } X

General α. For a general α, we can find τ ⋆ by using any method originally developed for α-entmax. However, instead of iterating over the actual scores, we can cheaply iterate over the histogram H. C.4. Block Mask Encoding and Traversal After computing the threshold τ in Phase 3 of Algorithm 1, we construct a binary mask M ∈ {0, 1}Tr ×Tc indicating which tile pairs (i, j) contain non-zero attention weights. Naively storing this as a boolean tensor would require Tr × Tc bytes. Instead, we use bitpacking encoded in int32 to reduce memory by a factor of 32×. That is, we pack 32 consecutive column indices into a single int32, yielding a compact representation Mpacked ∈ ZTr ×⌈Tc /32⌉ defined by: Mpacked i,⌊j/32⌋ =

31 X

Mi,32⌊j/32⌋+k · 2k .

k=0

As a result, each bit in Mpacked i,m indicates whether the corresponding tile pair has non-zero weights. 18

(48)

A DA S PLASH -2: Faster Differentiable Sparse Attention Table 4. Histogram capacity analysis for different integer types and bin configurations with Bc = 64. Capacity indicates maximum keys per query tile; resolution h = 1/B is the bin width affecting threshold accuracy (Proposition 1). Type (w)

Bins (B)

Bits/Bin (b)

Max/Bin

Capacity

Resolution (h)

uint64 64 64 64

4 8 16

16 8 4

65,535 255 15

4,194,240 16,320 960

0.25 0.125 0.0625

uint128 128 128 128

8 16 32

16 8 4

65,535 255 15

4,194,240 16,320 960

0.125 0.0625 0.03125

Efficient Traversal with GPU Instructions.

To iterate through active tiles, we use two GPU-native instructions:

1. popc(Mpacked i,m ): Returns the number of set bits (active tiles) in the int32. 2. fns(Mpacked i,m , k): Find the n-th set bit given by offset k. The traversal procedure processes only the |M| non-zero tiles rather than all Tr × Tc potential tiles. Concretely, for each packed word index m, we iterate only over set bits (active key blocks) using popc to count and fns to locate bits. If bit pos is the position of a set bit within word m, the corresponding key-block index is j = 32m + bit pos.

(49)

The operations operations are implemented via Triton’s inline PTX assembly, making the traversal cost negligible compared to the memory and compute operations. This traversal mechanism is used in Phase 4 of Algorithm 1 to load only the non-zero key and value tiles, and similarly in the backward pass to skip zero blocks. C.5. Histogram Capacity and Overflow Handling The bitpacking scheme naturally imposes capacity limits for our histogram since each bin can count up to 2b − 1 items before overflow. With Bc parallel accumulators per query tile, the maximum capacity is: Cmax = Bc × (2w/B − 1).

(50)

Table 4 summarizes capacity and resolution trade-offs for various configurations. The table reveals a fundamental trade-off: increasing B improves threshold resolution (tighter bound in Proposition 1) but reduces capacity per bin. For sequences up to 16K tokens per query, uint64 with B = 8 provides excellent resolution (h = 0.125) while avoiding overflow. For longer sequences, we turn to an overflow mitigation strategy, which we describe next.7 Periodic Flush. After processing Cmax keys, we extract all bin counts, accumulate them into shared memory array of size (Tr × B), and reset the SRAM histogram to zero. Since each query processes Tc key tiles and the histogram saturates after 2b − 1 tiles, the number of flushes required for a sequence of length n is: n Tc Nflush = = b . (51) Bc × (2b − 1) 2 −1 C.6. Backward Pass Algorithms In the backward pass, we exploit the block mask M to skip zero gradients, achieving asymptotic speedups when sparsity is high. The key difference from standard attention backpropagation is the use of the block mask M (and its transpose M⊤ ) to iterate only over non-zero blocks. For |M| ≪ Tr · Tc (high sparsity), this provides order-of-magnitude speedups in the backward pass, offsetting the higher forward pass cost of A DA S PLASH -2 relative to FlashAttention-2. 7 We note that uint128 is already available in CUDA: https://developer.nvidia.com/blog/ implementing-high-precision-decimal-arithmetic-with-cuda-int128/

A DA S PLASH -2: Faster Differentiable Sparse Attention Table 5. Downstream evaluations on short-context benchmarks. Models are trained with 50B tokens from DCLM-Edu and have 4K context length. Those marked with (32K) undergo context-length extension with 10B tokens of Prolong data. Best results are in bold. Model

LMB (ppl) LMB ARC-E ARC-C CSQA HS OBQA PIQA SocialQA WG Avg.

350M params. Softmax with RoPE Entmax with RoPE Softmax with NAPE Entmax with NAPE

23.93 22.36 19.23 18.62

40.6 39.5 41.2 42.4

61.9 63.0 61.8 61.9

32.1 34.8 33.4 33.0

50.5 50.5 47.6 51.2

41.6 40.2 40.5 41.0

38.4 39.0 38.0 39.2

66.6 64.8 64.8 66.0

45.1 44.1 43.5 44.4

48.9 53.4 52.7 53.7

47.3 47.7 47.1 48.1

1B params. Softmax with RoPE Entmax with RoPE Softmax with NAPE Entmax with NAPE

15.01 15.76 11.97 11.42

44.7 43.9 48.0 49.2

69.0 65.4 69.3 67.7

36.0 36.6 37.4 39.9

58.7 56.2 56.7 57.3

49.6 47.4 48.8 48.7

42.8 41.6 45.2 45.0

70.0 69.8 69.5 68.3

46.4 45.7 46.5 47.1

55.6 54.2 55.2 55.1

52.5 51.2 53.0 53.1

1B params. Softmax with RoPE (32K) Entmax with RoPE (32K) Softmax with NAPE (32K) Entmax with NAPE (32K)

14.95 12.82 12.04 13.01

46.8 48.2 50.7 49.0

61.9 63.7 62.4 64.1

34.0 34.1 33.6 34.2

50.0 50.2 49.8 49.2

44.4 45.0 45.4 45.8

42.0 40.4 41.0 41.8

68.0 67.1 68.2 67.6

45.1 45.8 46.5 46.9

53.2 52.5 51.9 52.9

49.5 49.7 49.9 50.2

D. Language Modeling We train our language models with the torchtitan library (Liang et al., 2025) and in a single node with 4x H100 NVIDIA GPUs. The models are trained for 50B tokens of DCLM-Edu data with the WSD scheduler (2000 warmup steps, stable until 80% of training and decay for the latter 10B tokens). Weight decay is set to 0.1 and maximum learning rate is set to 3 × 10−4 and 2 × 10−4 for 350M and 1B parameter models, respectively. Final learning rate is set to 10% of the maximum. Context length is set to 4096 tokens with an effective batch size of 524k tokens, resulting in 100k total steps. We decide to go beyond the Chinchilla (Hoffmann et al., 2022) optimal token counts since we saw steady increased model performance even with the additional tokens. Finally, for inference we use a simplified variant of A DA S PLASH -2 that operates on fully materialized logits z ∈ Rn . Despite its simplicity, in practice we found this approach to be faster than parallelizing over the K/V sequence dimension for the sequence lengths considered in this work. Context-length extension. For the long-context extension phase we swap the data mixture at the decay phase (10B tokens) with the ProLong dataset (Gao et al., 2025), keeping total training data at 50B tokens. Here, we increase the sequence length to 32k and keep the effective batch size fixed. For RoPE-based models, we increase their θ from the 50k default to 800k, consistent with the ProLong methodology. Evaluation. For long-context we adopt both the full RULER benchmark (Hsieh et al., 2024) and the In-Context Learning subset of HELMET (Yen et al., 2025) which is a capability not covered by RULER. For evaluation in short-context tasks, we also include Lambada (Paperno et al., 2016; Radford et al., 2019) for evaluating perplexity, and drop BoolQ (Clark et al., 2019) from the evaluation mix following Olmo3 (Olmo et al., 2025). For completeness, we provide the results of all of our models on short-context benchmarks in Table 5. The results on RULER for the long-context adapted models are presented in Table 1, while the HELMET-ICL results can be seen in Table 2. Discussion. Overall, we see that extended models show a small degradation on short-context benchmarks, which is typical of long-context extension (Gao et al., 2025). However, the scenario flips when we look at Table 1, where we observe that our long-context procedure was effective and 32K-context models are able to complete tasks at 32K context length to an effective degree. Analysis of In-Context Learning capabilities (Table 2) between dense and sparse models further confirm the strengths of A DA S PLASH -2 as it outperforms softmax models across all evaluated sequence lengths. D.1. Dense to Sparse Attention Conversion Having established the efficacy and efficiency of sparse attention, we make a preliminary study of how to convert a pretrained softmax model into a α-entmax one. We do so via continued pretraining: starting from a 1B-parameter checkpoint trained with softmax + NAPE at 4K context length, we replace softmax with α-entmax (implemented with A DA S PLASH -2) and 20

A DA S PLASH -2: Faster Differentiable Sparse Attention Table 6. Downstream evaluations on short-context benchmarks Models are trained with 50B tokens from DCLM-Edu and have 4K context length. Entmax (converted) implies the softmax-based model was evaluated with A DA S PLASH -2 without further training, while Entmax (CPT) is the converted model subjected to Continuous Pre-Training (CPT). Model (1B params.) Tokens LMB ARC-E ARC-C CSQA

OBQA PIQA SocialQA WG Avg.

50B token baselines Softmax (scratch) Entmax (scratch)

50B 50B

48.0 49.2

69.3 67.7

37.4 39.9

56.7 57.3

48.8 48.7

45.2 45.0

69.5 68.3

46.5 47.1

55.2 53.0 55.1 53.1

softmax → 1.5-entmax Entmax (converted) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT) Entmax (CPT)

50B 51B 52B 53B 54B 55B 56B 57B 58B 59B 60B

31.3 46.6 46.9 46.3 46.4 47.5 48.1 47.3 47.6 48.0 47.4

42.3 69.2 68.8 69.7 69.7 69.6 68.9 69.0 69.1 69.7 68.8

31.8 37.5 37.7 37.6 38.8 40.0 39.7 39.3 39.6 38.9 39.7

47.3 56.8 57.0 57.6 56.7 57.6 56.8 57.4 57.8 58.3 57.7

36.6 48.3 49.2 48.8 50.6 50.6 50.7 50.3 49.8 50.3 50.6

36.0 47.2 45.8 45.4 46.0 45.8 46.2 46.8 46.4 47.6 45.4

63.0 69.7 69.8 70.2 70.3 69.9 69.8 70.0 70.6 69.6 69.6

41.0 47.4 46.7 46.6 46.9 46.8 47.9 46.6 46.9 46.5 47.1

50.8 53.5 55.2 55.2 54.7 55.6 55.2 55.4 55.1 54.1 55.0

60B token baselines Softmax (scratch) Entmax (scratch)

60B 60B

48.5 47.8

68.9 69.5

37.1 41.6

58.4 58.8

51.2 50.7

46.6 47.2

70.0 68.7

47.4 48.7

55.2 53.7 57.0 54.4

42.2 52.9 53.0 53.0 53.3 53.7 53.7 53.6 53.7 53.7 53.5

continue training for 10B tokens on Nemotron-CC-V2 (Basant et al., 2025), keeping the final learning rate unchanged. We focus only on 1B NAPE model variants due to the seemingly sub-optimality of RoPE for α-entmax attention. Table 6 shows the full results. Directly replacing softmax with α-entmax in the attention modules without any additional training results in a performance degradation but not a total collapse (see “Entmax (converted)” entry). After only 2B tokens, at 52B tokens in total, the converted checkpoint already outperforms the baseline softmax scores on average. Finally, we continue training the baselines to measure the total performance lost due to the conversion, evaluating after successive increases of 1B tokens. We observe that the finals results are roughly on par with the results obtained by softmax and α-entmax models trained from scratch on the full 60B tokens. Overall, this strategy offers insight into a efficient recipe for building entmax-based models from existing softmax + NAPE checkpoints, which is particularly appealing for midtraining or long-context extension phases. Applying this methodology to other positional encodings or architectures could, additionally, require other adaptation phases (Gelberg et al., 2026).

E. Attention Sparsity and Efficiency Sparse attention patterns. In Figure 7, we analyze emerging attention block sparsity in the 1B Entmax (NAPE) model trained up to 32K context length. We compute the average 64 × 64 attention block sparsity ratios across 64 sequences sampled from the ProLong dataset, evaluated at multiple context lengths (4K, 8K, 16K, 32K, 64K and 128K). Across all context lengths, the sparsity pattern is strongly head-dependent. The first half of the heads (0-11; ALiBi) are consistently sparse across layers. In contrast, the second half (12-23; NoPE) is comparatively dense at 4K, but develops structured sparsity as context grows, most notably in mid-to-late layers. The overall mean sparsity increases monotonically with context length, indicating that longer contexts are handled with less dense attention on average.

Context scaling experiment details. We study the runtime behavior of A DA S PLASH -2 as context length increases using attention block sparsity patterns extracted from the 1B Entmax + NAPE model with 32k context length, reported in Table 1. To this end, we select a fixed transformer layer that exhibits high block sparsity and record its average block sparsity as a function of context length. For each context length, we benchmark A DA S PLASH -2 at the corresponding sparsity level, run forward and backward passes, and report average runtimes over repeated runs. Figure 4 reports the resulting average forward and backward runtimes as a function of context length. 21

Context 4k / Overall Mean: 0.48

Context 8k / Overall Mean: 0.52

Context 16k / Overall Mean: 0.56

1.0

A DA S PLASH -2: Faster Differentiable Sparse Attention

Context 32k / Overall Mean: 0.60

Context 64k / Overall Mean: 0.64

Context 128k / Overall Mean: 0.78 0.4

Layers

0.6

64x64 Block Sparsity Ratio

Layers

0.8

0.2

Heads

0.0

Figure 7. Average 64 × 64 attention block sparsity ratio for the Entmax (NAPE) model with 32K context length. Panels correspond to evaluated context lengths (4K / 8K / 16K / 32K / 64K / 128K). We report the overall average sparsity across all layers and heads in the title on each plot.

F. Full Algorithms We provide the full forward pass pseudo-code of A DA S PLASH -2 in Algorithm 1, and pseudo-code for our two backward kernels in Algorithm 2 and 3.

A DA S PLASH -2: Faster Differentiable Sparse Attention

Algorithm 1 A DA S PLASH -2 Forward Pass Require: Matrices Q, K, V ∈ Rn×d in HBM, parameter α > 1, bin count B, block sizes (Br , Bc ). Ensure: Output O ∈ Rn×d in HBM, threshold vector τ ∈ Rn in HBM, block mask M ∈ Z⌈n/Br ⌉×⌈n/Bc ⌉ in HBM. 1: Let Tr = ⌈n/Br ⌉ and Tc = ⌈n/Bc ⌉. 2: Initialize O ← 0 ∈ Rn×d , τ ← 0 ∈ Rn , and M ← 0 ∈ ZTr ×Tc in HBM. 3: Divide Q, O into Tr blocks Qi , Oi ∈ RBr ×d , and divide K, V into Tc blocks Kj , Vj ∈ RBc ×d . Divide τ into Tr blocks τi ∈ RBr . 4: for i = 1 to Tr do 5: Load Qi from HBM to on-chip SRAM. 6: // Phase 1: Compute row-wise maximum 7: Initialize mi ← −∞ ∈ RBr on SRAM. 8: for j = 1 to Tc do 9: Load Kj from HBM to SRAM. (j) 10: Compute Si ← Qi Kj⊤ ∈ RBr ×Bc . (j) 11: Update mi ← max mi , maxcol-wise (Si ) . 12: end for 13: 14: 15: 16: 17: 18: 19: 20: 21: 22: 23: 24: 25: 26: 27: 28: 29: 30: 31: 32: 33: 34: 35: 36: 37: 38: 39: 40:

// Phase 2: Build histogram in a single pass Initialize local histogram Hlocal ← 0 ∈ RBr on SRAM. ij for j = 1 to Tc do Load Kj from HBM to SRAM. (j) Compute Si ← Qi Kj⊤ ∈ RBr ×Bc . (j)

Update Hlocal ij with quantized scores from Zi . end for Br ×B Reduce Hlocal . ij into Hi ∈ R τi ← SolveHistogram(Hi , α) ∈ RBr

▷ See §C.3

// Phase 3: Single hybrid step and save block mask Initialize accumulators for f (τ0,i ), f ′ (τ0,i ), f ′′ (τ0,i ) (row-wise). for j = 1 to Tc do Load Kj from HBM to SRAM. (j) Compute Si ← Qi Kj⊤ ∈ RBr ×Bc . 1/α−1 (j) ∈ RBr ×Bc . Compute Pi (j) ← max 0, (α−1)Si − τ0,i Accumulate f (τi ), f ′ (τi ), f ′′ (τi ) using Pi (j) . if any(Pi (j) > 0) then Mij ← 1. end if end for τi ← HybridSolver(f, f ′ , f ′′ ) ∈ RBr

▷ See §3.2

// Phase 4: Compute output for nonzero blocks Initialize Oi ← 0Br ×d on SRAM. for j : Mij = 1 do Load Kj , Vj from HBM to SRAM. (j) 41: Compute Si ← Qi Kj⊤ ∈ RBr ×Bc . 1/α−1 (j) (j) 42: Compute Pi ← max 0, (α−1)Si − τi ∈ RBr ×Bc . (j) 43: Oi ← Oi + Pi Vj ∈ RBr ×d . 44: end for 45: Write Oi to HBM. 46: Write τi to HBM. 47: end for 48: Return: O (and saved τ , M). 23

A DA S PLASH -2: Faster Differentiable Sparse Attention

Algorithm 2 A DA S PLASH -2 Backward Pass for dK and dV r ⌉×⌈n/Bc ⌉ Require: Matrices Q, K, V , O, dO ∈ Rn×d and binary mask M ∈ Z⌈n/B in HBM, vector τ ∈ Rn in HBM, Pn 2−α block sizes Bc , Br , parameter α. Assume previously computed δi = j=1 Sij Vj /∥Ui ∥1 . 1: Divide Q into Tr = ⌈n/Br ⌉ blocks Q1 , . . . , QTr of size Br × d each, and divide K, V into Tc = ⌈n/Bc ⌉ blocks K1 , . . . , KTc , V1 , . . . , VTc of size Bc × d each. 2: Divide dO into Tr blocks dO1 , . . . , dOTr of size Br × d each. 3: Divide τ into Tr blocks τ1 , . . . , τTr of size Br each. 4: Initialize and divide dK, dV ∈ Rn×d into Tc blocks dK1 , . . . , dKTc and dV1 , . . . , dVTc of size Bc × d each. 5: Divide δ into Tr blocks δ1 , . . . , δTr of size Br . 6: for 1 ≤ j ≤ Tc do 7: Load Kj , Vj from HBM to on-chip SRAM. 8: Initialize dKj = 0Bc ×d on SRAM. 9: Initialize dVj = 0Bc ×d on SRAM. 10: for i : Mij = 1 do 11: Load Qi , dOi , τi , δi from HBM to on-chip SRAM. (j) 12: On chip, compute Si = Qi Kj⊤ ∈ RBr ×Bc .

13: 14: 15:

(j)

On chip, compute Pi = max(0, (α − 1)Si − τi )1/α−1 ∈ RBr ×Bc . (j) On chip, compute dVj ← dVj + (Pi )⊤ dOi ∈ RBc ×d . ⊤ On chip, compute dPi = dOi Vj ∈ RBr ×Bc . (j)

(j) 2−α

On chip, compute Ui = Pi ∈ RBr ×Bc . (j) (j) (j) 17: On chip, compute dSi = Ui ⊙ (dPi − δi ) ∈ RBr ×Bc . (j) 18: On chip, compute dKj ← dKj + (dSi )⊤ Qi ∈ RBc ×d . 19: end for 20: Write dKj , dVj to HBM. 21: end for 22: Return: Gradients dK, dV . 16:

Algorithm 3 A DA S PLASH -2 Backward Pass for dQ r ⌉×⌈n/Bc ⌉ Require: Matrices Q, K, V , O, dO ∈ Rn×d and binary mask M ∈ Z⌈n/B in HBM, vector τ ∈ Rn in HBM, Pn 2−α block sizes Bc , Br , parameter α. Assume previously computed δi = j=1 Sij Vj /∥Ui ∥1 . 1: Divide Q into Tr = ⌈n/Br ⌉ blocks Q1 , . . . , QTr of size Br × d each, and divide K, V into Tc = ⌈n/Bc ⌉ blocks K1 , . . . , KTc , V1 , . . . , VTc of size Bc × d each. 2: Divide dO into Tr blocks dO1 , . . . , dOTr of size Br × d each. 3: Divide τ into Tr blocks τ1 , . . . , τTr of size Br each. 4: Initialize dQ in HBM and divide it into Tr blocks dQ1 , . . . , dQTr of size Br × d each. 5: Divide δ into Tr blocks δ1 , . . . , δTr of size Br each. 6: for i = 1 to Tr do 7: Load Qi , dOi , δi , τi , from HBM to on-chip SRAM 8: Initialize dQi = 0Bc ×d on SRAM. 9: for j : Mij = 1 do (j) 10: On chip, compute Si = Qi Kj⊤ ∈ RBr ×Bc .

11: 12:

(j)

On chip, compute Pi = max(0, (α − 1)Si − τi )1/α−1 ∈ RBr ×Bc . On chip, compute dPi = dOi Vj⊤ ∈ RBr ×Bc . (j)

(j) 2−α

On chip, compute Ui = Pi ∈ RBr ×Bc . (j) (j) (j) 14: On chip, compute dSi = Ui ⊙ (dPi − δi ) ∈ RBr ×Bc . (j) 15: On chip, compute dQi ← dQi + dSi Kj ∈ RBr ×d . 16: end for 17: Write dQi to HBM 18: end for 19: Return: Gradient dQ 13:

Record · ID 19009 · SHA-256 887001b309451968
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.