KV 캐시 메모리 절감 혁신 PackKV

초록

PackKV는 대형 언어 모델의 KV 캐시를 대상으로 LLM 특성을 고려한 손실 압축 파이프라인을 제안한다. 양자화, 인코드‑인식 재배열, 비트‑패킹을 결합해 K와 V 캐시를 각각 평균 153 %·180 % 이상의 압축률을 달성하면서 정확도 손실을 최소화한다. 또한 압축·복원 과정이 연산 파이프라인에 거의 영향을 주지 않아 GPU 매트릭스‑벡터 곱 연산을 75 %~172 % 가속한다.

상세 분석

PackKV 논문은 LLM 추론 시 가장 큰 메모리 병목인 KV 캐시 관리 문제를 근본적으로 재설계한다는 점에서 의미가 크다. 기존 연구들은 주로 정밀도 손실을 최소화하기 위해 8‑bit 혹은 4‑bit 양자화에 의존했지만, KV 캐시 데이터는 토큰 간 상관관계와 레이어별 구조적 특성을 가지고 있어 단순 양자화만으로는 압축 효율에 한계가 있었다. PackKV는 이러한 특성을 “LLM‑aware”하게 분석하고, 세 단계 압축 파이프라인을 설계한다. 첫 번째 단계인 양자화는 기존 방법과 동일하게 부동소수점 텐서를 저비트 정수로 변환하지만, 레이어별 동적 범위와 토큰 위치 정보를 활용해 스케일링 팩터를 맞춤형으로 설정한다. 두 번째 단계인 인코드‑인식 재배열(encode‑aware repacking) 은 KV 쌍을 토큰 순서와 헤드 구조에 따라 그룹화하고, 동일한 통계적 특성을 공유하는 블록을 하나의 압축 단위로 묶는다. 이 과정에서 메모리 접근 패턴을 최적화해 연속적인 메모리 읽기를 보장한다. 마지막 단계인 비트‑패킹 인코딩(bit‑packing encoding) 은 재배열된 블록을 가변 길이 비트 스트림으로 압축한다. 여기서는 Huffman‑like 코드북을 사전 학습하고, GPU의 SIMD 연산에 최적화된 워프‑단위 비트 연산을 적용해 압축·복원 비용을 최소화한다.

핵심적인 설계 선택은 압축 후에도 KV 캐시를 동적으로 확장할 수 있게 만든 점이다. 기존 양자화 방식은 새로운 토큰이 추가될 때마다 전체 캐시를 재배열하거나 재양자화해야 하는 반면, PackKV는 블록 단위로 압축을 수행하므로 새로운 토큰이 들어오면 해당 블록만 추가 압축하면 된다. 이로써 실시간 생성 시 메모리 재할당 오버헤드가 크게 감소한다.

성능 평가에서는 LLaMA‑2‑7B, LLaMA‑2‑13B, Falcon‑40B 등 다양한 모델과 A100, RTX Pro 6000 GPU를 대상으로 실험했다. 동일한 정확도 손실(≤0.2 % BLEU 감소) 조건에서 K 캐시 압축률은 평균 2.53×, V 캐시 압축률은 2.79×에 달했으며, 이는 기존 최첨단 8‑bit 양자화 대비 각각 153 %·180 % 더 높은 수치이다. 또한 압축된 KV를 사용한 매트릭스‑벡터 곱 연산은 cuBLAS 대비 K는 75.7 %, V는 171.7 %의 처리량 향상을 보였고, 메모리 대역폭 사용량도 현저히 낮았다. 이는 비트‑패킹 단계가 GPU 레지스터와 공유 메모리에서 직접 연산을 수행하도록 설계돼, 메모리 복사 비용을 거의 없앴기 때문이다.

결과적으로 PackKV는 메모리 효율성과 연산 효율성을 동시에 달성한 최초의 LLM‑aware 손실 압축 프레임워크로 평가된다. 향후 초대형 모델(수백억 파라미터)이나 멀티‑GPU 분산 환경에서도 KV 캐시 압축을 적용하면, 메모리 제한으로 인한 컨텍스트 길이 제약을 크게 완화할 수 있을 것으로 기대된다.