HISA: 細粒度スパースアテンションのための効率的な階層的インデキシング

HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention

March 30, 2026

著者: Yufei Xu, Fanxu Meng, Fan Jiang, Yuxuan Wang, Ruijie Zhou, Jiexi Wu, Zhixin Pan, Zhaohui Wang, Xiaojuan Tang, Wenjie Pei, Tongxuan Liu, Di yin, Xing Sun, Muhan Zhang

cs.AI

要旨

トークンレベルのスパースアテンションメカニズムは、DeepSeek Sparse Attention（DSA）に代表されるように、軽量なインデクサを用いて各クエリに対し過去の全トークンをスコアリングし、選択されたサブセットのみでアテンション計算を行うことで、細粒度なキー選択を実現する。下流のスパースアテンションは効率的にスケールする一方で、インデクサは依然としてクエリごとに全文脈を走査するため、レイヤーあたりO(L^2)のボトルネックが生じ、文脈長の増大に伴い計算不能となる。本研究ではHISA（Hierarchical Indexed Sparse Attention）を提案する。これはインデクサのドロップイン代替品であり、検索プロセスを単一レベルのトークン走査から2段階の階層的手順に変換する。まず、ブロックレベルの粗いフィルタがプール化されたブロック代表値を評価し、無関係な領域を剪定する。次に、トークンレベルの精密化が残りの候補ブロック内でのみ元のインデクサを適用する。HISAは下流のSparse MLA演算子が要求する正確なトークンレベルのtop-kスパース性パターンを保持し、追加の学習を必要としない。カーネルレベルのベンチマークでは、HISAは32K文長で2倍、128K文長で4倍の高速化を達成した。Needle-in-a-HaystackとLongBenchでは、DeepSeek-V3.2のインデクサをファインチューニングなしで直接HISAに置換した。HISAは元のDSAと同等の品質を維持しつつ、ブロックスパースベースラインを大幅に上回った。さらに、HISAと元のDSAが生成するトークン選択セットは平均IoUが99%以上を示し、効率性向上が選択の忠実性に実質的に影響を与えないことを確認した。

English

Token-level sparse attention mechanisms, exemplified by DeepSeek Sparse Attention (DSA), achieve fine-grained key selection by scoring every historical token for each query using a lightweight indexer, and then computing attention only over the selected subset. While the downstream sparse attention scales efficiently, the indexer still scans the entire prefix for every query, introducing an O(L^2) per-layer bottleneck that becomes prohibitive as context length grows. We propose HISA (Hierarchical Indexed Sparse Attention), a drop-in replacement for the indexer that transforms the search process from a flat token scan into a two-stage hierarchical procedure. First, a block-level coarse filter scores pooled block representatives to prune irrelevant regions. Then, a token-level refinement applies the original indexer only within the remaining candidate blocks. HISA preserves the exact token-level top-k sparsity pattern required by the downstream Sparse MLA operator and requires no additional training. On kernel-level benchmarks, HISA achieves a 2times speedup at 32K context length and 4times at 128K. On Needle-in-a-Haystack and LongBench, we directly replace the indexer in DeepSeek-V3.2 with HISA, without any fine-tuning. HISA closely matches the original DSA in quality while significantly outperforming block-sparse baselines. Moreover, the token selection sets produced by HISA and the original DSA exhibit a mean IoU greater than 99%, indicating that the efficiency gains come with virtually no impact on selection fidelity.