PIVOT:令牌级稀疏注意力的高效查询组索引方法

PIVOT: Efficient Query-Group Indexing for Token-Level Sparse Attention

精选理由

PIVOT用共享前缀扫描替代逐查询扫描,加速4倍,端到端延迟降1.6倍,效果不打折。适合长上下文推理场景。

AI 摘要

DeepSeek稀疏注意力(DSA)在生成式系统中实现令牌级稀疏注意力,但索引器仍需对所有前序token打分,复杂度为O(L²)。PIVOT利用附近查询的token重叠性和索引器得分长尾分布,通过共享前缀扫描加速。PIVOT将一组查询聚合成一个代理查询,进行一次全前缀扫描得到候选集,再为每个查询选择top-k。两个变体:PIVOT-Reuse共享代理top-k,PIVOT-Refine重新评分候选集。在DeepSeek-V3.2和GLM-5.1上,在LongBench和RULER基准上,PIVOT匹配精度,索引器加速4倍,端到端延迟降低1.6倍。

原文 · arXiv: DeepSeek

PIVOT: Efficient Query-Group Indexing for Token-Level Sparse Attention

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.