AdaSplash: Adaptive Sparse Flash Attention

AdaSplash: Adaptive Sparse Flash Attention

Speakers: Nuno Gonçalves, Marcos V. Treviso, Andre Martins

Conference: ICML 2025

YouTube: https://slideslive.com/39043989

Overview

The talk "AdaSplash: Adaptive Sparse Flash Attention" introduces a groundbreaking approach to enhance the efficiency and scalability of attention mechanisms in large-scale transformer models. Presented by Nuno Gonçalves, Marcos V. Treviso, and Andre Martins from IST Lisbon, this work directly addresses the fundamental quadratic memory complexity inherent in traditional attention, a bottleneck that has historically limited the context lengths and overall scalability of transformers. While Flash Attention revolutionized softmax-based attention by eliminating the need to materialize large intermediate matrices, it left a critical gap for alternative, sparse attention mechanisms.

AdaSplash tackles this challenge head-on by providing a highly optimized, fused GPU kernel for alpha-Entmax attention, a sparse alternative to softmax. The core innovation lies in a novel Halley-bisection algorithm that dramatically accelerates the computation of the Entmax transformation itself, a process previously plagued by slow, GPU-unfriendly algorithms. By integrating this efficient Entmax computation into a Flash Attention-like architecture, AdaSplash not only makes sparse attention practical for long sequences but also demonstrates its empirical superiority over softmax attention in various downstream tasks, marking a significant step towards more efficient and focused transformer architectures.

This research is particularly pertinent in an era dominated by large language models, where context window limitations and computational costs remain major hurdles. By enabling adaptive, input-dependent sparsity, AdaSplash allows models to better focus on relevant tokens, potentially leading to improved performance, reduced memory footprint, and enhanced scalability for future generations of AI models. It offers a compelling alternative to the ubiquitous softmax, opening new avenues for designing more efficient and context-aware transformer architectures.

Background

▶ Watch: Introduction, attention's quadratic cost, and Flash Attention (0:00)

The attention mechanism is a cornerstone of modern transformer architectures, endowing them with their remarkable expressive power. At its essence, attention maps raw scores, derived from the dot product of queries and keys, to a probability distribution used to compute a weighted sum of values. However, this process traditionally involves the explicit materialization of two large matrices: the score matrix S and the probability matrix P. The memory consumption for these matrices scales quadratically with the sequence length, quickly becoming a prohibitive bottleneck for models processing long contexts.

In 2022, the introduction of Flash Attention by Dao et al. marked a pivotal advancement. This technique ingeniously re-engineered the attention computation to produce the exact softmax output without ever materializing the S and P matrices in global memory. By employing tiling and recomputing intermediate values, Flash Attention effectively reduced the memory complexity from quadratic to linear, enabling transformers to scale to significantly longer sequences while maintaining computational efficiency.

Despite its success, Flash Attention was designed specifically for the softmax function. The softmax, by its nature, assigns a non-zero probability to every token, even those entirely irrelevant to the current context. This can lead to attention dispersion, where the probability distribution becomes overly spread out as sequence length increases, hindering the model's ability to sharply focus on critical tokens. Recent studies have highlighted this as a fundamental limitation of softmax-based attention.

A promising alternative is to replace softmax with a transformation that naturally yields sparse, input-dependent probability distributions. Entmax is one such family of transformations. It introduces a parameter alpha that allows for explicit control over sparsity, interpolating between the dense softmax (alpha=1) and increasingly sparser distributions, notably recovering sparsemax when alpha equals two. The key advantage of Entmax is its ability to concentrate probability mass on important tokens, effectively ignoring irrelevant ones, which is a desirable property for long-context understanding.

However, the practical adoption of Entmax attention has been severely hampered by two core inefficiencies. Firstly, existing algorithms for computing the alpha-Entmax transformation are notoriously slow and poorly optimized for GPU execution. Secondly, even if the Entmax transformation itself were efficient, its integration into the attention mechanism still incurred the same quadratic memory cost that Flash Attention had already eliminated for softmax. This work directly addresses these limitations, seeking to unlock the full potential of alpha-Entmax attention in large-scale transformer models.

Key Findings

▶ Watch: AdaSplash's three core contributions to Entmax attention (2:30)

The "AdaSplash" project delivers several critical findings and contributions that significantly advance the state of the art for sparse attention mechanisms:

1. Novel and Faster Entmax Algorithm: The work introduces a new algorithm for computing the alpha-Entmax transformation that achieves dramatically faster empirical convergence compared to all prior methods. This Halley-bisection method reduces the required number of iterations by approximately 7x and translates to a roughly 15x speedup in the Entmax computation itself.

1. AdaSplash: Fused GPU Kernel for Entmax Attention: A highly optimized, fused GPU kernel, dubbed AdaSplash, is developed specifically for alpha-Entmax attention. This kernel finally enables models trained with Entmax attention to scale efficiently for longer contexts, eliminating the quadratic memory bottleneck that previously limited its practical use. It computes the exact output of Entmax attention without materializing the full score or probability matrices.

1. Competitive Performance with Flash Attention 2: AdaSplash demonstrates competitive performance against Flash Attention 2, and can even surpass it when sufficient sparsity is present. This is particularly relevant for longer sequences where Entmax's input-dependent sparsity naturally becomes more pronounced.

1. Empirical Superiority on Downstream Tasks: Entmax attention, powered by AdaSplash, is empirically shown to be a useful and potentially more efficient alternative to conventional softmax attention. It outperforms softmax-based attention across several downstream tasks, notably achieving significant gains in retrieval tasks (e.g., with ModernBERT) and slightly lower validation loss in language modeling (e.g., with GPT-2).

1. High Sparsity Without Performance Degradation: Models trained with Entmax attention exhibit extremely high levels of sparsity, sometimes as low as 1% non-zeros in their attention patterns. Crucially, this high sparsity does not lead to degraded performance; in fact, these models frequently outperform their softmax counterparts.

1. Sub-Quadratic Backward Pass Potential: The inherent sparsity of Entmax attention, combined with the optimized implementation, technically allows for a sub-quadratic backward pass, a significant advantage that softmax-based attention mechanisms do not inherently possess.

Technical Deep Dive

▶ Watch: Halley-Bisection: combining speed with bisection's convergence guarantee (4:10)

The technical innovations in AdaSplash primarily revolve around two areas: an accelerated algorithm for the alpha-Entmax transformation and a fused GPU kernel for its efficient integration into the attention mechanism.

Accelerated Alpha-Entmax Transformation

The core challenge in computing the alpha-Entmax transformation lies in finding a scalar tau such that the resulting probability vector sums to exactly one. This can be formulated as a root-finding problem for a function f(tau) = 0.

Traditionally, the bisection algorithm has been used for this purpose. While bisection guarantees convergence, it is slow, typically requiring over 20 iterations to achieve a good estimate of tau. This iterative process is highly inefficient for GPU execution, as it necessitates loading input data from global memory or L2 cache multiple times, making it very memory-bound.

The authors observed that the derivatives of the function f(tau) are relatively easy to calculate, simply being shifted powers of the original function times a constant. This insight prompted the exploration of Householder methods, a family of root-finding algorithms that use higher-order derivatives to achieve faster convergence rates. Specifically, for m=2, Halley's method offers cubic convergence, a substantial improvement over bisection's linear convergence. However, a known drawback of Householder methods, including Halley's, is their potential to diverge if the initial estimate is too far from the true root.

To overcome this, AdaSplash introduces a novel Halley-bisection method. This hybrid approach combines the rapid, cubic convergence of Halley's method with the guaranteed convergence of bisection. Similar to bisection, the algorithm maintains lower and upper bounds that are guaranteed to contain the solution. At each step, it attempts a Halley update. If the Halley update produces a new estimate within the current bounds, it is accepted, leveraging its fast convergence. If the Halley update falls outside the bounds (indicating potential divergence), the algorithm defaults to a bisection step, ensuring progress towards the root and maintaining the convergence guarantee. This strategy allows AdaSplash to achieve machine precision for tau in just two to three iterations, representing a 7x reduction in iteration count and translating to a 15x speedup in the Entmax computation over raw bisection.

AdaSplash: Fused GPU Kernel for Attention

The implementation of AdaSplash for attention closely follows the principles established by Flash Attention, adapting them for the Entmax transformation and its inherent sparsity. The goal is to compute the attention output O = P * V without materializing the full S and P matrices.

1. Tiling and Parallelization: Similar to Flash Attention, the large query, key, and value matrices (Q, K, V) are partitioned into smaller tiles that can fit into fast shared memory on the GPU. The computation is parallelized across the query dimension, assigning each thread block a specific query tile (Q_tile). This Q_tile is loaded once into shared memory and remains there throughout the execution for that thread block.

1. Iterative Entmax Computation: For each Q_tile, the algorithm streams consecutive K_tiles from global memory. For each K_tile, it computes the partial score S_tile = Q_tile * K_tile^T and importantly, the first two derivatives required for the Halley-bisection method. This process is repeated for all K_tiles to accumulate the necessary quantities (sums of S_tile and its derivatives row-wise) for the Halley-bisection algorithm. This iterative procedure is run two to three times to accurately estimate the optimal tau for the current Q_tile.

1. Final Output Accumulation: Once tau is determined for the Q_tile, the algorithm performs a second pass. It again streams K_tiles and their corresponding V_tiles. For each K_tile, it computes S_tile, then the probability tile P_tile using the calculated tau and the Entmax transformation. Finally, it accumulates the contribution P_tile * V_tile into the output O_tile. This process is repeated for all K_tiles, ensuring that the final O_tile contains the complete attention output for the given Q_tile.

This fused kernel design ensures that the forward pass of AdaSplash achieves linear memory complexity with respect to sequence length, effectively mirroring Flash Attention's memory efficiency but for Entmax.

Exploiting Sparsity

While Entmax naturally yields sparse distributions, this sparsity is often unstructured, making direct exploitation challenging. AdaSplash addresses this by further partitioning the probability tiles into even smaller micro-tiles, suitable for tensor core operations. For these micro-tiles, the system maintains a lookup table of indexes corresponding to non-zero entries. If a micro-tile is entirely zero, the system can skip loading its respective K and V tiles and bypass the corresponding computations, leading to further efficiency gains. This intelligent sparsity exploitation also technically enables a sub-quadratic backward pass, a significant theoretical advantage over dense softmax attention.

Experimental Setup & Results

▶ Watch: AdaSplash fused kernel: integrating Halley-Bisection into attention mechanism (5:00)

The evaluation of AdaSplash focused on comparing its performance against existing methods, particularly Flash Attention 2, and assessing its impact on downstream task performance.

Performance Benchmarks

• Memory Efficiency & Scalability: When comparing against an eager (non-fused) implementation of Entmax attention, AdaSplash demonstrates vastly superior memory efficiency. The eager version quickly runs out of memory as context length increases, whereas AdaSplash scales to much larger contexts without issues, confirming its ability to eliminate the quadratic memory bottleneck.
• Runtime Comparison with Flash Attention 2: AdaSplash exhibits competitive performance against Flash Attention 2. In scenarios with sufficient sparsity, which naturally occurs more frequently in longer sequences, AdaSplash can even surpass Flash Attention 2 in terms of runtime. Generally, in terms of overall runtime and memory requirements, AdaSplash closely matches the highly optimized implementation of Flash Attention 2. This suggests that the benefits of Entmax (sparsity, improved focus) come at virtually no additional computational cost once optimized with AdaSplash.

Downstream Task Performance

The impact of AdaSplash was evaluated on key NLP tasks, demonstrating that the theoretical benefits of sparse attention translate into tangible performance improvements.

• Retrieval Tasks: For retrieval tasks, an encoder-only model, ModernBERT, was trained. AdaSplash, particularly with alpha=2 (which corresponds to sparsemax and promotes high sparsity), significantly outperformed softmax attention by a considerable margin. This highlights Entmax's ability to better focus on relevant information for accurate retrieval.
• Language Modeling: In language modeling experiments using a GPT-2 architecture, AdaSplash-enabled Entmax attention consistently achieved lower validation loss across all training steps compared to the standard GPT-2 using softmax attention. This improvement was observed even with relatively small context lengths, suggesting a general benefit of the focused attention mechanism.

Sparsity Analysis

A crucial finding was the observation of extremely high sparsity levels in models trained with Entmax attention. The attention patterns exhibited as low as 1% non-zeros. This level of sparsity is remarkable, especially given that it does not lead to degraded performance. In fact, as demonstrated by the downstream task results, these highly sparse models frequently outperform their denser softmax-trained counterparts. This validates the hypothesis that distributing probability mass across irrelevant tokens is detrimental and that a focused, sparse attention mechanism can be more effective.

Backward Pass

While not exhaustively benchmarked, the authors note that the design of AdaSplash, with its ability to skip computations for all-zero tiles, technically allows for a sub-quadratic backward pass. This is a significant theoretical advantage over softmax-based attention, which typically incurs quadratic cost in the backward pass due to the dense nature of its gradients.

Practical Implications

▶ Watch: Achieving linear memory complexity for Entmax attention forward pass (6:10)

AdaSplash represents a significant leap forward for the practical adoption of sparse attention mechanisms in machine learning, offering profound implications for various stakeholders:

• For Practitioners and Model Builders: AdaSplash makes alpha-Entmax attention a viable and scalable alternative to the ubiquitous softmax. This means practitioners can now experiment with and deploy models that intrinsically promote focused, input-dependent sparsity, potentially leading to better performance on tasks requiring fine-grained attention (e.g., long-document understanding, retrieval, question answering). The empirical gains demonstrated in retrieval and language modeling suggest that Entmax attention can be a powerful tool for building more accurate and efficient transformer models. It frees model builders from the limitations of softmax's "attention dispersion" problem.

• For Infrastructure Teams and Deployers: The development of the fused GPU kernel means that Entmax attention can now be deployed at scale without incurring prohibitive memory or computational costs. AdaSplash's performance matches or even surpasses Flash Attention 2 under favorable conditions (high sparsity), indicating that the transition to sparse attention doesn't necessitate a performance penalty. This enables infrastructure teams to support models with longer context windows more efficiently, potentially reducing GPU memory requirements and increasing throughput in production environments. The possibility of a sub-quadratic backward pass also hints at future optimizations for training large sparse models.

• Tradeoffs and Limitations: While highly promising, AdaSplash does come with considerations. The performance gains over Flash Attention 2 are often contingent on the degree of sparsity, which is inherently input-dependent. For inputs where Entmax produces a dense distribution (e.g., alpha close to 1, or specific data characteristics), the performance might align more closely with Flash Attention 2 without significant speedups. Furthermore, integrating AdaSplash requires using its specific optimized kernel; it's not a drop-in replacement for existing softmax implementations without code changes. The question posed during the Q&A regarding simple thresholding of softmax outputs highlights a key distinction: Entmax's sparsity is mathematically derived from its properties, offering a principled approach, whereas arbitrary thresholding of softmax would require re-normalization and is not easily supported by existing optimized kernels like Flash Attention. AdaSplash provides a robust, principled solution for adaptive sparsity.

In essence, AdaSplash unlocks the potential of sparse attention, transforming it from a theoretically appealing but practically challenging concept into a deployable, high-performance solution that can drive the next generation of efficient and intelligent transformer models.

Key Takeaways

• AdaSplash enables scalable alpha-Entmax attention: It provides a highly optimized, fused GPU kernel that finally allows sparse, input-dependent Entmax attention to scale to long contexts in large transformer models, overcoming previous quadratic memory and computational limitations.
• Novel Halley-bisection algorithm for Entmax: A new hybrid root-finding method dramatically accelerates the core Entmax transformation, achieving machine precision in 2-3 iterations (7x fewer iterations, ~15x speedup) compared to traditional bisection.
• Competitive performance with Flash Attention 2: AdaSplash matches Flash Attention 2's efficiency for the forward pass and can even surpass it when high input-dependent sparsity is present, particularly in longer sequences.
• Empirical superiority on downstream tasks: Models trained with AdaSplash-enabled Entmax attention consistently outperform softmax-based attention on tasks like retrieval (e.g., ModernBERT) and language modeling (e.g., GPT-2), demonstrating improved focus and accuracy.
• High sparsity without performance degradation: Entmax attention leads to extremely high sparsity (as low as 1% non-zeros) in attention patterns, yet these models perform equally well or even better than dense softmax models.
• Potential for sub-quadratic backward pass: The inherent sparsity of Entmax, combined with AdaSplash's optimized implementation, technically allows for a sub-quadratic backward pass, offering a significant advantage for training efficiency.

About the Speaker(s)

The work on AdaSplash was a collaborative effort by Nuno Gonçalves, Marcos V. Treviso, and Andre Martins, all affiliated with IST Lisbon. Nuno Gonçalves delivered the presentation, outlining the technical details and experimental findings of their paper. Their research focuses on advancing the efficiency and practical applicability of attention mechanisms in deep learning, particularly within the context of large-scale transformer models.