• RAGCache
  • RAGCache의 작동과정
  • Cache retriever
    • Knowledge Tree Structure
  • RAG controller
    • Prefix-aware Greedy Double Size Frequency (PGDSF) replacement policy
    • Multi-level caching (tree 내 node 배치 전략)
    • Dynamic speculation pipeline
• Evaluation
• Reference

본 논문은 베이징 대학교와 ByteDance가 효율적인 RAG 수행을 위해 연구한 내용을 담고 있다. 최근 LLM 추론 시간 자체에 대한 가속화 연구가 많이 진행되고 있다(vLLM, SGLang …). RAG 작동 효율 향상을 위해 LLM 추론 자체가 빠른 것도 도움되지만 RAG 시스템 특성을 적절히 고려한 최적화도 필요하다.

RAG 시스템에서 보이는 주요한 문제는 아래와 같다:

1. Performance Bottleneck
RAG의 retirieval 단계는 일반적으로 millisecond 단위로 매우 빠르게 처리된다.

문제는 generation 단계이다. 검색된 문서가 input에 추가되면서 추론 시 계산 비용과 메모리 사용량이 증가한다.

• LLM 추론은 크게 두 단계로 진행된다: 1) prefill 단계: 입력 토큰의 key-value 텐서 계산 2) decoding 단계: 이전에 생성된 key-value 텐서 기반으로 출력 토큰 생성

RAG 시스템에서는 retrieval된 문서가 LLM input에 추가된다. 이때 LLM이 처리할 sequence가 길이가 매우 길어지게 되는데 이 부분이 RAG 시스템에서 주요한 병목 지점이다. 즉, prefill 단계에서 매우 긴 sequence에 대한 계산으로 많은 시간이 소요된다. 시퀀스 길이가 4000토큰 이상이면 처리 시간이 1초 이상 소요되는 것으로 알려져 있다.

2. Access Pattern
Retrieval request는 반복적인 패턴을 보인다. 즉, 소수의 문서가 자주 검색되는 특성이 있다. 예를 들어 전체 검색 요청의 60%이상이 전체 문서의 3%에 대한 질문이다. 따라서 이와 같은 패턴에 대해 따로 caching해 놓으면 속도가 빨라질 수 있다.

3. Optimized Space
자주 사용되는 문서들에 대해 중간 계산 단계(key-value tensor)를 caching함으로써 computational overhead를 크게 줄일 수 있다. 즉, LLM 이 처리해야 할 key-value tensor들을 미리 계산해서 cache에 저장해 둠으로써 시간을 크게 줄일 수 있다는 것이다. 이렇게 되면 자주 사용되는 문서를 처리할 때 저장된 tensor를 재사용할 수 있어 결과적으로 처리 시간이 단축되게 된다. 실험 결과에 따르면 이 방법을 도입한 후 pre-population 시간이 11.5배 감소했다고 한다.

RAGCache

RAGCache의 핵심 아이디어는 multi-layer dynamic caching system을 통해 검색된 문서들에 대한 중간 계산 상태를 효율적으로 caching하고 재사용함으로써 RAG을 효율적으로 수행하는 데에 있다.

RAGCache의 작동과정

1. Request(질문 요청) ->
2. Retrieval(문서 검색) ->
3. Cache 확인 (Cache retriever가 검색된 문서의 key-value tensor가 저장된 게 있는지 확인)
   • cache miss : LLM이 일반적인 generation 방식으로 추론 시작
   • cache hit : 저장된 tensor를 LLM engine에 전달
4. token generation 4.1) 첫 token 생성 이후 key-value tensor를 controller에 전달 4.2) controller가 tensor를 cache하고 state 업데이트 4.3) 답변 생성
5. 최종 반환

Cache retriever

저장된 Cache를 검색하는 Cache retriever의 주요 component는 Knowledge Tree Structure이다.

Knowledge Tree Structure

RAG은 문서 순서에 민감하다. 따라서 RAGCache는 knowledge tree 구조를 설계하여 caching된 문서들의 state를 구조화했다. 이와 같은 tree 형태의 구조는 RAG 시스템에서 문서 검색의 순서 민감성 문제를 해결한다.

Knowledge Tree는 document ID 기반의 prefix tree 형태를 가진다.

예를 들어 아래 그림에서 [D1,D3]와 [D2,D3] 두 개의 문서 sequence가 있을 때, D3가 두 sequence 모두에 나타나지만 앞에 나온 sequence가 다르기 때문에 각 sequence에서의 key-value tensor 값이 서로 다르다. knowledge tree 구조는 이와 같은 변화들을 효율적으로 관리하기 위해 prefix tree 형태를 가지며, 문서들의 순서를 유지함과 동시에 빠른 검색을 가능하게 한다.

(이 트리에서 각 경로는 특정 문서의 sequence를 표현하고, 각 node는 문서의 key-value tensor를 저장한다. 서로 다른 경로들이 node를 공유하는 것이 가능하다. 그리고 root node는 system prompt이다.)

RAG controller

Prefix-aware Greedy Double Size Frequency (PGDSF) replacement policy

RAG controller에서는 여러 요소들을 고려하는 복잡한 caching replacement 전략을 사용한다. 이 전략은 가장 가치있는 document state들이 cache에 유지되도록하며, cache hit rate를 최대화하고 중복 계산을 최소화 한다.

cache hit rate는 시스템이 필요한 데이터를 cache에서 성공적으로 찾는 비율이다.
높은 cache hit rate = 많은 요청을 캐시에서 바로 처리 가능 -> 시스템 전체의 성능이 향상

Cache Hit: 필요한 데이터가 캐시에 있어서 바로 가져다 쓸 수 있는 경우 Cache Miss: 필요한 데이터가 캐시에 없어서 처음부터 다시 계산해야 하는 경우

PGDSF 전략 목적은 하나다. cache hit 최대화, cache miss 최소화 이를 위해 아래와 같은 기준으로 계산의 우선순위를 둔다.

1. Frequency(접근 빈도. 시스템 재시작 혹은 cache 정리 시 초기화)
2. Size(tensor 크기. size 계산 시, 문서의 token 수와 메모리 사용량이 반양된다.)
3. Clock(마지막 접근 시간. 문서 검색될 때마다 업데이트됨)
4. Cost(prefix-aware를 다시 계산하데 들어가는 계산 비용. 즉, key-value tensor 계산 시간.)

Multi-level caching (tree 내 node 배치 전략)

RAG controller에서는 GPU와 host memory를 효율적으로 활용하기 위 접근 빈도가 높은 document state들은 빠른 접근을 위해 GPU memory에 저장하고, 접근 빈도가 낮은 document state들은 host memory에 저장한다. 이와 같이 계층화된 접근 방식은 속도와 용량 둘 다 최적화하는 데 도움이된다.

Host Memory는 컴퓨터의 주 메모리인 RAM으로, GPU memory에 비해 속도는 느리지만 용량이 크다.
GPU memory는 더 빠르지만 용량이 작다.

Dynamic speculation pipeline

RAG controller에서는 dynamic speculation pipeline을 통해 vector retrieval과 LLM 추론을 overlap시켜 실행한다. 즉, 검색결과가 나오는대로 LLM 추론을 시작한다. 이렇게 진행하면 지연시간을 최소화할 수 있다.

1. vector retrieval이 진행될 때 RAGCache는 중간 결과들을 LLM에 보내 speculative generation을 수행한다.
2. 검색된 문서가 변경되면 새로운 speculative generation을 시작한다.
3. 이 방법은 유휴시간을 최소화하여 전체 지연 시간을 크게 줄일 수 있다.

Evaluation

1. first token generation time (TTFT)이 vLLM+Faiss와 비교하여 최대 4배 빠르다
2. vLLM+Faiss와 비교하여 처리량이 최대 2.1배까지 증가했다.
3. SGLang과 비교하여 TTFT가 최대 3.5배까지 감소했다.
4. SGLang과 비교하여 처리량이 최대 1.8배까지 증가했다.

RAGCache 도입 이후 위와 같은 큰 성능 향상이 관찰되었다고 한다. 이와 같은 성능 향상은 다양한 모델, 다양한 retrieval setup에서도 일관되게 나타나 RAGCache 자체가 다양한 환경에서 견고히 사용될 수 있는 방법론임이 확인되었다고 한다.

Reference

RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation