CentroidKV: Efficient Long-Context LLM Inference via KV Cache Clustering

기술 요약: CentroidKV

문제 정의

최대 100만 토큰에 달하는 확장된 컨텍스트 창을 가진 대규모 언어 모델(LLM)은 키-값(KV) 캐시로 인해 심각한 배포 문제를 겪고 있습니다. 이전 토큰들의 키(Key)와 값(Value)을 저장하여 재계산을 방지하는 KV 캐시는 컨텍스트 길이에 따라 선형적으로 증가합니다. 이는 모든 생성 토큰마다 전체 캐시에 접근해야 하므로 추론 지연 시간과 처리량의 병목 현상을 초래합니다. 또한, 거대한 메모리 점유율은 배포를 더욱 복잡하게 만듭니다.

KV 캐시 크기를 줄이기 위한 기존 솔루션은 다음과 같습니다:

• 증발(Eviction): 어텐션 점수를 기반으로 덜 중요한 토큰을 버리는 방식입니다. 이는 처음에 중요하지 않다고 판단된 토큰이 향후 생성 과정에서 매우 중요해질 수 있기 때문에 성능 저하를 초래할 수 있습니다.
• 병합(Merging): 증발될 토큰들을 유지될 토큰들과 결합하는 방식입니다. 그러나 기존 방법들은 종종 증발을 먼저 수행한 후 병합 집합을 정의하기 때문에, 의미론적으로 먼 토근들이 그룹화되어 센트로이드 편향(centroid bias)과 정보 손실을 유발합니다.
• 양자화 및 프루닝(Quantization and Pruning): 이들은 제한적인 효율성 이득을 제공하거나 높은 계산 오버헤드를 발생시킵니다.

따라서 중요한 정보를 버리거나 과도한 계산 비용을 발생시키지 않으면서, 내재적 유사성에 기반하여 KV 캐시를 압축할 수 있는 방법이 필요합니다.

방법론: CentroidKV

본 논문은 온라인 KV 캐시 클러스터링 프레임워크인 CentroidKV를 소개합니다. 이 모델의 핵심 혁신은 낮은 계산 복잡도를 유지하면서 유사성에 기반해 KV 상태를 클러스터링하도록 설계된 청크 단위 소프트 매칭(Chunked Soft Matching, CSM) 알고리즘입니다.

주요 구성 요소

1. 관찰 사항:

   • 국소적 유사성(Localized Similarity): 키 상태는 시퀀스 차원을 따라 높은 코사인 유사성을 보이며, 유사한 토큰들은 국소적인 영역에 클러스터링됩니다.
   • 거리 상관관계(Distance Correlation): 토큰 간 거리가 멀어질수록 키 상태의 코사인 유사성은 일반적으로 단조 감소하며 볼록(convex)한 경향을 따릅니다.

2. 추론 파이프라인:

   • 프리필링(Pre-filling): 시퀀스가 처리되며, KV 캐시가 예산을 초과하는 경우 첫 번째 클러스터링 라운드가 실행됩니다.
   • 디코딩(Decoding): 캐시는 각 토큰과 함께 성장합니다. 캐시를 예산 내로 유지하기 위해 주기적으로($g$ 단계마다) 클러스터링이 호출됩니다.
   • 보존(Preservation): 모델의 안정성을 보장하기 위해 파티셔닝 전에 어텐션 싱크(Attention sinks)와 최근 토큰들을 보존합니다.

3. 청크 단위 소프트 매칭 (CSM):

   • 청킹(Chunking): 긴 컨텍스트를 효율적으로 처리하기 위해 시퀀스를 청크로 나눕니다.
   • 교차 파티션(Alternating Partition): 각 청크 내에서 토큰들을 교차 방식으로 두 개의 집합($A$와 $B$)으로 분할합니다. 이론적 분석에 따르면, 관찰된 단조 및 볼록 유사성 경향을 고려할 때 이 전략은 최적이며, 매우 유사한 토큰들이 서로 다른 집합에 배치되도록 하여 매칭을 용이하게 합니다.
   • 매칭(Matching): 비전 트랜스포머(Vision Transformer)에서 사용되는 이분 소프트 매칭(Bipartite Soft Matching, BSM)과 유사하게, CSM은 집합 $A$와 $B$ 사이에서 매우 유사한 토큰 쌍을 식별합니다.
   • 압축(Compression): 모든 쌍을 병합하는 대신, 후보 매치들을 유사도 순으로 정렬하고 압축 비율 $r$에 의해 제어되는 상위 부분만을 병합합니다.
   • 점진적 스케줄링(Progressive Scheduling): 압축 비율 $r$은 여러 클러스터링 라운드에 걸쳐 선형적으로 어닐링됩니다($r = r_{init} - j \cdot \delta_r$). 이를 통해 후기 라운드로 갈수록 병합을 더욱 선택적으로 수행하여 고신뢰 구조를 보존합니다.

4. 집계(Aggregation):

   • 각 클러스터에 대해, 키와 값 상태는 차수 가중 평균(degree-weighted averaging)을 사용하여 하나의 센트로이드(centroid)로 병합됩니다. 여기서 "차수"($n_t$)는 여러 라운드에 걸쳐 센로이드로 병합된 토큰들의 누적 기여도를 추적합니다.
   • 어텐션 메커니즘은 원래의 어텐션 출력을 근사하기 위해 클러스터 차수를 소프트맥스 계산의 로그 항으로 포함하도록 조정됩니다.

주요 기여

• CentroidKV 프레임워크: 긴 컨텍스트 시나리오에서 추론 효율성을 개선하는 단순하고 플러그 앤 플레이(plug-and-play)가 가능한 온라인 KV 캐시 클러스터링 솔루션입니다.
• 청크 단위 소프트 매칭 알고리즘: Bipartite Soft Matching을 KV 캐시 압축으로 확장한 새로운 클러스터링 알고리즘입니다. 이는 청크 내에서 교차 파티션 전략을 활용하며, 이는 관찰된 유사성 패턴에 대해 최적임이 이론적으로 증명되었습니다.
• 이론적 분석: 본 논문은 청크 내 파티셔닝 전략의 최적성에 대한 공식적인 증명을 제공하며, CSM이 K-Means($O(inkd)$)나 전체 쌍 메쉬 방식($O(n^2d)$)보다 낮은 복잡도($O(ncd)$)를 달성함을 입증합니다.
• 실증적 검증: RULER 및 LongBench 벤치마크에서 Llama-3.1-8B 및 Mistral-7B 모델을 대상으로 광범위한 실험을 수행했습니다.

결과

• 메모리 효율성: CentroidKV는 모델 성능을 유지하면서 KV 캐시 메모리 사용량을 최대 75% 감소(캐시 크기를 약 4배 축소)시킵-니다.
• 추론 속도:
  • 디코딩 지연 시간: 디코딩 단계의 속도를 최대 1.92배 가속화합니다(출력 토큰당 시간 감소로 측정).
  • 처리량: vLLM 프레임워크에 통합되었을 때 서빙 처리량을 최대 4배 증가시킵니다.
• 정확도:
  • RULR 벤치마크에서 CentroidKV는 다양한 캐시 예산(25%~75%)에 대해 최고 또는 최상위권의 성능을 달성하며, 특히 의미론적 이해가 필요한 작업(예: S-NIAH-1/2, MK-NIAH-1, VT)에서 뛰어난 성능을 보입니다.
  • LongBench에서는 Single-Doc QA, Multi-Doc QA 및 요약 작업에서 일관되게 우수한 성능을 보입니다.
  • 한계점: 정확한 토큰 매칭이 필수적인 합성(Synthetic) 작업 및 미세한 검색 작업(예: UUID 기반 키)에서는 성능이 약간 하락합니다. 저자들은 내용 인식(content-aware) 방식이 위치 기반 방식인 StreamingLLM에 비해 의미론적으로 임의적인 랜덤 문자열 처리에 어려움을 겪는다는 점을 언급했습니다.

의의 및 주장

본 논문은 CentroidKV가 정확도와 효율성 사이에서 유리한 절충안을 제공한다고 주장합니다. 키 상태의 내재적 유사성을 활용하고 경량화된 단일 패스 클러스터링 알고리즘을 사용함으로써, 단순 증발(naive eviction)의 성능 저하나 K-Means와 같은 반복적 클러스터링의 높은 계산 오버헤드를 피할 수 있습니다.

저자들은 CentroidKV가 모델 재학습이 필요 없는 경량화된 플러그 앤 플레이 솔루션임을 강조합니다. 이는 긴 컨텍스트 LLM의 메모리 병목 현상을 효과적으로 해결하여, 모델의 장거리 의존성 이해 능력을 크게 훼손하지 않으면서도 더 빠른 디코딩과 더 높은 처리량을 가능하게 합니다.

저자가 언급한 한계점:

• 현재 구현은 GPU 기반 압축에 집중되어 있으며, CPU 클러스터링과 같은 메모리 오프로딩 전략은 탐구하지 않았습니다.
• 압축 비율 스케줄은 수동으로 지정되며, 방식은 각 단계에서 캐시 감소를 최대 절반까지만 제한합니다. 향후 연구에서는 적응형 압축 전략을 탐구할 수 있습니다.