NVIDIA KVPress: Long-Context LLM에서의 KV 캐시 압축을 위한 툴킷

NVIDIA는 KVPress를 도입했는데, 이는 Large Language Models(LLM)의 Key-Value(KV) 캐시를 압축하도록 설계된 Python 툴킷으로, 긴 컨텍스트 윈도우의 효율적인 처리를 가능하게 합니다. 중간 주의 결과 저장과 관련된 메모리 오버헤드를 줄임으로써, KVPress는 모델이 더 큰 시퀀스—예를 들어 인-컨텍스트 검색 및 확장된 추론에 필요한 것들—을 금지적인 양의 GPU 메모리 없이 처리할 수 있게 합니다.

KV 캐시 메모리 병목

자기회귀 LLM에서는 텍스트 생성이 토큰 단위로 이루어집니다. 새로 생성되는 각 토큰에 대해 모든 이전 토큰의 표현을 다시 계산하지 않도록 하기 위해, 모델은 주의 레이어의 키(K)와 값(V)을 저장하는 KV Cache를 사용합니다. 이는 계산을 최적화하지만, 캐시 크기는 컨텍스트 윈도우에 선형적으로 증가합니다.

KV 캐시의 메모리 소비는 다음 공식에 의해 결정됩니다:

$$ ext{Size} \left( ext{KV} ight) = 2 imes ext{precision} imes n_{layers} imes n_{heads} imes d imes n_{tokens}$$

Llama 3-70B 모델이 bfloat16 정밀도로 1M 토큰 컨텍스트에서 실행되는 경우, KV 캐시만 약 327.6 GB가 필요합니다. 모델 가중치에 필요한 140 GB를 합치면 총 메모리 요구량이 약 470 GB에 달하며, 이는 KV 캐시가 총 메모리 사용량의 약 70%를 차지함을 의미합니다.

KVPress: 모듈식 압축 툴킷

KVPress는 메모리 사용량을 줄이기 위해 덜 중요한 KV 쌍을 가지치는 압축 알고리즘인 "presses" 모음을 제공합니다. 이러한 presses는 포워드 훅을 통해 모델의 주의 레이어에 통합되며, pre-filling 단계를 대상으로 하여 캐시가 가장 클 때 압축합니다.

사용 가능한 압축 기법

• KnormPress: 키-값 노름이 가장 낮은 KV 쌍을 기준으로 가지치기합니다.
• SnapKVPress: 가장 최근 쿼리에 대한 낮은 주의 가중치와 관련된 KV 쌍을 가지치기합니다.
• ExpectedAttentionPress: KVPress 저자들이 만든 새로운 미공개 기법으로, 미래 쿼리에 대한 예상 주의 가중치가 가장 낮은 KV 쌍을 가지치기합니다.

KVPress는 모듈식으로 설계되어 연구자들이 새로운 방법을 쉽게 추가하고, transformers 라이브러리에서 제공되는 KV 캐시 양자화와 같은 다른 메모리 절약 기술과 통합할 수 있도록 합니다.

성능 및 메모리 영향

KV 캐시를 압축하면 피크 메모리 사용량이 감소하고 생성 속도가 향상됩니다. bfloat16 정밀도의 128k 토큰 컨텍스트를 가진 Llama 3.1 8B 모델에 50% 압축 비율로 KVPress를 적용하면 다음과 같은 개선이 이루어집니다:

• 메모리 감소: 피크 메모리 사용량이 45GB에서 37GB로 감소합니다.
• 속도 증가: A100 GPU에서 디코딩 속도가 초당 11 토큰에서 초당 17 토큰으로 증가합니다.

벤치마킹 및 정확도 트레이드오프

KVPress는 RULER, InfiniteBench, Loogle과 같은 표준 장문 컨텍스트 데이터셋에서 압축 기법을 벤치마크할 수 있는 CLI를 포함합니다.

RULER 데이터셋(4k 토큰 컨텍스트)에서의 벤치마크 결과, 가장 효과적인 조합은 AdaKVPress와 ExpectedAttentionPress입니다. 그러나 툴킷은 중요한 트레이드오프를 강조합니다: 압축 비율이 증가할수록 모델 정확도가 일반적으로 낮아집니다. 이는 모델 성능에 미치는 영향을 최소화하는 압축 알고리즘에 대한 지속적인 연구의 필요성을 강조합니다.