AttentionIndex 用倒排索引加速长上下文稀疏注意力,吞吐最高提升 1.8 倍
AttentionIndex: Accelerating Learned Sparse Attention with a Causal Inverted-File Index
解决 sparse attention 索引慢的老问题,百万 token 下 prefill 快 1.3-1.8 倍,召回还比 HiSA 高十几个百分点
DeepSeek-V4、GLM-5 等生产模型使用的 learned sparse attention 中,索引器仍需给前缀中所有 key 打分,在 100 万 token 上下文时选择阶段占 prefill 时间一半以上。AttentionIndex 通过倒排文件索引对 indexer key 按相似度分组路由查询,在 prefill 期间构建索引并复用先前聚类。在 GLM-5.3 和 DeepSeek-V4-Flash 的 100 万 token 测试中,prefill 吞吐达到穷举选择的 1.3-1.8 倍。在与 HiSA 吞吐相差 5% 以内的条件下,top-k 召回分别多恢复 17.0 和 10.0 个百分点,长上下文基准准确率最多高出 18.4 个百分点。
AttentionIndex: Accelerating Learned Sparse Attention with a Causal Inverted-File Index
Learned sparse attention, used in production models such as DeepSeek-V4 and GLM-5, lets each query attend to only k keys chosen by a lightweight indexer, but the indexer still scores every key in the prefix. At one million tokens, this selection takes more than half of prefill time. Two-stage selection over contiguous blocks, as in HiSA, reduces this cost, but its recall declines as context grows at a fixed candidate budget. We present AttentionIndex, which routes each query over an inverted-file index that groups indexer keys by similarity and refines the candidates with the native indexer score. AttentionIndex builds the index during prefill, making new keys searchable immediately, reusing previous clustering to keep maintenance cheap, and exposing each query only to clusters built from its own prefix. On GLM-5.3 and DeepSeek-V4-Flash at one million tokens, AttentionIndex delivers 1.3-1.8 times the prefill throughput of exhaustive selection, and at throughput within 5% of HiSA, it recovers 17.0 and 10.0 percentage points more of the exhaustive top-k, respectively. AttentionIndex also improves accuracy on long-context benchmarks by up to 18.4 percentage points over HiSA at the same candidate budget.