🤖 machine learning

AdaSplash-2: Faster Differentiable Sparse Attention

AdaSplash-2 introduces a novel histogram-based initialization and sparsity-aware GPU implementation to drastically accelerate $\alpha$ -entmax sparse attention, enabling training speeds comparable to FlashAttention-2 at moderate-to-high sparsity levels while achieving superior performance in long-context settings.

Original authors: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

Published 2026-04-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive library (a Transformer AI model) where every book (a word in a sentence) needs to read every other book to understand the context.

In a traditional library, this is a nightmare. If you have 1,000 books, you have to read 1,000,000 pairs of connections. If you have 100,000 books, the number of connections explodes to 10 billion. This is the quadratic bottleneck that slows down modern AI when it tries to read long documents.

To fix this, researchers invented Sparse Attention. Instead of reading every book, the AI learns to ignore the irrelevant ones and only focus on the important ones. This is like a librarian who only pulls out the top 5 most relevant books for a query, ignoring the rest.

However, there was a catch. The "smart librarian" algorithm they were using (called $\alpha$ -entmax) was too slow to figure out which books to pick. It had to do a lot of heavy math, like a librarian trying to calculate the perfect probability for every single book before deciding what to ignore. This calculation was so slow that it canceled out the speed gains from ignoring the books.

Enter ADASPLASH-2: The "Super-Speedy Librarian."

This paper introduces a new, faster way to run this smart librarian algorithm. Here is how it works, using some everyday analogies:

1. The Problem: The "Guessing Game"

To decide which books to ignore, the algorithm needs to find a specific "cutoff score" (called $\tau$ ).

Old Way: Imagine trying to find the exact temperature where water boils. You have to keep guessing, checking, and adjusting your guess over and over again. This takes a long time.
The Issue: In the previous version (ADASPLASH), the computer had to scan the entire list of scores multiple times to find this cutoff. It was accurate but slow.

2. The Solution: The "Histogram Shortcut"

The authors of ADASPLASH-2 came up with a brilliant trick: The Histogram.

Instead of looking at every single score one by one, imagine you have a giant bucket with 100 slots (bins).

As the computer scans the scores, it doesn't store the exact numbers. Instead, it just drops a marble into the slot that matches the score's range.
The Magic: By the time it finishes scanning, it has a quick "map" (a histogram) of where all the scores are clustered.
The Result: This map gives the computer a very good guess at the cutoff score immediately. It's like looking at the distribution of marbles and saying, "Okay, the boiling point is definitely between 90 and 95 degrees," without needing to measure every single drop of water.

Because this "map" is built instantly in the computer's super-fast internal memory (SRAM), the algorithm only needs 1 or 2 quick checks to find the exact answer, instead of dozens.

3. The "Skip the Empty Boxes" Trick

Once the computer knows the cutoff, it knows which books (or blocks of data) are irrelevant.

The Old Way: The computer would still drive past the empty shelves, check them, realize they are empty, and then move on. This wastes fuel (energy and time).
The New Way: ADASPLASH-2 uses a bit-packed mask. Think of this as a highly efficient checklist where a single "1" or "0" tells the computer exactly which shelves to visit. If a shelf is empty, the computer literally skips the drive to that shelf. It doesn't even load the data into memory.

Why Does This Matter?

Speed: In situations where the AI is reading very long documents (like a whole novel or a legal contract), the "smart librarian" (ADASPLASH-2) is now faster than the standard "dumb librarian" (FlashAttention-2) that reads everything.
Smarter Long-Context: Because the AI can handle long texts more efficiently without getting bogged down by math, it actually gets smarter at long tasks. It stops getting confused by irrelevant details and focuses on the story.
No Trade-off: Usually, making things faster makes them less accurate. But here, the new method is so efficient that the AI performs just as well (or better) on short tasks and significantly better on long tasks.

The Bottom Line

ADASPLASH-2 is like upgrading a librarian from someone who reads every single page of every book to find the answer, to a super-intelligent assistant who:

Quickly scans the spines to make a rough guess (The Histogram).
Instantly knows exactly which books to pull off the shelf.
Ignores the empty shelves entirely.

This allows AI models to read longer books, understand complex stories, and do it all in the blink of an eye.

1. Problem Statement

Transformer models face a computational bottleneck due to the quadratic time and memory complexity of the self-attention mechanism ( $O(n^2)$ ), particularly in long-context scenarios. While FlashAttention (and its successors) optimized the standard softmax attention by fusing operations and utilizing on-chip memory, softmax remains inherently dense (every token receives non-zero probability). This density causes attention mass to spread over irrelevant tokens in long contexts, leading to representational collapse and poor generalization.

$\alpha$ -entmax attention is a promising alternative that produces input-dependent sparse outputs (exact zeros) while remaining differentiable. However, computing the normalization constant ( $\tau$ ) for $\alpha$ -entmax requires iterative root-finding methods (unlike the closed-form log-sum-exp of softmax). Previous attempts to implement this on GPUs (e.g., ADASPLASH) suffered from high computational overhead due to the multiple passes required to converge on $\tau$ , preventing them from matching the speed of optimized softmax kernels like FlashAttention-2.

2. Methodology: ADASPLASH-2

The authors propose ADASPLASH-2, a hardware-aware algorithm that drastically reduces the cost of computing the $\alpha$ -entmax normalizer $\tau$ by leveraging on-chip SRAM and histogram-based initialization.

Core Innovations

On-Chip Histogram Initialization:
- Instead of iterating over raw scores to find $\tau$ , the algorithm streams attention scores through the GPU's on-chip SRAM.
- It constructs a compact histogram of the scores (discretized into $B$ bins) entirely in SRAM. This histogram represents the distribution of scores without materializing dense intermediate matrices.
- Theoretical Guarantee: The authors prove that solving for $\tau$ on this histogram yields a provable lower bound ( $\tau_h \leq \tau^*$ ) on the true normalizer. This provides a high-quality initial estimate.
Rapid Convergence (1–2 Iterations):
- Because the histogram initialization is extremely close to the true root, a subsequent hybrid solver (combining Halley, Newton, and Secant methods depending on $\alpha$ ) converges to the exact solution in typically 1 to 2 iterations.
- This contrasts with previous methods that required many iterations, effectively turning a multi-pass problem into a near-single-pass operation.
Sparsity-Aware GPU Implementation:
- Bit-Packed Block Masking: During the refinement phase, the algorithm constructs a binary mask indicating which blocks of the attention matrix contain non-zero weights. This mask is stored using bit-packing (32 blocks per int32), minimizing memory overhead.
- Efficient Traversal: The backward pass and the final output computation utilize native GPU instructions (fns - find next set, popc - population count) to traverse only the non-zero blocks, skipping zero blocks with negligible overhead.
Hybrid Solver Strategy:
- The paper derives smoothness properties of the $\alpha$ $α$ -entmax function to select the optimal root-finding step:
  - $\alpha \leq 1.5$ : Uses Halley's method (uses second derivative).
  - $1.5 < \alpha \leq 2$ : Uses Newton's method (uses first derivative).
  - $\alpha > 2$ : Uses Secant method.
- A safeguarded fallback to bisection ensures stability if an update step falls outside valid bounds.

3. Key Contributions

Algorithmic Efficiency: Introduced a histogram-based initialization that reduces the number of iterations for $\tau$ computation from many to typically 1–2, enabling $\alpha$ -entmax to compete with softmax in speed.
Hardware Optimization: Designed a Triton kernel that exploits GPU memory hierarchies (HBM $\to$ SRAM) and bit-packed data structures to handle dynamic sparsity efficiently.
Theoretical Guarantees: Provided proofs for the lower-bound property of the histogram approximation and the continuity of derivatives for different $\alpha$ values, ensuring the stability of the hybrid solver.
Performance Breakthrough: Demonstrated that ADASPLASH-2 matches or exceeds FlashAttention-2 in training throughput when block sparsity is moderate-to-high (>60%), a condition common in long-context training.

4. Experimental Results

The authors evaluated ADASPLASH-2 on synthetic benchmarks, language modeling tasks, and long-context evaluations.

Runtime Efficiency:
- Moderate-to-High Sparsity: In regimes with >60% block sparsity, ADASPLASH-2 outperforms both CUDA and Triton implementations of FlashAttention-2.
- Speedup: At high sparsity, it achieves >2x speedup over FlashAttention-2.
- Backward Pass: The backward pass sees the most significant gains because the sparsity-aware traversal skips zero blocks, offsetting the slightly higher cost of the forward pass normalization.
Language Modeling (LLaMA-3 based models):
- Long-Context (RULER & HELMET): Models trained with $\alpha$ -entmax (specifically with NAPE positional encoding) significantly outperform softmax baselines on long-context tasks (up to 32K tokens). They show substantial gains in Variable Tracking and Word Extraction, tasks requiring precise attention to specific spans.
- Short-Context: Models trained with ADASPLASH-2 match or slightly improve upon softmax baselines on standard short-context benchmarks (e.g., ARC, CSQA, LAMBADA), proving that sparsity does not degrade performance in dense regimes.
- In-Context Learning (ICL): $\alpha$ -entmax + NAPE shows superior few-shot performance as context length increases, suggesting sparsity helps the model focus on relevant in-context examples.

5. Significance and Impact

Scalability: ADASPLASH-2 removes a major barrier to training sparse attention mechanisms at scale. It allows researchers to utilize the theoretical benefits of $\alpha$ -entmax (dynamic sparsity, better long-context generalization) without sacrificing training speed.
Long-Context Solutions: By enabling efficient training of sparse attention, this work offers a practical path to solving the "attention dispersion" problem in large language models, where softmax fails to distinguish relevant tokens in very long sequences.
Hardware Co-Design: The paper serves as a blueprint for designing attention kernels that are not just algorithmically sound but specifically optimized for modern GPU memory hierarchies (SRAM utilization, bit-packing, and instruction-level parallelism).

In summary, ADASPLASH-2 successfully bridges the gap between the theoretical advantages of sparse, differentiable attention and the practical requirements of high-performance GPU training, making it a viable replacement for softmax in long-context transformer training.