Pivotal方法让稀疏注意力索引提速4倍,同时保持精度不变。
PIVOT: Efficient Query-Group Indexing for Token-Level Sparse Attention

- 通过合并邻近查询为代理查询,仅扫描一次前缀提升效率
- 在长上下文场景下推理延迟降低1.6倍,最高加速4倍
- 无需训练、可直接替换现有系统,适合大模型部署优化
Token级稀疏注意力虽提升下游计算效率,但将瓶颈转移至索引器。为选出每查询的前k个候选,传统索引需对所有历史标记评分,导致每层复杂度达O(L²)。本文发现:相邻查询选择的高分标记高度重叠,且索引得分沿键轴呈长尾分布。基于此提出PIVOT——一种无需训练、可直接替换DSA索引器的方法。其将一组邻近查询聚合为单个代理查询,仅执行一次全前缀扫描生成候选集,再从该集中为每个查询选出前k项。两种变体兼顾速度与精度:PIVOT-Reuse共享代理查询的前k结果以最大化速度;PIVOT-Refine对候选集重新评分后独立选前k,代价略增但精度等同稠密索引。同一算法覆盖预填充和解码阶段:前者使用固定大小连续组,后者在多标记预测(MTP)中联合处理。在DeepSeek-V3.2和GLM-5.1上,结合LongBench和RULER测试,PIVOT保持稠密索引器精度的同时,实现最高4倍加速与1.6倍端到端延迟降低。
原文摘要 · Abstract (English)
Token-level sparse attention, as implemented by DeepSeek Sparse Attention (DSA) in production systems, makes the downstream attention efficient but shifts the bottleneck to the indexer that feeds it. To select the top-k tokens for each query, the indexer must still score every preceding token, incurring a cost of O(L^2) per layer for a sequence of length L. We observe that this per-query scan is largely redundant: nearby queries select highly overlapping top-k tokens, and the indexer scores are long-tailed along the key axis. We exploit these properties in PIVOT, Proxy Indexing Via One full-prefix Traversal, a training-free, drop-in replacement for the DSA indexer that shares one prefix scan across a group of nearby queries. PIVOT aggregates a group into a single proxy query, performs one shared full-prefix scan to obtain a candidate set, and then selects a top-k for each query from that set. Two variants trade speed for fidelity: PIVOT-Reuse shares the proxy top-k across the group for maximum speed, whereas PIVOT-Refine re-scores the candidate set with the indexer of each query and then selects an individual top-k, matching the dense indexer at a small additional cost. A single algorithm covers both inference phases, differing only in how groups are formed: fixed-size groups of consecutive queries in prefill, and the queries decoded together in one multi-token prediction (MTP) step in decode. On DeepSeek-V3.2 and GLM-5.1 across LongBench and RULER, PIVOT matches the accuracy of the dense DSA indexer while accelerating it by up to 4x and reducing end-to-end latency by up to 1.6x at long context.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。