HISA: 미세한 희소 어텐션(Sparse Attention)을 위한 효율적인 계층적 인덱싱

HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention

토큰 레벨의 희소 어텐션 메커니즘은 DeepSeek Sparse Attention (DSA)와 같이, 각 쿼리에 대해 모든 과거 키에 대한 점수를 계산하여 가벼운 인덱서를 사용하고, 선택된 부분집합에 대해서만 어텐션을 계산하여 미세한 수준의 키 선택을 달성합니다. 다운스트림의 희소 어텐션 자체는 확장성이 좋지만, 인덱서는 여전히 각 쿼리에 대해 전체 접두사를 스캔해야 하며, 이는 컨텍스트 길이가 증가함에 따라 비례하여 증가하는 레이어별 병목 현상을 유발합니다. 본 논문에서는 HISA (Hierarchical Indexed Sparse Attention)를 제안합니다. HISA는 기존 인덱서를 대체하는 플러그 앤 플레이 방식으로, 평면적인 토큰 스캔을 두 단계의 계층적 절차로 재구성합니다. (1) 블록 레벨의 초기 필터링 단계에서는 풀링된 블록 표현에 대한 점수를 계산하여 관련 없는 영역을 제거하고, (2) 토큰 레벨의 세밀 조정 단계에서는 보존된 후보 블록 내에서 원래 인덱서를 적용합니다. HISA는 다운스트림의 Sparse MLA 연산자가 사용하는 동일한 토큰 레벨의 최상위 희소 패턴을 유지하며, 추가적인 학습이 필요하지 않습니다. 커널 레벨 벤치마크에서 HISA는 64K 컨텍스트에서 최대 속도 향상을 달성합니다. Needle-in-a-Haystack 및 LongBench 데이터셋에서, 우리는 DeepSeek-V3.2 및 GLM-5의 인덱서를 직접적으로 HISA 인덱서로 대체했으며, 추가적인 파인튜닝은 수행하지 않았습니다. HISA는 원래 DSA와 유사한 품질을 유지하면서, 블록 희소(block-sparse) 기반 모델보다 훨씬 뛰어난 성능을 보입니다.

Token-level sparse attention mechanisms, exemplified by DeepSeek Sparse Attention (DSA), achieve fine-grained key selection by scoring every historical key for each query through a lightweight indexer, then computing attention only on the selected subset. While the downstream sparse attention itself scales favorably, the indexer must still scan the entire prefix for every query, introducing an per-layer bottleneck that grows prohibitively with context length. We propose HISA (Hierarchical Indexed Sparse Attention), a plug-and-play replacement for the indexer that rewrites the search path from a flat token scan into a two-stage hierarchical procedure: (1) a block-level coarse filtering stage that scores pooled block representations to discard irrelevant regions, followed by (2) a token-level refinement stage that applies the original indexer exclusively within the retained candidate blocks. HISA preserves the identical token-level top-sparse pattern consumed by the downstream Sparse MLA operator and requires no additional training. On kernel-level benchmarks, HISA achieves up to speedup at 64K context. On Needle-in-a-Haystack and LongBench, we directly replace the indexer in DeepSeek-V3.2 and GLM-5 with our HISA indexer, without any finetuning. HISA closely matches the original DSA in quality, while substantially outperforming block-sparse baselines.