자기 인덱싱 KVCache: 압축된 키로부터 희소 어텐션을 예측하는 방법

Self-Indexing KVCache: Predicting Sparse Attention from Compressed Keys

자기 어텐션에서 KV 캐시는 LLM의 장문 컨텍스트 및 대규모 배치 추론에서 주요 병목 현상으로 등장했습니다. 기존 접근 방식은 종종 희소성 예측과 압축을 별개의 모듈로 처리하며, 관련 토큰을 선택하기 위해 보조 인덱스 구조에 의존하고, 메모리 사용량을 줄이기 위해 복잡한 양자화 방식을 사용합니다. 이러한 분산된 설계는 불필요한 오버헤드를 발생시키고 확장성을 제한합니다. 본 논문에서는 새로운 패러다임을 제안합니다. 압축된 키 표현을 단순히 저장 공간으로 취급하는 것이 아니라, 효율적인 희소 어텐션을 직접적으로 가능하게 하는 자기 인덱싱 구조로 활용하는 것입니다. 부호 기반 1비트 벡터 양자화(VQ) 방식을 설계하여, 저희 방법은 압축과 검색을 단일하고 하드웨어 친화적인 형식으로 통합합니다. 이러한 접근 방식은 외부 인덱스나 학습 기반 예측기의 필요성을 없애고, 메모리 제약이 있는 추론을 위한 가볍고 견고한 솔루션을 제공합니다. 모든 구성 요소는 하드웨어 효율성을 고려하여 설계되었으며 구현이 용이합니다. 맞춤형 CUDA 커널을 구현하여, 저희 방법은 FlashAttention과 완벽하게 통합되어 추가적인 런타임 및 메모리 오버헤드를 최소화합니다. 실험 결과는 저희 접근 방식이 효과성과 효율성을 모두 제공한다는 것을 보여줍니다.

The KV cache in self-attention has emerged as a major bottleneck in long-context and large-batch inference for LLMs. Existing approaches often treat sparsity prediction and compression as separate modules, relying on auxiliary index structures to select relevant tokens, and on complex quantization schemes to reduce memory usage. This fragmented design introduces redundant overhead and limits scalability. In this paper, we propose a novel paradigm: treating the compressed key representation not merely as storage, but as a self-indexing structure that directly enables efficient sparse attention. By designing a sign-based 1-bit vector quantization (VQ) scheme, our method unifies compression and retrieval in a single, hardware-friendly format. This approach eliminates the need for external indices or learning-based predictors, offering a lightweight yet robust solution for memory-constrained inference. All components are designed to be hardware-efficient and easy to implement. By implementing custom CUDA kernels, our method integrates seamlessly with FlashAttention, minimizing additional runtime and memory overhead. Experimental results demonstrate that our approach delivers both effectiveness and efficiency.