PIVOT: 토큰 레벨 희소 어텐션의 효율적인 쿼리 그룹 인덱싱

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

딥시크 희소 어텐션(DSA)과 같이 실제 시스템에서 구현되는 토큰 레벨 희소 어텐션은 다운스트림 어텐션을 효율적으로 만들지만, 이를 위한 인덱서에 병목 현상이 발생합니다. 각 쿼리에 대해 상위 k개의 토큰을 선택하기 위해, 인덱서는 여전히 모든 이전 토큰에 대한 점수를 매겨야 하며, 이는 시퀀스 길이 L에 대해 레이어당 O(L^2)의 비용을 초래합니다. 우리는 이 쿼리별 스캔이 대부분 불필요하다는 것을 관찰했습니다. 왜냐하면 인접한 쿼리는 매우 중복되는 상위 k개의 토큰을 선택하며, 인덱서 점수는 키 축을 따라 긴 꼬리를 갖기 때문입니다. 우리는 PIVOT(Proxy Indexing Via One full-prefix Traversal)에서 이러한 특성을 활용합니다. PIVOT는 학습이 필요 없는 DSA 인덱서를 대체하는 방식으로, 일련의 인접한 쿼리에 대해 하나의 전역 스캔을 공유합니다. PIVOT는 그룹을 단일 프록시 쿼리로 집계하고, 하나의 공유된 전체-프리픽스 스캔을 수행하여 후보 집합을 얻은 다음, 각 쿼리에 대해 해당 집합에서 상위 k개를 선택합니다. 두 가지 변형이 속도와 정확성 간의 균형을 맞춥니다. PIVOT-Reuse는 그룹 내에서 프록시 상위 k개를 공유하여 최대 속도를 제공하는 반면, PIVOT-Refine은 각 쿼리의 인덱서를 사용하여 후보 집합을 다시 평가한 다음 개별 상위 k개를 선택하여 밀집 인덱서와 거의 동일한 성능을 제공하지만 약간의 추가 비용이 발생합니다. 단일 알고리즘이 추론 단계 모두를 처리하며, 그룹 형성 방식만 다릅니다. 사전 채우기에서는 고정 크기의 연속된 쿼리 그룹을 사용하고, 디코딩에서는 멀티-토큰 예측(MTP) 단계에서 함께 디코딩되는 쿼리를 사용합니다. DeepSeek-V3.2 및 GLM-5.1 모델에서 LongBench 및 RULER 데이터셋에 대해 PIVOT는 밀집 DSA 인덱서와 동일한 정확도를 달성하면서 최대 4배의 속도 향상과 긴 컨텍스트 환경에서 최대 1.6배의 엔드투엔드 지연 감소를 제공합니다.

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.