PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference

Large Language Models (LLMs) have shown remarkable comprehension abilities but face challenges in GPU memory usage during inference, hindering their scalability for real-time applications like chatbots. To accelerate inference, we store computed keys and values (KV cache) in the GPU memory. Existing methods study the KV cache compression to reduce memory by pruning the pre-computed KV cache. However, they neglect the inter-layer dependency between layers and huge memory consumption in pre-computation. To explore these deficiencies, we find that the number of crucial keys and values that influence future generations decreases layer by layer and we can extract them by the consistency in attention weights. Based on the findings, we propose PyramidInfer, a method that compresses the KV cache by layer-wise retaining crucial context. PyramidInfer saves significant memory by computing fewer keys and values without sacrificing performance. Experimental results show PyramidInfer improves 2.2x throughput compared to Accelerate with over 54% GPU memory reduction in KV cache.

한 줄 요약

PyramidInfer는 LLM 추론 시 KV 캐시를 피라미드 구조로 압축하여 54% 이상의 GPU 메모리 절감과 2.2배의 처리량 향상을 달성한 방법이다.

핵심 기여도

KV 캐시 압축을 prefill 단계부터 적용하여 기존 방법 대비 54% 이상의 GPU 메모리 절감.
Layer-wise PvC 선택을 통해 전체 캐시 계산량을 줄이며 2.2배 처리량 향상.
ICR(추론 컨텍스트 중복)과 RAC(최근 어텐션 일관성) 현상을 기반으로 PvC 추출.
OpenCompass 벤치마크에서 다양한 모델에서 성능 개선을 검증.

핵심 아이디어

기존 KV 캐시 압축 방법은 대부분 generation 단계에서 이미 계산된 캐시를 줄이는 데 집중했으나, prefill 단계에서의 메모리 소비를 고려하지 못했다. PyramidInfer는 prefill 단계부터 캐시를 압축함으로써 전체 메모리 사용량을 줄이는 새로운 접근을 제시한다. 이는 ICR 가설과 RAC 현상을 기반으로 한다. ICR은 추론 시 마지막 토큰을 제외한 나머지 토큰들이 필요 없는 예측 정보를 포함한다는 점을 지적하며, RAC는 최근 토큰들이 동일한 컨텍스트 키-값에 일관된 어텐션을 보인다는 현상을 설명한다. 이를 통해 각 레이어에서 점차 줄어드는 PvC를 선택적으로 유지함으로써 캐시 크기를 피라미드 형태로 압축한다.

기술적 접근법

**ICR 가설**: 추론 시 마지막 토큰 외의 토큰들이 불필요한 예측 정보를 포함.
**RAC 현상**: 최근 토큰들이 동일한 컨텍스트 키-값에 일관된 어텐션을 보임.
**PvC 선택**: 레이어가 깊어질수록 PvC 길이를 점진적으로 줄임.
**KV 캐시 압축**: prefill 및 generation 단계에서 PvC만 계산하여 전체 캐시 크기를 줄임.
**LLaMA 2-13B 모델**을 사용한 실험과 **OpenCompass** 벤치마크를 통해 성능 검증.

주요 결과

**Accelerate 대비 2.2배 처리량 향상**, **KV 캐시 메모리 54% 절감**.
**H2O 대비 2.4배 처리량 향상**.
**LLaMA 2-13B 모델**에서 성능 저하 없이 메모리 효율성 향상.
**OpenCompass**에서 다양한 모델과 태스크에서 일관된 개선.

의의 및 한계

PyramidInfer는 LLM 추론 시 GPU 메모리 사용량을 줄이는 데 기여하며, 특히 대규모 모델과 긴 프롬프트 처리에 유리하다. prefill 단계부터 캐시 압축을 적용함으로써 기존 방법의 한계를 극복했으며, 다양한 모델에서 벤치마크 성능을 개선했다. 그러나 PvC 선택 기준은 어텐션 가중치의 일관성에 의존하므로, 어텐션 패턴이 불규칙한 입력에서는 효과가 줄어들 수 있다. 또한, 모델 종류와 입력 길이에 따라 최적의 PvC 비율이 달라질 수 있어, 추가적인 튜닝이 필요할 수 있다.

실용적 활용

PyramidInfer는 대규모 언어 모델을 실시간 챗봇, 고객 지원 시스템, 대화형 인터페이스 등에 배포할 때 GPU 메모리 효율성을 향상시키는 데 유용하다. 특히, 서버 자원이 제한된 환경에서 여러 사용자의 동시 요청 처리를 가능하게 하며, 클라우드 기반 LLM 서비스의 확장성 향상에도 기여할 수 있다.