提出分层索引机制,让长文本注意力计算更快更省资源。
HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention
- 用分块粗筛+细粒度精筛两阶段替代逐字扫描,降低计算开销。
- 在64K上下文长度下提速最高达数倍,性能接近原始模型。
- 可直接替换现有模型索引器,无需重新训练,适合长文本任务。
基于标记级别的稀疏注意力机制(如DeepSeek Sparse Attention, DSA)通过轻量级索引器对每个查询评分所有历史键,仅计算选定子集的注意力。尽管下游稀疏注意力计算高效,但索引器需遍历完整前缀,随上下文长度增长而成为瓶颈。本文提出HISA(分层索引稀疏注意力),作为索引器的即插即用替代方案,将搜索路径从扁平标记扫描重构为两阶段分层过程:(1) 块级别粗筛选,对聚合块表示评分以剔除无关区域;(2) 标记级别精筛,仅在保留候选块内应用原索引器。HISA保持与下游稀疏多头注意力(Sparse MLA)一致的标记级顶稀疏模式,无需额外训练。在内核级基准测试中,HISA在64K上下文长度下实现最高提速。在Needle-in-a-Haystack和LongBench上,直接替换DeepSeek-V3.2与GLM-5中的索引器,无需微调,性能接近DSA,显著优于块稀疏基线。
原文摘要 · Abstract (English)
Token-level sparse attention mechanisms, exemplified by DeepSeek Sparse Attention (DSA), achieve fine-grained key selection by scoring every historical key for each query through a lightweight indexer, then computing attention only on the selected subset. While the downstream sparse attention itself scales favorably, the indexer must still scan the entire prefix for every query, introducing an per-layer bottleneck that grows prohibitively with context length. We propose HISA (Hierarchical Indexed Sparse Attention), a plug-and-play replacement for the indexer that rewrites the search path from a flat token scan into a two-stage hierarchical procedure: (1) a block-level coarse filtering stage that scores pooled block representations to discard irrelevant regions, followed by (2) a token-level refinement stage that applies the original indexer exclusively within the retained candidate blocks. HISA preserves the identical token-level top-sparse pattern consumed by the downstream Sparse MLA operator and requires no additional training. On kernel-level benchmarks, HISA achieves up to speedup at 64K context. On Needle-in-a-Haystack and LongBench, we directly replace the indexer in DeepSeek-V3.2 and GLM-5 with our HISA indexer, without any finetuning. HISA closely matches the original DSA in quality, while substantially outperforming block-sparse baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。