- Title: PackKV Reducing KV Cache Memory Footprint through LLM-Aware Lossy Compression
- ArXiv ID: 2512.24449
- 발행일: 2025-12-30
- 저자: Bo Jiang, Taolue Yang, Youyuan Liu, Xubin He, Sheng Di, Sian Jin
📝 초록
트랜스포머 기반 대형 언어 모델(LLMs)은 광범위한 실제 응용 분야에서 놀라운 잠재력을 보여주고 있습니다. 그러나 키-밸류(KV) 캐시의 큰 메모리 요구 사항으로 인해 긴 문맥 추론이 여전히 중요한 도전 과제입니다. KV 캐시는 시퀀스 길이와 배치 크기가 증가함에 따라 몇 GB까지 확장될 수 있습니다. 본 논문에서는 긴 문맥 생성을 위한 최적화된 일반적이고 효율적인 KV 캐시 관리 프레임워크인 \textbf{PackKV}를 제시합니다. PackKV는 KV 캐시 데이터의 특성에 맞게 조정된 손실 압축 기법을 도입하며, 압축 알고리즘과 시스템 아키텍처 간의 세밀한 공학 설계를 특징으로 합니다. 우리의 접근 방법은 KV 캐시가 동적으로 성장하는 성격을 유지하면서 높은 연산 효율성을 보존합니다. 실험 결과는, 현존하는 양자화 방법과 같은 정확도 하락을 감안할 때 PackKV가 평균적으로 K 캐시에 대해 \textbf{153.2}\%의 메모리 감소율을, V 캐시에 대해서는 \textbf{179.6}\%를 달성함을 보여줍니다. 또한 PackKV는 분해압축 오버헤드를 효과적으로 제거하고 행렬-벡터 곱 연산을 가속화하여 매우 높은 실행 통과율을 제공합니다. 구체적으로, A100 및 RTX Pro 6000 GPU에서 cuBLAS 행렬-벡터 곱 커널에 비해 K는 평균 \textbf{75.7}\%, V는 \textbf{171.7}\%의 통과율 향상을 이루며, 이는 더 적은 GPU 메모리 대역폭을 필요로 합니다. 코드는 https://github.com/BoJiang03/PackKV에서 이용 가능합니다.
💡 논문 해설
1. **새로운 압축 파이프라인 설계:** PackKV는 양자화, 인코딩에 대한 재포장, 비트 패킹 인코딩을 통합한 첫 번째 대형 언어 모델(Large Language Model)에 특화된 손실 압축 프레임워크입니다. 이 방법은 기존의 양자화 방법과 비교해 K 캐시와 V 캐시에서 각각 153.2%와 179.6% 더 높은 압축 비율을 달성합니다.
메타포: PackKV는 큰 책을 작은 가방에 넣듯이, 모델의 메모리 사용량을 줄이는 기술입니다.
계산 인식 디컴프레션 통합: 이 방법은 매트릭스-벡터 곱셈 커널 내부에서 직접 디컴프레션을 통합하여 압축된 데이터를 GPU 공유 메모리와 레지스터에 바로 불러올 수 있게 합니다. 이를 통해 전역 메모리 쓰기백과 중복 읽기를 줄여, 총 처리량을 향상시킵니다.
메타포: PackKV는 책장을 여는 대신 필요한 정보만 빠르게 찾아내어 더 효율적으로 데이터를 사용하는 것입니다.
이론적 기초와 실용적인 알고리즘: KV 캐시 재정렬 문제를 집합 분할 문제의 변형으로 정식화하고, 시간 복잡도가 낮은 그리디 솔루션과 중위값 기반 효율적인 알고리즘을 제공합니다.
메타포: PackKV는 책장을 효과적으로 재정렬하여 필요한 정보를 빠르게 찾을 수 있도록 도와줍니다.
📄 논문 발췌 (ArXiv Source)
손실 압축, KV 캐시, 대형 언어 모델, GPU
서론
Transformer 기반의 LLMs는 자연어 처리를 혁명화하여 다양한 작업에서 혁신을 가능하게 했습니다. 자기 주의 메커니즘은 모델이 장거리 종속성과 문맥 정보를 포착할 수 있게 합니다. 그러나 이러한 능력은 긴 입력 컨텍스트에서 추론 중에 상당한 계산 및 메모리 비용을 초래합니다.
이러한 발자국이 증가함에 따라 추론 성능이 크게 제약받고, 가능한 컨텍스트 길이를 제한하거나 배치 크기를 줄이는 등의 문제가 발생하며, 메모리 제약을 받는 하드웨어에서 LLMs의 배포를 방해합니다. KV 캐시는 모델이 처리하는 각 토큰에 대한 중간 key와 value 텐서를 저장하고 있으며 이후 디코딩 단계에서 재사용되어 불필요한 계산을 피하기 위해 사용됩니다. 시퀀스 길이와 배치 크기가 증가함에 따라 캐시 크기는 선형적으로 성장하여 GPU 메모리의 상당 부분을 차지할 수 있습니다.
이러한 문제를 해결하기 위해 최근 연구는 세 가지 접근 방식을 주로 활용합니다: 양자화, 필터링, 및 GPU-CPU 이동. KIVI와 같은 양자화 기반 방법은 정확도에 미치는 영향을 최소화하는 양자화 전략을 신중하게 설계하려고 합니다. 그러나 이러한 기술들은 추가 인코딩 없이는 제한적인 압축 비율만을 달성하며, 이는 오버헤드를 초래하고 지연에 민감한 LLM 추론에서 적용 범위를 제한합니다.
Q-Hitter와 같은 필터링 기반 방법은 미래의 디코딩에 중요하지 않다고 예측되는 KV 쌍을 선택적으로 삭제합니다. 일부 경우에는 효과적이지만, 이러한 방법들은 불확실한 주의 콜백으로 인해 비용이 많이 드는 KV 재계산이나 상당한 정확도 저하를 초래할 수 있습니다.
GPU-CPU 이동은 메모리 오버플로우 처리를 위한 전통적인 접근 방식으로, KV 캐시 데이터를 CPU 메모리에 할당합니다. 이를 통해 GPU 메모리 압박이 완화되지만, 데이터 전송 지연과 복잡한 스케줄링 오버헤드로 인해 추론 성능이 크게 저하됩니다.
본 논문에서는 추론 중 KV 캐시 최적화를 위한 고성능 LLM에 특화된 손실 압축 프레임워크인 PackKV를 제안합니다. PackKV는 오류 제어 양자화, KV 캐시 특수 압축 방식 및 캐시 상주 디컴프레션을 통합하여 상당한 메모리 절감과 계산 효율성을 향상시키도록 설계되었습니다. 패킹 앤 디코딩 파이프라인의 시스템 수준 실행과의 공동 설계를 통해, PackKV는 압축해제된 데이터가 GPU 공유 메모리 및 레지스터 내에서 바로 사용될 수 있도록 하여 전역 메모리 쓰기백을 제거하고 처리량을 최대화합니다. 이러한 고성능 디자인으로, PackKV는 LLM 추론에 거의 오버헤드 없이 배포 가능하며 대부분의 경우 추론 계산을 가속시킵니다.
주요 기여:
신규 압축 파이프라인 설계: PackKV를 소개합니다. 이는 양자화, 인코딩 인식 재포장, 비트 패킹 인코딩을 통합한 첫 번째 LLM에 특화된 손실 압축 프레임워크입니다.
계산 인식 디컴프레션 통합: 매트릭스-벡터 곱셈 커널 내부에서 직접 디컴프레션을 통합하여 전역 메모리 쓰기백과 중복 읽기를 제거합니다.
이론적 기초와 실용적인 알고리즘: KV 캐시 재정렬 문제를 집합 분할 문제의 변형으로 정식화하고, 시간 복잡도가 낮은 그리디 솔루션과 중위값 기반 효율적인 알고리즘을 제공합니다.
포괄적 시스템 구현: 완전한 GPU 기반 프레임워크를 개발하여 동적 KV 캐시 성장에 대한 원활한 추가, 여러 블록 간 단일 커널 디컴프레션 및 영 제로 간섭을 갖춘 완벽한 멀티-GPU 확장성을 제공합니다.
확실한 실험 검증: 여섯 개의 LLM 모델과 여섯 개의 벤치마크 데이터셋에 대한 평가를 통해 PackKV는 압축 비율에서 기존 양자화 방법을 능가하고, cuBLAS 커널보다 처리량이 향상된다는 것을 보여줍니다.
본 논문은 다음과 같이 구성되어 있습니다: 섹션 2에서는 LLM 추론, 손실 압축 및 GPU 아키텍처에 대한 배경을 제공합니다. 섹션 3에서는 PackKV 프레임워크의 설계를 제시합니다. 섹션 4에서는 실험적 평가 및 분석을 상세히 설명합니다. 마지막으로, 섹션 5에서는 논문의 결론과 향후 연구 방향을 제시합니다.
배경 및 동기
대형 언어 모델 추론
LLMs는 자연어 처리에서 핵심적인 역할을 하며, 텍스트 생성에 뛰어난 능력을 보여줍니다. 추론 과정은 자동 회귀 디코딩을 포함하며 여기서 모델은 각 토큰을 순차적으로 생성합니다.
이러한 과정을 효율적으로 수행하기 위해 Transformer 모델은 Key-Value(KV) 캐시를 유지합니다. KV 캐시는 주의 메커니즘에 의해 구축되며, 추론 중 각 토큰에 대해 Key(K)와 Value(V) 텐서가 계산되고 이를 통해 미래의 토큰을 처리하는 데 재사용됩니다.
KV 캐시는 [context_len, head_num, head_dim] 형식으로 구성되며, 각각 시퀀스 길이, 주의 헤드 수, 헤드 차원을 나타냅니다.
/>
KV 캐시의 prefill 단계(3 토큰 입력 가정). />
디코딩 단계에서 KV 캐시(1 토큰 생성).LLM 추론 중 KV Cache 동작, $`M_{Q,K,V}`$는 $`K,Q,V`$ 벡터로 매핑하는 행렬입니다. 각 $`Q`$ 벡터는 모든 $`K`$ 벡터와 내적을 수행하여 소프트맥스 연산 후 각 $`V`$ 벡터에 대한 가중치를 생성합니다. 모든 $`V`$ 벡터는 그들의 가중치와 곱해져 Attn Output으로 집계됩니다.
추론 중 KV 캐시는 Prefill Stage에서 초기 프롬프트 처리를 통해 초기 KV 캐시를 생성하고, Decode Stage에서는 새로운 KV 벡터가 캐시에 추가되어 전체 캐시를 주의 연산을 위한 매트릭스-벡터 곱셈으로 사용합니다.
KV 캐시는 GPU 메모리 사용량을 급격히 증가시키며, 시퀀스 길이와 배치 크기에 따라 선형적으로 성장하므로 GPU 메모리를 모델 가중치보다 더 많이 소비할 수 있습니다. 이러한 문제를 해결하기 위해 KV 캐시 양자화 및 GPU-CPU 이동 등의 방법이 사용되지만, 이들은 제한적인 효과만을 보이고 성능 오버헤드가 발생합니다.
KV 캐시는 추론 중 메모리 발자국과 트래픽의 주요 부분을 차지하며, 이를 관리하는 데 필요한 메모리 대역폭이 계산보다 주된 병목 지점입니다. 디코딩 연산은 매트릭스-벡터 곱셈에 의해 주도되며, 이는 낮은 산술 강도와 데이터 재사용 부족으로 인해 메모리 제약을 가집니다.
LLMs에서의 데이터 축소
/>
PackKV 개요. vanilla 및 이전 양자화 솔루션과 비교. (1) 인코딩 후 비트 패킹을 위해 양자화된 kv 캐시 벡터의 순서를 재정렬합니다. (2) 채널 방향으로 양자화된 kv 캐시 정수에 대한 비트 패킹을 적용하여 추가 메모리 절감을 달성합니다. 그리고 디컴프레션 오버헤드를 제거하기 위해 계산 인식 디컴프레션 + 매트릭스-벡터 곱셈 방법을 개발했습니다.
대형 신경망, 특히 LLMs에서는 손실 압축이 메모리 병목 현상을 완화하는 데 사용될 수 있습니다. 양자화는 연속 또는 고정밀도 값(예: float32)을 더 작은 이산값 집합(예: 정수)으로 매핑하여 데이터의 엔트로피를 줄이는 형태의 손실 압축입니다.
이러한 후 양자화 인코더의 낮은 처리량은 실제 LLM 추론에서 지연에 민감하기 때문에 실용적이지 않습니다. CacheGen과 같은 방법은 이러한 인코더를 KV 캐시 양자화 파이프라인에 통합하여 1GB/s 미만의 처리량을 제공하지만, GPU-CPU 전송 및 GPU 메모리 대역폭에는 부족합니다.
손실 압축을 적용할 때는 압축 비율과 모델 정확도 사이의 균형 관리를 해야 합니다. 양자화를 적용하는 데는 데이터 특성과 특정 하류 작업에 대한 신중한 고려가 필요하며, 다양한 그레인ularity(예: 토큰/채널/블록별) 또는 적응적 오차 경계 사용 등의 다른 양자화 전략을 사용할 수 있습니다.
GPU 아키텍처
GPU 메모리 계층은 성능 최적화에 중요한 역할을 합니다. 전역 메모리는 높은 용량과 고 지연으로 제공되므로 효율적인 접근 패턴이 필수적입니다. 각 스트림 멀티프로세서(SM)의 공유 메모리는 블록 내의 쓰레드들 사이에서 공유되는 프로그래머 관리형 캐시입니다.
GPU 실행은 단일 명령 다중 스레드(SIMT) 모델을 따르며, 워프(보통 32개 스레드) 내부의 스레드는 동일한 명령을 잠금 상태에서 동시에 수행합니다. 스레드가 서로 다른 실행 경로를 따라야 하는 경우 발생하는 분기 발산은 직렬화를 강제하고 성능을 저하시킵니다. 쓰레드 간 조정을 위해 GPU는 간섭 없이 읽기-수정-쓰기를 완료할 수 있는 원자 연산을 제공합니다.
설계 방법론
본 섹션에서는 LLM 추론을 위한 KV 캐시 압축 프레임워크인 PackKV를 제시합니다. Figure 4은 우리의 다섯 단계 파이프라인을 보여줍니다: 양자화, 인코딩 인식 재포장, 비트 패킹, 원활한 추가, 그리고 계산 인식 디컴프레션.
우리의 프레임워크는 고정 크기 버퍼에 저장된 블록 단위로 KV 캐시를 처리합니다. 첫 번째 단계에서는 각 토큰의 KV 캐시가 공유 양자화 스케일을 사용하여 양자화됩니다. 두 번째 단계에서는 압축 효율성을 최대화하면서 K-V 쌍을 함께 유지하는 방식으로 양자화된 벡터를 재정렬합니다. 세 번째 단계에서는 신중하게 관리되는 메타데이터와 함께 가벼운 비트 패킹 인코딩이 적용됩니다. 네 번째 단계에서는 압축 블록을 추가하여 단일 커널 디컴프레션을 가능하도록 합니다. 마지막으로, 다섯 번째 단계에서는 매트릭스-벡터 곱셈과 결합된 디컴프레션을 통합하여 압축해제된 데이터가 GPU 레지스터에 바로 불러와져 메모리 접근을 줄이고 LLM 추론의 주요 병목 현상을 가속화합니다.
알고리즘 설계
우리의 압축 파이프라인은 두 가지 주요 구성 요소로 구성됩니다: 양자화를 통한 손실 압축 및 인코딩 인식 재포장과 비트 패킹을 통한 무손실 압축. 양자화 단계에서는 고정밀도 값이 이산 정수로 매핑되어 데이터 엔트로피가 줄어듭니다. 무손실 압축을 위해 먼저 인코딩 인식 재포장을 적용합니다: 양자화된 KV 캐시 벡터는 그룹으로 재배치되며, 이는 후속 인코더가 더 짧은 per-pack 인코딩 길이를 사용하여 높은 압축 비율을 달성하도록 합니다. 그런 다음 이러한 팩들을 가벼운 비트 패킹 스키ema로 인코딩합니다.
