HISA: 세밀한 희소 주의를 위한 효율적 계층적 인덱싱

초록

토큰 수준 희소 주의 메커니즘은 DeepSeek 희소 주의(DSA)에서 보듯이, 경량 인덱서를 사용하여 각 쿼리마다 모든 과거 토큰에 점수를 매기고 선택된 부분 집합에 대해서만 주의를 계산함으로써 세밀한 키 선택을 달성합니다. 다운스트림 희소 주의는 효율적으로 확장되지만, 인덱서는 여전히 모든 쿼리에 대해 전체 접두사를 스캔하여 계층당 O(L^2) 병목 현상을 초래하며, 이는 컨텍스트 길이가 증가함에 따라 치명적으로 됩니다. 우리는 HISA(계층적 인덱싱 희소 주의)를 제안합니다. 이는 인덱서를 대체하는 드롭인 솔루션으로, 검색 과정을 평면적인 토큰 스캔에서 두 단계의 계층적 절차로 전환합니다. 먼저, 블록 수준의 코스 필터가 풀링된 블록 대표값에 점수를 매겨 관련 없는 영역을 제거합니다. 그런 다음, 토큰 수준 정제 단계에서 남은 후보 블록 내에서만 원본 인덱서를 적용합니다. HISA는 다운스트림 희소 MLA 연산자가 요구하는 정확한 토큰 수준 상위 k개 희소 패턴을 유지하며 추가 학습이 필요하지 않습니다. 커널 수준 벤치마크에서 HISA는 32K 컨텍스트 길이에서 2배, 128K에서 4배의 속도 향상을 달성했습니다. Needle-in-a-Haystack 및 LongBench에서 우리는 DeepSeek-V3.2의 인덱서를 어떠한 미세 조정 없이 HISA로 직접 대체했습니다. HISA는 원본 DSA의 품질을 근접하게 유지하면서 블록 희소 베이스라인을 크게 능가합니다. 더욱이, HISA와 원본 DSA가 생성한 토큰 선택 집합은 99% 이상의 평균 IoU를 보여주어, 효율성 향상이 선택 정확도에 거의 영향을 미치지 않음을 입증합니다.

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.

HISA: 세밀한 희소 주의를 위한 효율적 계층적 인덱싱

HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention

초록

Support