动态分层稀疏注意力让大模型在低内存下高效处理长文本。
Long-Context Modeling with Dynamic Hierarchical Sparse Attention for Memory-Constrained LLM Inference
- 根据输入内容实时预测注意力稀疏模式，分块评估重要性并传递到词元级。
- 在128K上下文长度下，比块稀疏注意高出12%~20%相对准确率，预填充成本相当。
- 支持单卡24GB显存运行100K上下文，适配多硬件与多种开源模型。
注意力计算的二次复杂度限制了长上下文大模型在有限硬件内存下的可扩展性。尽管注意力本身具有稀疏性，现有静态稀疏方法无法适应任务或输入变化，而近期动态方法依赖预设模板或启发式规则，可能牺牲通用性。本文提出动态分层稀疏注意力（DHSA），一种数据驱动框架，在保持大模型主干冻结的前提下在线预测注意力稀疏性。DHSA通过分层路由机制，在块级别估计重要性并传播至词元级交互，保留因果关键依赖的同时实现高效稀疏化。在Needle-in-a-Haystack、LongBench和RULER测试中，DHSA在高度稀疏环境下仍保持接近稠密注意力的精度，相比块稀疏注意力在相似预填充成本下获得12%–20%相对准确率提升。结合内存高效的分块后端，DHSA在128K上下文长度下实现高达10倍的预填充加速。在LLaMA-3.1-8B（4-bit）上，仅需单张24GB显卡即可支持100K上下文，而稠密注意力在此条件下无法运行。我们提供互补的GPU与CPU后端，使DHSA可在多样硬件环境与多个开源模型家族中部署。结果表明，DHSA是内存受限场景下长上下文大模型推理的高效且自适应解决方案。
原文摘要 · Abstract (English)
The quadratic cost of attention limits the scalability of long-context LLMs, especially under limited hardware memory budgets. While attention is often sparse, existing static sparse methods cannot adapt to task- or input-dependent variations, and recent dynamic approaches rely on predefined templates or heuristics that may sacrifice generality. We propose Dynamic Hierarchical Sparse Attention (DHSA), a data-driven framework that predicts attention sparsity online while keeping the LLM backbone frozen. DHSA performs hierarchical routing by estimating importance at the chunk level and propagating it to token-level interactions, preserving causally important dependencies while enabling efficient sparsification. Across Needle-in-a-Haystack test, LongBench and RULER, DHSA maintains near-dense accuracy in highly sparse regimes, achieving 12--20% relative accuracy gains over Block Sparse Attention at comparable prefill cost. With a memory-efficient tiled backend, DHSA delivers up to $10\times$ prefill speedup at 128K context length. On LLaMA-3.1-8B (4-bit), DHSA scales to 100K context on a single 24GB GPU, where dense attention fails. We provide complementary GPU and CPU backends, enabling DHSA to run across diverse hardware environments and multiple open-weight model families. These results demonstrate DHSA as an efficient and adaptable solution for memory-constrained long-context LLM inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。