ICML 2025

ADASPLASH: Adaptive Sparse Flash Attention
Nuno Gonc¸alves 1 Marcos Treviso 2 Andr´e F. T. Martins 1 2 3
Abstract
The computational cost of softmax-based atten-
tion in transformers limits their applicability to
long-context tasks. Adaptive sparsity, of which
α-entmax attention is an example, offers a flexi-
ble data-dependent alternative, but existing 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 α-entmax. We first
introduce a hybrid Halley-bisection algorithm, re-
sulting in a 7-fold reduction in the number of itera-
tions needed to compute the α-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
our method achieves substantial improvements in
runtime and memory efficiency compared to ex-
isting α-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.1
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-
abilities 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 to dispersion (Veliˇckovi´c et al., 2025).
1Instituto Superior T´ecnico, Universidade de Lisboa, Portu-
gal 2Instituto de Telecomunicac¸˜oes, Lisbon, Portugal 3Unbabel,
Lisbon, Portugal.
Correspondence to:
Nuno Gonc¸alves
<nuno.m.goncalves@tecnico.ulisboa.pt>.
Proceedings of the 42 nd ...