[논문 리뷰] RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation

Chao Jin, Zili Zhang외 5명|arXiv (Cornell University)|2024. 04. 18.

Topic Modeling인용 수 5

한 줄 요약

RAGCache는 다단계 동적 캐싱 시스템을 도입하여 검색된 문서의 중간 키-값 텐서를 저장하고 공유함으로써 RAG의 재계산을 줄이고 지연 및 처리량을 크게 향상시킨다.

ABSTRACT

Retrieval-Augmented Generation (RAG) has shown significant improvements in various natural language processing tasks by integrating the strengths of large language models (LLMs) and external knowledge databases. However, RAG introduces long sequence generation and leads to high computation and memory costs. We propose RAGCache, a novel multilevel dynamic caching system tailored for RAG. Our analysis benchmarks current RAG systems, pinpointing the performance bottleneck (i.e., long sequence due to knowledge injection) and optimization opportunities (i.e., caching knowledge's intermediate states). Based on these insights, we design RAGCache, which organizes the intermediate states of retrieved knowledge in a knowledge tree and caches them in the GPU and host memory hierarchy. RAGCache proposes a replacement policy that is aware of LLM inference characteristics and RAG retrieval patterns. It also dynamically overlaps the retrieval and inference steps to minimize the end-to-end latency. We implement RAGCache and evaluate it on vLLM, a state-of-the-art LLM inference system and Faiss, a state-of-the-art vector database. The experimental results show that RAGCache reduces the time to first token (TTFT) by up to 4x and improves the throughput by up to 2.1x compared to vLLM integrated with Faiss.

연구 동기 및 목표

• 검색 보강 생성(RAG) 시스템의 성능 병목을 식별한다.
• 검색된 지식의 중간 상태를 캐시하는 기회를 특성화하여 계산량을 줄인다.
• cross-request 공유를 지원하는 key-value 텐서의 다단계 동적 캐싱 시스템을 설계한다.
• RAG의 접두사 민감(attention) 동작에 맞춘 교체 정책과 스케줄링 기법을 개발한다.
• 프로토타입을 시연하고 최신 기준선 대비 지연 및 처리량 향상을 정량화한다.

제안 방법

• 검색 RAG 시스템을 특성화하고 지식 주입으로 인한 장문의 시퀀스 생성이 병목인지를 확인한다.
• RAGCache를 제안하여 요청 간에 중간 key-value 텐서를 저장하고 공유하는 지식 트리 기반 캐시를 도입한다.
• GPU/호스트 메모리 계층을 관리하기 위한 prefix-aware Greedy Dual-Size Frequency(PGDSF) 교체 정책을 도입한다.
• 히트율을 높이고 쓰래싱을 줄이기 위해 캐시 인지 가능한 요청 재정렬을 구현한다.
• 검색과 추론을 겹치게 하여 엔드-투-엔드 지연을 감소시키는 동적 예측적 파이프라이닝을 개발한다.
• Faiss와 함께 vLLM에서 RAGCache를 평가하고 SGLang 및 vLLM 베이스라인과 비교한다.

실험 결과

연구 질문

• RQ1 RAG 시스템의 주요 성능 병목은 무엇이며, 검색된 지식의 중간 상태를 캐시하는 것이 이를 완화할 수 있는가?
• RQ2 요청 간 공유와 다단계 캐시가 TTFT와 처리량을 향상시킬 수 있는가?
• RQ3 RAG에서 LLM 토큰 생성의 접두사 민감한 특성을 존중하도록 교체 정책과 스케줄링을 설계하는 것이 가능한가?
• RQ4 예측적 파이프라이닝이 검색과 추론을 효과적으로 겹치게 하여 시스템 안정성을 해치지 않는가?
• RQ5 RAGCache를 기존 LLM 추론 및 벡터 데이터베이스에 통합했을 때 얻을 수 있는 실증적 이득은 무엇인가?

주요 결과

• RAG 생성 병목은 선충전(prefill) 단계가 시퀀스 길이와 문서 크기에 따라 확장되므로 이 부분이 지배적이다.
• 검색된 문서의 key-value 텐서를 캐시하면 재계산 전체보다 선충전 지연이 크게 감소한다.
• RAGCache는 TTFT를 최대 4× 감소시키고 Faiss를 사용하는 vLLM에 비해 처리량을 최대 2.1× 향상시켰다.
• RAGCache는 SGLang 대비 TTFT를 최대 3.5× 감소시키고 처리량도 최대 1.8× 증가시킨다.
• 지식 트리는 문서 순서 및 LLM 상태 의존성에 맞춘 접두사 민감 캐싱을 가능하게 한다.
• 동적 예측적 파이프라이닝은 검색과 추론을 겹치게 하여 엔드-투-엔드 지연을 추가로 감소시킨다.