HySparse: 오라클 토큰 선택 및 KV 캐시 공유를 갖는 하이브리드 희소 어텐션 아키텍처

HySparse: A Hybrid Sparse Attention Architecture with Oracle Token Selection and KV Cache Sharing

본 연구에서는 Hybrid Sparse Attention (HySparse)이라는 새로운 아키텍처를 소개합니다. HySparse는 각 전체 어텐션 레이어를 여러 개의 희소 어텐션 레이어로 번갈아 배치합니다. 개념적으로 간단하지만, HySparse는 각 희소 레이어의 토큰 선택 및 KV 캐시를 이전 전체 어텐션 레이어로부터 직접 파생시킵니다. 이러한 아키텍처는 기존의 희소 어텐션 방법의 두 가지 근본적인 한계를 해결합니다. 첫째, 기존 방식은 일반적으로 토큰의 중요도를 예측하기 위한 추가적인 프록시를 사용하며, 이는 추가적인 복잡성을 야기하고 잠재적으로 최적의 성능을 달성하지 못할 수 있습니다. 반면에, HySparse는 전체 어텐션 레이어를 중요한 토큰을 식별하는 정확한 '오라클'로 활용합니다. 둘째, 기존의 희소 어텐션 설계는 종종 계산량을 줄이지만 KV 캐시를 절약하지 못합니다. HySparse는 희소 어텐션 레이어가 전체 어텐션의 KV 캐시를 재사용할 수 있도록 하여 계산량과 메모리 모두를 줄입니다. 우리는 7B 밀집 모델과 80B MoE 모델 모두에서 HySparse를 평가했습니다. 모든 설정에서 HySparse는 전체 어텐션 및 하이브리드 SWA 기준 모델보다 일관되게 우수한 성능을 보였습니다. 특히, 49개의 총 레이어를 가진 80B MoE 모델에서 전체 어텐션을 사용하는 레이어가 5개에 불과했지만, HySparse는 상당한 성능 향상을 달성하면서 KV 캐시 저장 공간을 거의 10배 줄였습니다.

This work introduces Hybrid Sparse Attention (HySparse), a new architecture that interleaves each full attention layer with several sparse attention layers. While conceptually simple, HySparse strategically derives each sparse layer's token selection and KV caches directly from the preceding full attention layer. This architecture resolves two fundamental limitations of prior sparse attention methods. First, conventional approaches typically rely on additional proxies to predict token importance, introducing extra complexity and potentially suboptimal performance. In contrast, HySparse uses the full attention layer as a precise oracle to identify important tokens. Second, existing sparse attention designs often reduce computation without saving KV cache. HySparse enables sparse attention layers to reuse the full attention KV cache, thereby reducing both computation and memory. We evaluate HySparse on both 7B dense and 80B MoE models. Across all settings, HySparse consistently outperforms both full attention and hybrid SWA baselines. Notably, in the 80B MoE model with 49 total layers, only 5 layers employ full attention, yet HySparse achieves substantial performance gains while reducing KV cache storage by nearly 10x.