You Only Index Once: 공유 라우팅을 통한 교차 레이어 희소 어텐션 (Cross-Layer Sparse Attention)

현대 LLM (Large Language Models)의 긴 문맥 추론 (Long-context inference)은 디코딩 효율성에 의해 점점 더 제약을 받고 있으며, 특히 모델이 긴 중간 사고 사슬 (chains of thought)을 생성하는 추론 중심의 환경에서 더욱 그러합니다. 기존의 희소 어텐션 (sparse attention) 방법들은 종종 실질적인 효율성과 품질 사이의 트레이드오프 (trade-off) 문제에 직면합니다. 구조화된 블록 희소 (Structured block sparse) 방법들은 일반적으로 더 강력한 가속을 제공하지만 눈에 띄는 품질 저하를 초래하는 반면, 토큰 희소 (token sparse) 방법들은 대개 더 정확하지만 전체 캐시 (cache)에 대한 top-k 라우팅 (routing) 비용이 여전히 높기 때문에 제한적인 엔드투엔드 (end-to-end) 속도 향상만을 제공합니다.

본 연구에서는 YOCO와 같은 KV 공유 아키텍처를 기반으로 구축된 교차 레이어 희소 어텐션 (CLSA, Cross-Layer Sparse Attention)을 제안합니다. 핵심 아이디어는 교차 디코더 레이어 (cross-decoder layers) 간에 KV 캐시 (KV cache)를 공유할 뿐만 아니라, 라우팅 인덱스 (routing index) 또한 공유하는 것입니다. 단일 인덱서 (indexer)가 토큰 수준의 top-k 선택을 한 번만 계산하고 결과로 나온 인덱스를 레이어 전반에 걸쳐 재사용함으로써, 라우팅 오버헤드 (routing overhead)를 분할 amortizing 하는 동시에 토큰 희소 어텐션의 미세한 선택성 (fine-grained selectivity)을 보존합니다.

결과적으로 이 아키텍처는 프리필링 (pre-filling), KV 캐시 저장 (KV-cache storage), 긴 문맥 디코딩 (long-context decoding)을 포함한 모든 주요 추론 병목 현상을 공동으로 개선합니다. 짧은 문맥 및 긴 문맥 벤치마크 전반에 걸친 실험 결과, CLSA는 정확하면서도 효율적임을 보여주었으며, 128K 문맥에서 최대 7.6배의 디코딩 속도 향상과 17.1배의 전체 처리량 (throughput) 개선을 달성했습니다. 이러한 결과는 모델 품질과 추론 효율성을 동시에 발전시키는 긴 문맥 LLM을 위한 더욱 완전한 아키텍처 솔루션을 시사합니다.