Long-Context Modeling with Dynamic Hierarchical Sparse Attention for On-Device LLMs

TL;DR

This work targets the quadratic attention cost that hinders long-context LLMs in resource-constrained, on-device settings. It introduces Dynamic Hierarchical Sparse Attention (DHSA), a data-driven plug-in that dynamically predicts attention sparsity via hierarchical chunking and boundary detection, without retraining the base model. DHSA computes chunk-level similarity $S_c ∈ ℝ^{N_c×N_c}$ and upsamples to token-level $S_t ∈ ℝ^{L×L}$, applying a per-token budget $N_b$ to form a sparsity mask, and it uses dynamic chunk boundaries learned through a boundary predictor comprising an Encoder, a Feature Fusion module, and an MLP. Empirically, DHSA matches dense-attention accuracy on Needle-in-a-Haystack and LongBench while reducing prefill latency by about $25$–$45ig%$ and peak memory by $30$–$35ig%$, with consistent 6–18% relative gains over other dynamic baselines, demonstrating practical, adaptable long-context efficiency for on-device LLMs. This approach offers a scalable path toward efficient long-context reasoning with content-aware sparsity and budget-aware inference.

Abstract

The quadratic cost of attention hinders the scalability of long-context LLMs, especially in resource-constrained settings. Existing static sparse methods such as sliding windows or global tokens utilizes the sparsity of attention to reduce the cost of attention, but poorly adapts to the content-dependent variations in attention due to their staticity. While previous work has proposed several dynamic approaches to improve flexibility, they still depend on predefined templates or heuristic mechanisms. Such strategies reduce generality and prune tokens that remain contextually important, limiting their accuracy across diverse tasks. To tackle these bottlenecks of existing methods for long-context modeling, we introduce Dynamic Hierarchical Sparse Attention (DHSA), a data-driven framework that dynamically predicts attention sparsity online without retraining. Our proposed DHSA adaptively segments sequences into variable-length chunks, then computes chunk representations by aggregating the token embeddings within each chunk. To avoid the bias introduced by varying chunk lengths, we apply length-normalized aggregation that scales the averaged embeddings by the square root of the chunk size. Finally, DHSA upsamples the chunk-level similarity scores to token level similarities to calculate importance scores that determine which token-level interactions should be preserved. Our experiments on Gemma2 with Needle-in-a-Haystack Test and LongBench show that DHSA matches dense attention in accuracy, while reducing prefill latency by 20-60% and peak memory usage by 35%. Compared to other representative baselines such as block sparse attention, DHSA achieves consistently higher accuracy (6-18% relative gains) with comparable or lower cost, offering an efficient and adaptable solution for long-context on-device LLMs.

Figures (11)

• Figure 1: Overview of the proposed Dynamic Hierarchical Sparse Attention (DHSA) framework.
• Figure 2: Needle-in-a-haystack results on Gemma2-2b-it (maximum context length = 8k). The budget is specified per query per layer.
• Figure 3: Architecture of the boundary predictor, consisting of a shared encoder, a feature fusion module, and an MLP classifier.
• Figure 4: Illustration of our automatic labeling strategy. The core intuition is that tokens within a coherent span tend to receive similar attention distributions.
• Figure 5: Token-level similarity matrices with static (left) and dynamic (right) chunking in Gemma2-2b-it. Static chunking uses fixed-size chunks, while dynamic chunking adapts boundaries to content, shown for a local attention layer (layer 0) and a global attention layer (layer 1).