단 한번의 인덱싱: 공유 라우팅을 통한 레이어 간 희소 어텐션

You Only Index Once: Cross-Layer Sparse Attention with Shared Routing

최신 대규모 언어 모델(LLM)에서 긴 문맥 추론은 디코딩 효율성, 특히 모델이 복잡한 사고 과정을 거치는 추론 중심 환경에서 점점 더 큰 제약 요인이 되고 있습니다. 기존의 희소 어텐션 방법들은 종종 실질적인 효율성과 품질 간의 균형 문제를 안고 있습니다. 구조화된 블록 희소 방식은 일반적으로 더 빠른 속도를 제공하지만 눈에 띄는 품질 손실을 야기하는 반면, 토큰 희소 방식은 보통 정확도가 높지만 전체 시스템 속도 향상이 제한적입니다. 왜냐하면 전체 캐시에 대한 top-k 라우팅 연산이 여전히 비용이 많이 들기 때문입니다. 본 연구에서는 KV-sharing 아키텍처(예: YOCO)를 기반으로 레이어 간 희소 어텐션(CLSA)을 제안합니다. 핵심 아이디어는 디코더 레이어 전체에서 KV 캐시뿐만 아니라 라우팅 인덱스까지 공유하는 것입니다. 단일 인덱서가 토큰 수준의 top-k 선택을 한 번 수행하고, 그 결과를 여러 레이어에 재사용하여 토큰 희소 어텐션의 미세한 선택성을 유지하면서 라우팅 오버헤드를 줄입니다. 결과적으로 얻어지는 아키텍처는 사전 채움(pre-filling), KV 캐시 저장 공간, 그리고 긴 문맥 디코딩을 포함한 주요 추론 병목 현상을 동시에 개선합니다. 짧은 문맥과 긴 문맥 벤치마크를 통해 실험한 결과, CLSA는 정확성과 효율성 모두에서 우수한 성능을 보이며, 128K 문맥 길이에서 최대 7.6배의 디코딩 속도 향상 및 전체 처리량 17.1배 개선 효과를 달성했습니다. 이러한 결과는 긴 문맥 LLM에 대한 보다 완벽한 아키텍처 솔루션을 제시하며, 모델 품질과 추론 효율성을 동시에 향상시킬 수 있음을 시사합니다.

Long-context inference in modern LLMs is increasingly constrained by decoding efficiency, especially in reasoning-heavy settings where models generate long intermediate chains of thought. Existing sparse attention methods often face a practical efficiency-quality trade-off. Structured block sparse methods typically provide stronger acceleration but incur noticeable quality loss, while token sparse methods are usually more accurate yet deliver limited end-to-end speedup because top-k routing over the full cache remains expensive. In this work, we propose cross-layer sparse attention (CLSA), which is built on top of KV-sharing architectures such as YOCO. The core idea is to share not only the KV cache across cross-decoder layers, but also the routing index. A single indexer computes token-level top-k selection once and reuses the resulting index across layers, thereby preserving the fine-grained selectivity of token sparse attention while amortizing the routing overhead. The resulting architecture improves all major inference bottlenecks jointly, including pre-filling, KV-cache storage, and long-context decoding. Experiments across short-context and long-context benchmarks show that CLSA is both accurate and efficient, achieving up to 7.6x decoding speedup and 17.1x overall throughput improvement at 128K context. These results suggest a more complete architectural solution for long-context LLMs that jointly advances model quality and inference efficiency.