---
title: 'AdaSplash-2: Hardware-Aware Sparse Attention'
url: https://www.emergentmind.com/topics/adasplash-2
type: topic
---

# AdaSplash-2: Hardware-Aware Sparse Attention

AdaSplash-2 is a hardware-aware implementation of differentiable sparse attention based on the $\alpha$-entmax transformation, targeting the elimination of the quadratic computational bottleneck in long-context transformer models. By introducing a novel histogram-based initialization for the entmax normalization root and a GPU kernel that efficiently exploits block sparsity, AdaSplash-2 achieves competitive or superior runtimes compared to FlashAttention-2 in settings where attention is highly sparse. This method demonstrates its effectiveness both in synthetic benchmarks and large-scale language modeling tasks, where it not only matches softmax-based baselines on short contexts but also realizes significant gains as input lengths and sparsity increase [2604.15180].

## 1. $\alpha$-entmax Attention and Motivation

Standard softmax-based attention, defined for scores $s\in\mathbb{R}^n$ by
$$
\text{softmax}(s) = \exp(s-\tau\mathbf{1}), \qquad \tau = \log\sum_j \exp(s_j),
$$
assigns nonzero mass to all tokens, inducing $O(n^2)$ work per layer and encouraging distributed, often diffuse, attention which can impede learning in long-context settings.

$\alpha$-entmax attention [Peters et al. 2019] generalizes softmax and sparsemax by allowing a tunable entropic regularization through the Tsallis entropy:
$$
\alpha\text{-entmax}(s) := \arg\max_{p\in\Delta_n}~p^\top s + H_\alpha(p), \qquad \Delta_n = \{p \geq 0,\ \mathbf{1}^\top p = 1\},
$$
leading to the closed-form solution
$$
\mathrm{entmax}_\alpha(s) = \left[(\alpha-1)s - \tau\mathbf{1}\right]_+^{1/(\alpha-1)},
$$
subject to $\sum_i \mathrm{entmax}_\alpha(s)_i = 1$ and $[x]_+ = \max(0, x)$. The normalizer $\tau\in\mathbb{R}$ is found by solving the root of
$$
f(\tau) = -1 + \sum_k [(\alpha-1)s_k - \tau]_+^{1/(\alpha-1)} = 0.
$$

A key property is input-dependent sparsity: for each $\alpha>1$, entmax assigns exact zeros wherever $(\alpha-1)s_i \leq \tau$, generating probability vectors with adaptive support. This behavior allows attention computation and memory usage to scale with the true, contextual support size rather than the full $n^2$ space, addressing both computational and representational inefficiencies in long-context transformers.

## 2. Histogram-Based Normalizer Initialization

A practical challenge for $\alpha$-entmax layers is the efficient solution of $f(\tau)=0$ per row. Traditional root-finding methods such as bisection are robust but converge slowly, whereas Halley or Newton methods are fast but require a good starting point.

AdaSplash-2 introduces a hardware-friendly histogram-based initialization that stores a binned summary of transformed scores in on-chip SRAM. The method comprises:

- Centering scores as $z = (\alpha-1)s - (m-1)\mathbf{1}$ with $m=(\alpha-1)\max(s)$, normalizing $z$ so $\max(z)=1$ and ensuring $\tau^\star\in[0,1]$.
- Discretizing $[0,1]$ into $B$ bins of width $h=1/B$ and assigning each $z_j$ to its appropriate bin.
- Constructing a histogram $H \in \mathbb{N}^B$ where $H_k$ counts the number of $z_j$ falling into each bin.
- Approximating the normalizer by replacing $z_j$ with its bin's left edge in the normalizer equation, yielding a reduced monotone root-finding problem:
  $$
  f_h(\tau) = -1 + \sum_{k=0}^{B-1} H_k [k/B - \tau]_+^{1/(\alpha-1)}
  $$
- By mathematical proposition, the root $\tau_h$ of $f_h$ provides a lower bound within $h$ of the exact $\tau^\star$: $\tau^\star-h < \tau_h \leq \tau^\star$.

A single safeguar ded hybrid root-finding step (Halley for $\alpha\leq 1.5$, Newton if $1.5<\alpha\leq 2$, secant for $\alpha > 2$, fallback to bisection if needed) refines $\tau_h$ to the true root, typically converging within 1–2 passes over the data. The histogram method requires only $O(B)$ words of fast on-chip memory and substantially accelerates normalization compared to standard techniques.

## 3. Sparsity-Aware GPU Pipeline

AdaSplash-2 is implemented as a Triton GPU kernel organized into four key phases per query block $Q_i$ (of shape $B_r\times d$) over key blocks $K_j$ ($B_c\times d$):

1. **Row Maximum Computation:** Compute $m_i = \max_j \max_\text{col}(Q_i K_j^\top)$ per query block.
2. **Histogram Construction:** For each tile $Q_i$ vs $K_j$, scale the score tile to $[0,1]$, bin indices, and build bit-packed local histograms of shape $B_r\times B$ in SRAM.
3. **$\tau$ Refinement and Block Masking:** Solve for $\tau_h$ using special-case or general histogram solvers; refine to final $\tau$ with a hybrid step; simultaneously, build bit-packed masks $M_{ij}\in\{0,1\}$ per block, indicating which blocks contain nonzero attention.
4. **Sparse MatMul:** Using the mask, load only nonzero key and value blocks; GPU native population-count instructions enable efficient traversal, accumulating $O_i += P_{ij} V_j$ for nonzero attention blocks.

The computational complexity scales with nonzero block fraction: in the worst-case, $O(n^2d)$ (as for dense attention), but the actual work is proportional to $n^2(1-s)d$, where $s$ is block sparsity. Histogram and tile management overhead is $O(nd + nB)$, which is negligible when $B \ll n$. At high sparsity ($s \gtrsim 60\%$), especially for long-contexts ($n \gtrsim 16K$), backward passes are up to $2\times$ faster than FlashAttention-2.

## 4. Empirical Results

AdaSplash-2 was evaluated on NVIDIA A6000 and H100 GPUs using Triton-based kernels, with baselines including CUDA/Triton FlashAttention-2 ("FA2"). Synthetic and language modeling experiments were conducted:

- **Root-Finder Evaluation:** For $n=4096$ sampled scores $\sim \mathcal{N}(0,1)$, histogram initialization with $B\in \{4,8,16\}$ drastically reduces normalizer error $|\tau-\tau^\star|$ to $10^{-6}$ after only one iteration.
- **Sparsity-Sensitivity:** For causal attention with $n=16\,\text{K}$, at block sparsity $s>60\%$, AdaSplash-2 outperforms FA2 by $10{-}50\%$ and achieves $2\times$ speedup at $s\to90\%$.
- **Context Scaling:** Using block sparsity patterns extracted from a 1B $\alpha$-entmax-NAPE LM, backward speedups emerge even at $4\,\text{K}$ (with $20\%$ sparsity); step time surpasses FA2 beyond $16\,\text{K}$ length.
- **Large-Scale Language Modeling:** LLaMA-3 models (350M, 1B) trained on $50$B DCLM-Edu tokens ($4\text{K}$ context, bf16 precision). At short context ($4$K), entmax+NAPE obtains best average scores: 48.1 (350M) vs. 47.3 (softmax+RoPE) and 47.1 (softmax+NAPE); 1B model: ppl 11.42 vs 11.97 (softmax+NAPE), avg accuracy 53.1 vs 53.0. On long-context tasks (RULER at up to 32K), entmax+NAPE outperforms softmax variants by +2–6 points average and +2.2 avg at 32K for HELMET ICL.

## 5. Limitations, Trade-offs, and Practical Considerations

While AdaSplash-2 achieves significant speedups for backward propagation in high-sparsity regimes, its forward pass is slower than FA2 for dense attention due to histogram management overhead. However, this gap narrows as block sparsity increases above 30%. Notably, while $\alpha$-entmax enables dynamic, differentiable sparsity, current kernels still require scanning all keys at inference time; highly efficient inference kernels remain an open engineering challenge.

The histogram initialization scheme requires that $B \underline{\wedge} n$ fit per-row in SRAM, which becomes a constraint for extremely long sequences. To address this, AdaSplash-2 incorporates an overflow handling scheme for periodic histogram flushing. The hybrid solver’s refinement still necessitates a secondary pass over scores, although this could potentially be fused with the sparse matmul to improve efficiency.

## 6. Scenarios of Maximal Benefit and Future Directions

AdaSplash-2 is particularly advantageous in:

- Long-context transformer training where sparsity emerges organically (e.g., document-level QA, generative modeling at scale).
- Context lengths of 8K–32K, where block sparsity greater than 60% is commonly observed early in training.
- Tasks where static patterns or top-$k$ sparsity baselines are surpassed by differentiable, dynamic sparsity.

Future research and engineering directions include: (i) fused inference kernels aligning entmax computation with key-value retrieval, (ii) mixed-precision and hardware-specific optimizations (e.g., NVIDIA Hopper TMA/TMAMMA), (iii) adapting the $\alpha$ parameter per head or per layer, and (iv) extending techniques to encoder–decoder and cross-attention modules.

By integrating rapid, provable initialization and fine-tuned GPU kernels, AdaSplash-2 delivers expressive differentiable sparse attention for large-scale models, achieving or exceeding FlashAttention-2 speed in moderate and high sparsity settings and providing robust generalization for both short and long-context tasks [2604.15180].

Source: https://www.emergentmind.com/topics/adasplash-2