FastKV: 토큰 선택적 전파를 통한 빠른 장문맥 처리를 위한 KV 캐시 압축

FastKV: KV Cache Compression for Fast Long-Context Processing with Token-Selective Propagation

February 3, 2025

저자: Dongwon Jo, Jiwon Song, Yulhwa Kim, Jae-Joon Kim

cs.AI

초록

대규모 언어 모델(LLM)은 장기 문맥 시퀀스를 처리하는 데 뛰어나지만, 문맥 정보를 저장하기 위해 상당한 키-값(KV) 캐시가 필요하여 계산 효율성과 메모리 사용량에 부담을 줄 수 있습니다. 이전 KV 캐시를 압축하는 노력은 주로 메모리 요구량을 줄이는 데 초점을 맞추었지만, 대기 시간을 개선하는 데 제한이 있었습니다. 이 문제를 해결하기 위해 우리는 FastKV를 소개합니다. FastKV는 장기 문맥 시퀀스의 대기 시간을 향상시키기 위해 설계된 KV 캐시 압축 방법입니다. 정확도를 유지하면서 처리 속도를 향상시키기 위해 FastKV는 LLM의 초기 레이어에서 전체 문맥 정보를 유지하고 깊은 레이어에서는 사전 채우기 단계에서도 이 정보의 일부만 선택적으로 전파하는 혁신적인 Token-Selective Propagation (TSP) 접근 방식을 채택합니다. 또한 FastKV는 그룹화된 쿼리 어텐션(GQA)을 고려한 KV 캐시 압축을 통해 GQA의 메모리 및 계산 효율성 장점을 활용합니다. 실험 결과는 FastKV가 HeadKV, 최첨단 KV 캐시 압축 방법과 비교하여 첫 번째 토큰까지의 시간(TTFT) 및 처리량에서 각각 2.00배 및 1.40배의 개선을 달성한다는 것을 보여줍니다. 또한 FastKV는 장기 문맥 벤치마크에서 기준선과 비교 가능한 수준의 정확도를 성공적으로 유지합니다. 코드는 https://github.com/dongwonjo/FastKV에서 확인할 수 있습니다.

English

While large language models (LLMs) excel at handling long-context sequences, they require substantial key-value (KV) caches to store contextual information, which can heavily burden computational efficiency and memory usage. Previous efforts to compress these KV caches primarily focused on reducing memory demands but were limited in enhancing latency. To address this issue, we introduce FastKV, a KV cache compression method designed to enhance latency for long-context sequences. To enhance processing speeds while maintaining accuracy, FastKV adopts a novel Token-Selective Propagation (TSP) approach that retains the full context information in the initial layers of LLMs and selectively propagates only a portion of this information in deeper layers even in the prefill stage. Additionally, FastKV incorporates grouped-query attention (GQA)-aware KV cache compression to exploit the advantages of GQA in both memory and computational efficiency. Our experimental results show that FastKV achieves 2.00times and 1.40times improvements in time-to-first-token (TTFT) and throughput, respectively, compared to HeadKV, the state-of-the-art KV cache compression method. Moreover, FastKV successfully maintains accuracy on long-context benchmarks at levels comparable to the baselines. Our code is available at https://github.com/dongwonjo/FastKV.