DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention

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 $α$-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, we show that DashAttention is non-dispersive, translating to better long-context modeling ability. Experiments with LLMs show that DashAttention achieves comparable accuracy as full attention with 75% sparsity and a better Pareto frontier than NSA and InfLLMv2, especially in high-sparsity regimes. We also provide an efficient, GPU-aware implementation of DashAttention in Triton, which achieves a speedup of up to over FlashAttention-3 at inference time. Overall, DashAttention offers a cost-effective strategy to model long contexts.

• The paper presents a three-stage hierarchical attention mechanism that adaptively sparsifies context using differentiable entmax routing.
• It achieves up to 3.3x speedup over FlashAttention-3 while maintaining competitive accuracy even at 90% sparsity.
• The method preserves gradient flow across hierarchical stages, enabling scalable long-context modeling on modern GPUs.

DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention

Motivation and Context

Long-context language modeling critically depends on efficiently retrieving pertinent information from extensive contexts, balancing between selective focus and adaptive flexibility. Standard dense softmax attention fails to restrict computation to relevant tokens, resulting in excessive dispersion and computational inefficiency, especially as the context grows. Alternatively, hard sparse methods based on top-$k$ block selection constrain flexibility, introduce hyperparameters for fixed budgets, and sever gradients between coarse and fine attention stages. There is a practical and theoretical gap for an attention mechanism that is adaptively sparse, fully differentiable, and hardware-efficient.

DashAttention Architecture

DashAttention introduces a three-stage hierarchical attention framework designed to address the shortcomings of both dense softmax and static sparse attention methods.

Figure 1: High-level pipeline for DashAttention: Stage 0 summarizes token chunks with local SDPA; Stage 1 uses $\alpha$-entmax for adaptive sparse routing; Stage 2 refines softmax over routed tokens, applying bias derived from routing probabilities.

• Stage 0: Local Chunk Summarization Chunks of tokens are summarized into compact representations via a learned local softmax-weighted pooling. This enables smooth interpolation between mean pooling and learned, expressive chunk summaries, facilitating seamless adaption from pretrained dense attention models.
• Stage 1: Entmax Chunk Routing Chunk summaries are scored with respect to the query. Rather than applying a top-$k$ truncation, DashAttention applies a differentiable $\alpha$-entmax transformation, producing a sparse, query-adaptive probability distribution over the chunks. As a result, sparsity is dynamic and data-driven, allowing more or fewer chunks per query/head/layer.

  Figure 2: Visualization of adaptive sparsity under varying $\alpha$ values and comparison with top-$k$ softmax, showing the dynamic support offered by entmax versus the rigid support of top-$k$.

• Stage 2: Prior-Induced Sparse Softmax Attention The prior from Stage 1 is incorporated into a token-level softmax via a KL-divergence regularizer, yielding an attention pattern that is block-sparse, adaptively biased, and fully differentiable. The masked softmax operates over only the routed chunks, with per-row bias terms derived from entmax routing probabilities.

Theoretical Analysis

DashAttention addresses the dispersion problem typical of long-context softmax attention, where entropy of the attention distribution scales as $\log n$ with sequence length, hampering long-range dependency modeling. By employing adaptive sparse routing with entmax and avoiding dispersive softmax-based head aggregation, DashAttention maintains low entropy and non-dispersive properties even as contexts become extremely long.

Unlike fixed-budget top-$k$ attention, where selection support is discontinuous and non-differentiable, DashAttention’s entmax routing is continuous and enables direct gradient flow from token-level outputs to chunk-level summarization parameters.

Efficiency and Implementation

DashAttention is engineered for high efficiency on GPU hardware, implemented in Triton as a series of specialized fused kernels:

• Stage 0: Produces chunk representations in one pass, caching after chunk completion.
• Stage 1: Applies entmax routing via AdaSplash-2, storing bit-packed active chunk masks.
• Stage 2: Uses a masked FlashAttention kernel, biasing the scores per routed chunk, supporting efficient prefill and incremental decoding.

Across extensive benchmarks, DashAttention achieves up to $3.3\times$ speedup over FlashAttention-3 and $\alpha$0 over InfLLMv2, with the gap widening at higher sparsity and longer context lengths due to avoidance of inefficiencies in top-$\alpha$1 index materialization.

Empirical Evaluation and Results

DashAttention is evaluated on long-context (RULER, HELMET) and general NLP benchmarks, using Minicpm-4 LLMs of 1B, 3B, and 8B parameters, with models initialized from full attention pretraining and subject to long-context continual training.

Key results include:

• Comparable or superior accuracy to full attention at 75% sparsity on hard long-context tasks, notably outperforming NSA and InfLLMv2 in high-diffusion retrieval settings.
• Strong cost-effectiveness: DashAttention retains 39.4% overall accuracy on HELMET at $\alpha$290% sparsity, a 9–19% absolute improvement over prior hierarchical sparse baselines.
• Short-context evaluation yields no degradation relative to dense or other sparse baselines, confirming the method's generality and absence of negative side effects.

Dynamic sparsity allocation is evidenced by layerwise sparsity profiles: early layers are denser, with increasingly aggressive sparsification deeper in the model, a self-organizing property without manual schedule design.

Implications and Future Directions

Practical implications:

DashAttention enables cost-effective, high-throughput inference for LLMs in settings demanding ultra-long context, with dynamic memory and compute allocation. The adaptability of the architecture to both training and inference obviates the need for manual tuning of support budgets or retraining with static sparsity.

Theoretical implications:

The integration of entmax into hierarchical routing provides a template for future adaptive sparse architectures, particularly in domains where head- and token-level mixture-of-experts or structured selection is beneficial. The maintenance of gradient pathways through the entire attention hierarchy is likely advantageous for optimization, meta-learning, and low-data adaptation scenarios.

Future work:

There is clear applicability of DashAttention’s architecture and techniques to hybrid (e.g., Mamba-Transformer) models, expansion to elementwise or multi-scale dynamic sparsity, and, critically, integration into modern inference servers (e.g., vLLM, SGLang) to close the last-mile efficiency gap in practical LLM deployment.

Conclusion

DashAttention establishes a new paradigm for hierarchical sparse attention: adaptively sparse, fully differentiable, and highly efficient. Empirical evidence shows robust long-context retrieval performance and superior efficiency on modern hardware, with favorable theoretical properties regarding dispersion and optimization. DashAttention forms a concrete foundation for the next generation of scalable, adaptive attention mechanisms in large-scale language modeling and beyond.