NOSA: 네이티브 및 오프로드 가능한 희소 어텐션

초록

훈련 가능한 희소 주의 메커니즘은 장문맥 처리에서 대형 언어 모델(LLM)의 디코딩 효율성 병목 현상을 해결하기 위한 유망한 솔루션으로 부상하며, 작업 성능에 최소한의 영향을 미치면서 메모리 접근을 크게 절약합니다. 그러나 기존의 희소 주의 메커니즘은 중요한 한계를 해결하지 못하고 있습니다: 키-값(KV) 캐시의 크기가 줄어들지 않아 GPU 상의 배치 크기가 제한되고, 특히 대규모 배치 추론에서 디코딩 처리량이 저하됩니다. 본 논문에서는 훈련 가능한 희소 주의 메커니즘이 인접한 디코딩 단계 간 토큰 선택에서 강력한 지역성을 자연스럽게 나타내어, 기본적인 주의 계산을 변경하지 않고도 KV 캐시 오프로딩을 가능하게 함을 보여줍니다. 그러나 내재된 지역성만으로는 효율적인 오프로딩을 달성하기에 부족하며, 선택된 KV 쌍의 CPU와 GPU 간 전송이 전체 디코딩 비용을 지배하는 문제가 여전히 남아 있습니다. 이러한 통찰을 바탕으로, 본 논문은 KV 캐시 오프로딩을 기본적으로 지원하도록 설계된 훈련 가능한 희소 주의 프레임워크인 NOSA를 제안합니다. NOSA는 토큰 선택을 쿼리 인식 및 쿼리 무관 구성 요소로 분해하여 명시적인 지역성 제약을 도입함으로써 KV 전송을 줄이면서도 훈련 중 사용된 동일한 주의 계산을 보존합니다. 우리는 NOSA를 사용하여 10억 파라미터 모델을 사전 훈련하고 광범위한 벤치마크를 수행하여, 기존의 훈련 가능한 희소 주의 메커니즘 기준(InfLLM-V2) 대비 최대 2.3배의 디코딩 처리량 향상을 달성하면서도 거의 손실 없는 성능을 유지함을 보여줍니다.

English

Trainable sparse attention has emerged as a promising solution to address the decoding efficiency bottleneck of LLMs in long-context processing, significantly saving memory accesses while minimally impacting task performance. However, existing sparse attention methods leave a crucial limitation unresolved: the size of the key-value (KV) cache remains unreduced, which constrains on-GPU batch sizes and throttles decoding throughput, especially in large-scale batched inference. In this paper, we show that trainable sparse attention naturally exhibits strong locality in token selection across adjacent decoding steps, thereby enabling KV cache offloading without altering the underlying attention computation. However, the inherent locality remains insufficient to achieve efficient offloading, as the transfer of selected KV pairs between the CPU and GPU continues to dominate the overall decoding cost. Building on this insight, we present NOSA, a trainable sparse attention framework designed to natively support KV cache offloading. NOSA introduces explicit locality constraints by decomposing token selection into query-aware and query-agnostic components, thereby reducing KV transfers while preserving the same attention computation as used during training. We pretrain a 1B-parameter model with NOSA and conduct extensive benchmarks, showing that it preserves near-lossless performance while achieving up to a 2.3x improvement in decoding throughput compared with the vanilla trainable sparse attention baseline (InfLLM-V2).

NOSA: 네이티브 및 오프로드 가능한 희소 어텐션

NOSA: Native and Offloadable Sparse Attention

초록

Support