기술 요약: RDKV - KV 캐시 공동 추방 및 양자화를 위한 왜곡-비트 할당

1. 문제 제기

대형 언어 모델 (LLM) 은 긴 문맥 추론 중에 메모리 대역폭과 용량이라는 치명적인 병목 현상에 직면합니다. 키 - 값 (KV) 캐시 크기는 시퀀스 길이에 비례하여 선형적으로 증가하므로, 매 디코딩 단계마다 온칩 메모리로 온칩 메모리에서 수십 기가바이트의 데이터를 다시 읽어야 합니다. 이로 인해 추론은 계산 중심이 아닌 메모리 중심이 됩니다.

기존 압축 전략은 일반적으로 두 가지 고립된 범주로 나뉩니다:

1. 추방 (Eviction): 토큰이나 채널에 점수를 매겨 가장 낮은 순위의 것을 완전히 폐기 (이진 유지 또는 폐기) 합니다. 이는 중간 중요도의 토큰을 위한 정밀도 계층을 남기지 않습니다.
2. 양자화 (Quantization): 모든 토큰을 유지하지만 비트 너비를 줄입니다. 이는 미미한 토큰까지 유지하여 메모리를 낭비합니다.

최근의 하이브리드 접근법은 이들을 결합하려 시도하지만, 종종 오프라인 검색을 통해 결정된 휴리스틱 라우팅 점수나 사전 할당된 비율에 의존하며, 추방과 양자화를 통합된 최적화 문제로 다루지 못합니다.

2. 방법론: RDKV

본 논문은 KV 캐시 압축을 왜곡 - 비트 (Rate-Distortion) 문제로 공식화하는 RDKV(Rate-Distortion KV cache compression) 프레임워크를 제안합니다. 이러한 관점에서 추방과 양자화는 별개의 동작이 아닌 단일 비트 할당 곡선의 두 끝점이 됩니다.

2.1 이론적 공식화

RDKV 는 제거 또는 양자화가 어텐션 계산에 미치는 교란에 기반하여 각 캐시 단위 (V 캐시의 토큰, K 캐시의 채널) 에 대한 왜곡 가중치를 유도합니다.

• V 캐시 (토큰): 토큰에 대한 왜곡 가중치 $w_t$는 추방 후 어텐션 분포의 총변동 거리입니다. 이는 모든 쿼리 위치에서 해당 토큰이 받은 어텐션 점수의 합으로 계산됩니다: $w_t = \sum_{\tau} a_{\tau,t}$.
• K 캐시 (채널): 채널에 대한 왜곡 가중치 $w_c$는 해당 채널을 0 으로 설정함으로써 발생하는 로그트 편차의 스펙트럼 노름입니다. 이는 해당 쿼리 및 키 열의 노름의 곱으로 계산됩니다: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

이러한 가중치는 비트 - 왜곡 목적 함수의 비용 계수로 작용합니다:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{s.t.} \quad \sum_{u} b_u \leq B$$
여기서 $b_u$는 단위 $u$에 할당된 비트 너비이며, $B$는 총 비트 예산입니다.

2.2 비트 할당 파이프라인

할당은 이산적인 하드웨어 지원 비트 너비 $\mathcal{B} = \{0, 2, 4, 8, 16\}$에 적응된 **역수 (reverse) 수분 채우기 (water-filling, 라그랑주 완화)**를 통해 해결됩니다.

1. 1 단계 (가중치 부여): 프리필 단계 후 전체 압축되지 않은 캐시를 사용하여 토큰 가중치 ($w_t$) 와 채널 가중치 ($w_c$) 를 한 번 계산합니다.
2. 2 단계 (토큰 할당): V 캐시에 대해 다중 선택 배낭 문제 (MCKP) 를 해결합니다. 토큰은 $\{0, 2, 4, 8, 16\}$에서 비트 너비가 할당됩니다. 비트 너비 0 은 추방에 해당합니다.
3. 3 단계 (채널 할당): K 캐시에 대해 MCKP 를 해결하지만, 오직 2 단계에서 유지된 토큰에 대해서만 수행합니다. 채널은 토큰의 V 비트 너비와 독립적으로 비트 너비가 할당되어, 단일 토큰이 혼합 정밀도의 K 및 V 구성 요소를 가질 수 있습니다.

2.3 효율적인 패킹 디코드 (TriZone)

실제 메모리 및 지연 시간 절감을 실현하기 위해 RDKV 는 디양자화를 어텐션 커널에 융합하여 처리 전 FP16 으로 데이터를 언패킹할 필요가 없는 TriZone캐시 레이아웃을 도입합니다.

• 존 A: 비트 너비 $\{2, 4, 8\}$을 가진 유지된 K 행 및 V 행을 포함하는 패킹된 구캐시 스트림입니다. 이러한 행은 균일한 디양자화 경로를 허용하도록 비트 너비별로 그룹화됩니다.
• 존 B: 16 비트 정밀도가 할당된 토큰에 대한 FP16 V 행입니다.
• 존 C: 새로 생성된 토큰에 대한 FP16 K 및 V 행입니다.

어텐션 커널은 행렬 곱셈에 디양자화를 직접 융합하여 이 세 가지 존을 결합하여 로그트와 출력을 계산하며, 추가 메모리 대역폭 오버헤드를 제거합니다.

3. 주요 기여

1. 왜곡 - 비트 공식화: 본 논문은 KV 캐시 압축을 통합된 왜곡 - 비트 문제로 제시하여, 최적 해가 본질적으로 제로 - 비트 (추방) 와 유한 - 비트 (양자화) 단위를 혼합함을 증명합니다. 이는 휴리스틱 비율이나 오프라인 검색의 필요성을 제거합니다.
2. 역수 수분 채우기를 통한 이산 할당: 저자들은 할당을 라그랑주 완화를 통해 해결 가능한 이산 배낭 문제로 구체화하여, 어텐션 동역학에서 유도된 왜곡 가중치에 기반하여 $\{0, 2, 4, 8, 16\}$에서 최적 비트 너비를 할당합니다.
3. TriZone 패킹 레이아웃: 디양자화를 어텐션 커널에 융합하는 하드웨어 인식 구현으로, 이론적 비트 너비 절감을 실제 피크 메모리 및 디코드 지연 시간 감소로 전환합니다.

4. 실험 결과

실험은 LongBench, RULER, InfiniteBench를 대상으로 LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B, Qwen2.5-72B에서 수행되었습니다.

• 정확도: RDKV 는 평가된 최상의 베이스라인 (Ada-SnapKV) 보다 평균 9.1% 더 우수한 성능을 보입니다. LongBench 에서 전체 캐시 정확도의 **97.81%**를 **2.48%**의 캐시 유지율로 회복합니다.
• 긴 문맥 성능:
  • RULER: 128K 문맥에서 RDKV 는 66.95% 정확도를 달성하여 가장 강력한 베이스라인 (AdaKV) 인 58.43% 보다 8.5 포인트 높은 성능을 보입니다.
  • InfiniteBench: RDKV 는 평균 점수 39.46 으로 압축 방법 중 1 위를 차지하여 AdaKV 를 1.40 포인트 상회합니다.
  • Needle-in-a-Haystack: RDKV 는 깊이 전반에 걸쳐 균일한 검색 정확도 (평균 0.99) 를 유지하는 반면, 추방 전용 베이스라인은 중간 깊이에서 현저히 실패합니다.
• 효율성:
  • 지연 시간: 128K 문맥에서 RDKV 는 전체 캐시 FlashAttention-2 디코딩 대비 4.5 배의 디코드 속도 향상을 달성합니다.
  • 메모리: 128K 문맥에서 1.9 배의 피크 메모리 감소를 달성하며, FullKV 가 메모리 부족 (OOM) 에 빠지는 단일 A100 64GB GPU 에서 256K 문맥을 성공적으로 실행합니다.

5. 중요성 및 주장

본 논문은 단일 최적화 프레임워크 내에서 추방과 양자화를 통합적으로 다루는 것이 단계별 또는 비율 기반 접근법보다 우월하다고 주장합니다. 어텐션 계산에 유도된 왜곡에서 직접 할당 점수를 도출함으로써 RDKV 는 중간 토큰을 폐기하는 추방 전용 방법의 "이진 함정"과 미미한 토큰을 유지하는 양자화 전용 방법의 "낭비"를 피합니다.

저자들은 그들의 방법이 다음과 같음을 강조합니다:

• 모델 재학습이 필요하지 않습니다.
• 아키텍처에 구애받지 않으며, 다양한 모델 크기와 계열에서 일관된 성능을 발휘합니다.
• 긴 문맥 길이에서 두 지표 모두에서 베이스라인을 엄격히 상회하는 지연 시간과 정확도 간의 파레토 최적 트레이드오프를 제공합니다.
• 하드웨어를 위해 이산화된 왜곡 - 비트 이론의 "연속 최적해"가 현재 최첨단 하이브리드보다 우수한 혼합 비트 전략을 자연스럽게 산출함을 입증합니다.

본 논문은 이러한 통합 접근법이 효율성을 희생하지 않고 정확도를 향상시켜, 리소스 제약이 있는 하드웨어에서 긴 문맥 추론을 가능하게 한다고 결론지었습니다.