DHSA: Dynamic Hierarchical Sparse Attention

This repository contains the official implementation of [ICML 26 Spotlight, NeurIPS 25 Efficient Reasoning Workshop] paper Long-Context Modeling with Dynamic Hierarchical Sparse Attention for Memory-Constrained LLM Inference.

Table of Contents

• Overview
• Results
• Repository Structure
• Installation
• Predictor Checkpoints
• Quick Start
• Acknowledgements

Overview

LLMs face efficiency limits from the quadratic cost of dense attention. Static sparse methods (e.g., sliding windows, global tokens) reduce computation but cannot adapt to content. DHSA is a data-driven plugin that predicts attention sparsity on the fly without retraining the base model. By extending Block-Sparse Attention with variable-sized query and key blocks and effective sparsity prediction, DHSA better preserves accuracy under highly sparse regimes.

This release supports Qwen2/Qwen2.5 and Llama with the following sparse attention mask:

DHSA handles prompt prefill. Autoregressive decoding falls back to the model's original FlashAttention-2 implementation.

Results

NIAH uses the repository runner's score and is not on the same scale as LongBench or RULER. Compare values vertically within each column.

Llama-3.1-8B-Instruct BF16

Llama-3.1-8B-Instruct 4-bit

Qwen2.5-3B-Instruct BF16

Kernel-level Latency

Latency is measured on one A100 40 GB GPU for a single attention layer, including mask construction and the attention kernel. Values are median milliseconds from five measured iterations after two warmups. Parentheses show speedup over Full FA2 measured in the same run. We do not distinguish between BF16 and 4-bit settings because quantization does not affect kernel-level latency.

Llama-3.1-8B-Instruct

Qwen2.5-3B-Instruct

Repository Structure

DHSA/
├── Block-Sparse-Attention/
├── checkpoints/
├── data/
├── results_latency/
├── results_longbench/
├── results_NIAH/
├── results_ruler/
├── scripts/
├── tests/
├── utils/
├── eval_longbench.py
├── eval_ruler.py
├── measure_latency.py
├── run_longbench.py
├── run_needle_in_haystack.py
├── run_ruler.py
├── topk_predictor.py
└── visualize_needle_in_haystack.py

Installation

# Clone the repository
git clone https://github.com/xiongsiheng/DHSA.git
cd DHSA

# Create and activate a conda environment with Python 3.12
conda create -n DHSA python=3.12 -y
conda activate DHSA

pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt
pip install flash-attn==2.7.4.post1 --no-build-isolation

# Build the Block-Sparse-Attention kernel
bash scripts/build_kernel.sh

Our default build supports A100 (sm_80), RTX 3090 (sm_86), and RTX 4090 (sm_89). Override the architecture list for a machine-specific build:

BLOCK_SPARSE_ATTN_CUDA_ARCHS=86 bash scripts/build_kernel.sh

The sparse kernel is BF16, causal, forward-only, and requires head dimension 128. The validated Transformers version is 4.45.0.

Predictor Checkpoints

The released predictor weights are provided in the DHSA collection.

Download the repositories into checkpoints/:

hf download sxiong/DHSA-Qwen2.5-3B-Instruct-BF16 \
  --local-dir checkpoints/DHSA-Qwen2.5-3B-Instruct-BF16

hf download sxiong/DHSA-Llama-3.1-8B-Instruct-BF16 \
  --local-dir checkpoints/DHSA-Llama-3.1-8B-Instruct-BF16

hf download sxiong/DHSA-Llama-3.1-8B-Instruct-NF4 \
  --local-dir checkpoints/DHSA-Llama-3.1-8B-Instruct-NF4

The downloaded directory structure is:

checkpoints/
├── DHSA-Qwen2.5-3B-Instruct-BF16/
│   ├── predictor_static.pt
│   └── predictor_dynamic.pt
├── DHSA-Llama-3.1-8B-Instruct-BF16/
│   ├── predictor_static.pt
│   └── predictor_dynamic.pt
└── DHSA-Llama-3.1-8B-Instruct-NF4/
    ├── predictor_static.pt
    └── predictor_dynamic.pt

All predictor checkpoints support both 6.25% and 12.5% density. Static checkpoints use a fixed row budget, while dynamic checkpoints use query-dependent row budgets and include density-specific budget configurations.

Quick Start

Needle-in-a-Haystack Test

To evaluate the model's ability to retrieve specific information from a long context, run:

bash scripts/run_needle_in_haystack_llama3.sh

bash scripts/run_needle_in_haystack_llama3_100k.sh

bash scripts/run_needle_in_haystack_qwen2_5.sh

To visualize the results, run:

bash scripts/visualize_NIAH_res.sh

LongBench

To test performance on the comprehensive LongBench suite, run:

bash scripts/run_longbench_llama3.sh

bash scripts/run_longbench_qwen2_5.sh

To evaluate the results, run:

bash scripts/eval_longbench.sh

RULER

To evaluate performance on the controlled synthetic RULER suite, you can either generate the data using the official repository or download our pre-generated data, then run:

bash scripts/run_ruler_llama3.sh

bash scripts/run_ruler_qwen2_5.sh

To evaluate the results, run:

bash scripts/eval_ruler.sh

Latency Measurement

To benchmark latency and memory usage against Flash Attention2, run:

bash scripts/measure_latency.sh

bash scripts/measure_latency_batched.sh

Acknowledgements

The implementation is built upon Block-Sparse Attention and KVCache-Factory. We sincerely appreciate these teams for their open-source contributions.

Citation

@inproceedings{xionglong,
  title={Long-Context Modeling with Dynamic Hierarchical Sparse Attention for Memory-Constrained LLM Inference},
  author={Xiong, Siheng and Zou, Joe and Fekri, Faramarz and Cho, Yae Jee},
  booktitle={Forty-third International Conference on Machine Learning}
}