Paper Overview
- Field: NLP
- Authors: Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti
- Published: 2026-05-19
- arXiv: 2505.14310
- Problem: Hierarchical attention (NSA, InfLLMv2) relies on a fixed top-k selection of KV blocks, which assumes a uniform number of relevant tokens per query and blocks gradient flow between sparse and dense stages.
- Method: DashAttention replaces top-k with the $\alpha$-entmax adaptive sparse transformation, enabling per-query variable block selection and a fully differentiable two-stage hierarchy (sparse prior + softmax attention).
- Theoretical property: DashAttention is proven non-dispersive, distinguishing it from prior hierarchical attention and improving long-context modeling.
- Empirical results: Matches full attention accuracy at 75% sparsity; achieves a better accuracy-speed Pareto frontier than NSA and InfLLMv2 at high sparsity on large language models.
- System contribution: A Triton-based GPU-aware implementation delivers inference faster than FlashAttention-3.
- Practical impact: Provides a cost-effective pathway for deploying long-context LLMs by reducing KV computation while preserving accuracy and end-to-end differentiability for training.
Abstract
Current hierarchical attention methods, such as NSA and InfLLMv2, select the top-k relevant key-value (KV) blocks based on coarse attention scores and subsequently apply fine-grained softmax attention on the selected tokens. However, the top-k operation assumes the number of relevant tokens for any query is fixed and it precludes the gradient flow between the sparse and dense stages. In this work, we propose DashAttention (Differentiable and Adaptive Sparse Hierarchical Attention), which leverages the adaptively sparse $\alpha$-entmax transformation to select a variable number of blocks according to the current query in the first stage. This in turn provides a prior for the second-stage softmax attention, keeping the entire hierarchy fully differentiable. Contrary to other hierarchical attention methods, the authors prove that DashAttention is non-dispersive, leading to better long-context modeling capability. Experiments on large language models demonstrate that DashAttention matches full attention accuracy at 75% sparsity and shows a better Pareto frontier than NSA and InfLLMv2 in the high-sparsity regime. The authors also provide a Triton-based GPU-aware implementation that is faster than FlashAttention-3 at inference. Overall, DashAttention offers a cost-effective strategy for long-context modeling.