AdaSplash-2: Faster Differentiable Sparse Attention

Abstract: 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 AdaSplash-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, AdaSplash-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.

• The paper presents a histogram-based normalization for entmax sparse attention that guarantees a lower bound and reduces iterative computation.
• It details sparsity-aware block masking using GPU bitwise operations to confine computation to active regions, enhancing throughput in long-context models.
• Empirical results show up to 2x speedup in training with increased block sparsity and maintained performance compared to dense softmax baselines.

AdaSplash-2: Hardware-Efficient Differentiable Sparse Attention

Introduction and Motivation

Transformer-based models are fundamentally constrained by the quadratic complexity of their attention mechanisms, especially when scaling to extended context lengths. While FlashAttention and its descendants have enabled high-throughput softmax-based attention by exploiting GPU memory hierarchies and single-pass online normalization, they remain coupled to dense attention distributions—each token receives nonzero mass due to the properties of softmax normalization. This uniform density degrades long-context generalization, impedes distinguishing salient tokens, and exacerbates computational bottlenecks.

Entmax, a parameterized family of differentiable sparse transformations, provides a principled alternative by yielding input-adaptive dynamic sparsity and exact zeros in the attention output. Despite demonstrated empirical gains in long-context settings, entmax attention has not been competitive in training throughput due to the iterative cost of computing its normalization threshold, $\tau$, which lacks the additive structure of logsumexp. The original AdaSplash kernel [goncalves2025adasplash] made strides in this direction but incurred multiple passes over large key blocks, limiting its efficacy in moderate sparsity regimes.

AdaSplash-2 (2604.15180) introduces a hardware-aware approach to address both normalization efficiency and sparse memory alignment on modern GPUs, providing end-to-end throughput that matches or exceeds heavily optimized dense attention baselines in practical scenarios.

---

Algorithmic Innovations

Histogram-Based Normalization

At the core of AdaSplash-2 is a histogram-based approximation for the entmax normalizer. Rather than iteratively evaluating the normalization function $f(\tau)$ over all $n$ scores, AdaSplash-2 discretizes the score range, constructs a compact histogram in on-chip SRAM during block-wise streaming, and solves the approximate normalization over $B\ll n$ bin counts. The histogram is constructed such that the normalization threshold obtained is a guaranteed lower bound on the true value, with the absolute error bounded by the bin width $h = 1/B$.

This approximation is theoretically justified to never overestimate the support set, ensuring that sparse block masking never excludes genuine non-zeros and that further refinement can proceed safely from this initialization. Higher-order refinement based on properties of $f$ (e.g., Halley or Newton steps, selected by $\alpha$) converges to the exact root in typically one or two Passes.

Sparsity-Aware Block Masking and GPU Implementation

AdaSplash-2 exploits the histogram to construct a lightweight, bitpacked mask indicating nonzero $64 \times 64$ attention blocks, efficiently encoded and traversed by native GPU bitwise instructions. Both the forward and backward passes are restricted to active regions, yielding substantial throughput gains with naturally arising input-dependent sparsity.

The forwarding sequence consists of: (1) streaming blocks for rowwise maxima, (2) on-chip histogram construction, (3) root finding for normalization with a hybrid higher-order/bisection solver, and (4) sparse output computation using the generated block mask.

The entire design is orchestrated within Triton kernels, with meticulous attention to memory alignment, avoidance of atomic contention in histogramming, and periodic buffer flushes to prevent overflow at extreme input lengths.

---

Empirical Results

Sparsity-Efficiency Tradeoff

AdaSplash-2 outperforms a highly-optimized CUDA implementation of FlashAttention-2 across a wide spectrum of block sparsity ratios and maintains its advantage as sparsity increases. Notably, in contexts where block sparsity exceeds 60%—a regime prevalent in long-context language modeling—training steps are up to twice as fast as the FlashAttention-2 baseline, despite additional root-finding computations.

Figure 1: AdaSplash-2's runtime as a function of input sparsity exhibits significantly improved efficiency in moderate and high sparsity regimes relative to FlashAttention-2.

Histogram Initialization Efficacy

By analyzing the mean absolute error of root-finding methods, AdaSplash-2 demonstrates that histogram-based initialization dramatically reduces the number of solver iterations required to reach machine-precision thresholds. The hybrid solver initialized from histogram estimates converges effectively in one or two steps for $\alpha=1.5$, orders of magnitude faster than standard bisection.

Figure 2: The hybrid root-finding method with histogram initialization achieves near zero error after a single iteration, contrasting with prior iterative solvers.

Long-Context LLM Performance

Entmax attention models (with NAPE position encoding) trained via AdaSplash-2 achieve average accuracy improvements over softmax-trained baselines on RULER and HELMET ICL benchmarks, especially as sequence lengths increase to 32K tokens. These models also match or slightly surpass dense softmax models even in short-context settings, indicating no trade-off in short-range benchmarks.

Context Length Scaling

AdaSplash-2's forward+backward efficiency scales favorably with length, driven by naturally emerging sparsity in LLMs. As context length increases, empirical mean block sparsity grows, and AdaSplash-2 increasingly outpaces dense baselines—reflecting its strong practical scaling attributes.

Figure 3: Normalized step times across varying context lengths and sparsity, showing AdaSplash-2's advantage, especially in the backward pass as context scales.

---

Practical and Theoretical Implications

The AdaSplash-2 paradigm pushes the boundaries of hardware-aligned sparse attention computation for training large-scale models, especially as demand for longer contexts and efficient memory usage intensifies. The guaranteed correctness of histogram-based normalization supports safe sparse execution, while the bitpacked masking ensures minimal wasted work on zero regions. This unlocks efficiency scaling with data-dependent sparsity, not just through static pruning or limiting window size.

Theoretically, the method generalizes to other convex, monotonic, sparse normalization functions amenable to histogram discretization. Empirical findings suggest that model performance on long-context reasoning, variable/scenario tracking, and in-context learning is measurably enhanced by leveraging entmax sparsity.

Potential Future Directions:

• Extending the histogram-based normalization paradigm to other architectures, including Mixture-of-Experts and attention variants beyond entmax.
• Optimizing for inference by amortizing multi-pass costs, exploring approximate sparse attention during generation, and aligning with hardware developments (e.g., Hopper architectural features).
• Further automating and adapting bin sizes and block configurations dynamically based on observed input distributions, possibly leveraging learned block partitioning.

---

Conclusion

AdaSplash-2 demonstrates that hardware-aware, histogram-accelerated sparse attention is feasible at scale with theoretical guarantees and measurable empirical gains. It closes the optimization gap between academically motivated dynamic sparse attention and production-grade throughput, making differentiable input-dependent sparsity a practical tool for efficient long-context model training and deployment.

(2604.15180)