CacheBlend Fast Large Language Model Serving for RAG with Cached Knowledge Fusion

논문 : https://arxiv.org/abs/2405.16444

개인의견 : rag가 앞으로 정말 많이 사용될거 같은데 어떻게보면 이러한 방법론을 사용하여 kv cache를 partial하게 사용하고 update를 하는 방법에 대한 논문이군요

아래글은 Scholar GPT와 대화를 나눈 내용입니다.

Q : 논문의 핵심 내용과 강점, 알고리즘 설명, 그리고 한계점

Answer

논문의 강점과 독창성

논문 “CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion”는 대형 언어 모델(LLM)의 추론 속도를 획기적으로 개선하기 위해 CacheBlend라는 새로운 KV Cache 재활용 방법을 제안합니다. 특히, 이 방식은 다중 텍스트 청크가 포함된 입력에 대해 효율적인 KV Cache 처리를 통해 생성 품질을 유지하면서도 최초 토큰 생성 지연(TTFT)을 줄이는 것이 목표입니다.

주요 강점과 독창성

1. 고속 KV 캐시 융합: CacheBlend는 KV Cache를 특정 토큰의 KV만 선택적으로 재계산하면서 융합하여, 기존 방법보다 훨씬 빠르게 작동합니다. 이를 통해 LLM의 입력 프리필(prefill) 속도가 향상됩니다.
2. 유연한 KV 캐시 활용: CacheBlend는 재사용되는 텍스트 청크가 입력의 프리픽스(접두사)인지와 관계없이 기존의 모든 KV 캐시를 재활용할 수 있는 유연성을 제공합니다.
3. 파이프라이닝을 통한 처리 지연 감소: 선택적인 KV 재계산을 KV 캐시 로드와 병렬로 수행하여 추론 지연을 최소화하고, NVMe SSD와 같은 느린 저장 장치에서도 효율적으로 작동할 수 있습니다.
4. 생성 품질 유지: 기존 방법에서 간과된 교차 주의(cross-attention) 정보의 손실을 최소화하여, 생성 품질을 높게 유지하면서도 캐시 융합을 수행합니다.

핵심 알고리즘: CacheBlend의 선택적 KV 재계산

CacheBlend의 알고리즘은 여러 텍스트 청크의 프리컴퓨팅된 KV 캐시를 재활용하면서도 일부 토큰에 대한 KV 값을 선택적으로 재계산하여 전체 재계산 없이도 높은 품질의 출력을 보장하는 데 중점을 둡니다. 이 과정을 예시와 함께 설명하면 다음과 같습니다.

예시: CacheBlend를 통한 선택적 KV 재계산 과정

1. KV 캐시 로드 및 초기화:
   • 예를 들어, “Chunk A”, “Chunk B”, “Chunk C”라는 3개의 청크가 재사용된다고 가정합니다. CacheBlend는 각 청크의 KV 캐시를 저장 장치에서 로드합니다.
2. 선택적 KV 재계산:
   • CacheBlend는 이전 층에서 KV 편차가 큰 토큰을 기준으로, 각 층에서 주의가 필요한 토큰을 선택하고 이 토큰에 대한 KV를 재계산합니다. 예를 들어, 3개의 청크 중에서 Chunk A에서 재계산이 필요한 토큰들을 식별하여 해당 토큰만 선택적으로 재계산합니다.
   • 이 과정에서 전체 토큰 중 일부 토큰만을 재계산하여 연산 비용을 크게 절감합니다.
3. KV 병합과 결과 출력:
   • CacheBlend는 선택적으로 재계산된 KV와 기존 KV 캐시를 융합하여 최종적인 KV 캐시를 생성합니다. 이를 기반으로 각 청크 간 교차 주의 정보가 유지된 채로 고품질 응답을 생성할 수 있습니다.

논문의 한계점

1. 일부 모델 구조에 대한 제한: CacheBlend는 Transformer 기반의 언어 모델에 최적화되어 있어, 다른 아키텍처(예: Mamba 또는 Griffin)에서는 효과가 제한적일 수 있습니다.
2. 다양한 LLM 서빙 엔진과의 통합 제한: 최신 LLM 서빙 엔진(e.g., Distserve, StableGen)과의 호환성을 테스트하지 않아, 향후 확장성이 일부 제한될 가능성이 있습니다.
3. 선택적 KV 재계산 설정 복잡성: 특정 토큰의 선택적 재계산에 따른 최적화 과정이 복잡하여, 실제 구현 및 적용에 있어 성능 최적화를 위한 추가 설정이 필요할 수 있습니다.

이 논문은 CacheBlend의 선택적 KV 캐시 재계산 방식을 통해 LLM의 처리 속도와 효율성을 높이면서도 교차 주의가 유지된 고품질 응답을 제공하는 데 강점을 두고 있습니다

Q: 이 논문이 promtCache와 다른점이 어떤점일까?

CacheBlend와 PromptCache는 KV 캐시를 재활용하여 LLM 추론 속도를 높이고 효율을 극대화하려는 점에서 유사하지만, 캐시를 재활용하는 방식과 캐시의 구조에 있어 몇 가지 중요한 차이점이 있습니다.

CacheBlend와 PromptCache의 주요 차이점

1. KV 캐시 활용 방식
   • CacheBlend: CacheBlend는 여러 텍스트 청크가 포함된 입력에 대해 일부 중요한 토큰의 KV만 선택적으로 재계산하여 KV 캐시를 융합합니다. 이 방식을 통해 전체 KV 캐시를 반복적으로 로드하거나 계산하지 않고도 고품질의 응답을 빠르게 생성할 수 있습니다.
   • PromptCache: PromptCache는 특정 프롬프트 모듈을 사전에 정의하고 이를 기반으로 모듈화된 프롬프트의 주의 상태를 캐싱합니다. 예를 들어, 법적 서류나 템플릿이 자주 사용되는 환경에서 자주 등장하는 텍스트 모듈을 프롬프트 캐시로 저장하여 재사용함으로써 추론을 가속화합니다.

   주요 차이점: CacheBlend는 일부 토큰만 선택적으로 재계산하여 기존 KV 캐시를 유지한 채 품질을 높이는 데 중점을 두고 있고, PromptCache는 프롬프트 모듈 전체의 재사용에 초점을 맞추고 있습니다.

2. 캐시의 유연성과 적용 범위
   • CacheBlend: CacheBlend는 프리픽스(접두사) 위치에 상관없이 다양한 입력 청크의 KV 캐시를 융합할 수 있습니다. 즉, CacheBlend는 입력 청크 간 융합을 통해 서로 다른 위치의 KV 캐시를 동시에 활용하는 유연성을 갖추고 있습니다. 이를 통해 LLM이 긴 텍스트와 다양한 청크 조합을 다룰 때 고품질의 결과를 빠르게 생성합니다.
   • PromptCache: PromptCache는 모듈화된 프롬프트를 사전 정의하는 프롬프트 마크업 언어(PML)를 통해 모듈이 위치한 프리픽스 위치에서 재사용이 용이하도록 설계되었습니다. 주로 템플릿이나 반복적인 프롬프트를 사용하는 환경에서 성능이 극대화되지만, CacheBlend와 달리 입력 청크의 위치와 상관없이 재활용할 수 있는 유연성은 상대적으로 제한적입니다.
3. 처리 지연 감소 방식
   • CacheBlend: CacheBlend는 파이프라인 처리를 통해 일부 토큰의 선택적 재계산과 캐시 로드를 병렬로 수행합니다. 예를 들어, KV 캐시 로드와 재계산이 동시에 진행되어 지연 시간을 줄이고, 속도와 효율을 극대화합니다.
   • PromptCache: PromptCache는 반복적인 프롬프트 모듈의 캐시만 불러와 사용하며, 특정 토큰에 대한 재계산을 별도로 수행하지 않으므로, 각 모듈의 캐시 상태를 유지하는 방식으로 처리 지연을 줄입니다.
4. 응답 품질에 대한 초점
   • CacheBlend는 교차 주의(cross-attention) 정보를 유지하여 품질 손실을 최소화하는 것을 목표로 합니다. 선택적으로 KV를 재계산하여, 청크 간의 주의 정보를 융합하는 방식으로 고품질의 응답을 제공합니다.
   • PromptCache는 모듈화된 프롬프트 캐싱에 중점을 두며, 정보가 사전에 정의된 텍스트 모듈일 때 캐시의 재사용으로 효율성을 높이되, 캐시된 모듈 간의 교차 주의는 직접 다루지 않습니다.

요약

• CacheBlend는 다양한 입력 청크에 걸쳐 일부 토큰의 KV만 선택적으로 재계산하여 캐시를 융합하는 방식을 통해, 위치에 관계없이 고품질의 응답을 유지하며 속도 향상을 목표로 합니다.
• PromptCache는 반복적인 프롬프트 모듈을 모듈화하여 전체 모듈을 캐싱하는 방식으로 주로 템플릿 구조의 텍스트에 최적화되어 있습니다.

따라서 CacheBlend는 보다 유연하게 다양한 위치의 텍스트 청크를 융합하는 데 강점을 가지며, PromptCache는 사전 정의된 모듈화된 프롬프트를 재사용하는 환경에서 효율적입니다.

Q: 사용자 질의와 직접적으로 관련된 특정 KV Cache만 우선적으로 로드 하는 방법에 대해서 구체적으로 설명해줘

Answer

CacheBlend가 사용자 질의와 관련된 특정 KV Cache를 우선적으로 로드하는 방식은 사용자 질의에 있는 키워드와 각 청크(Chunk) 내의 중요한 토큰들 간의 연관성을 이용하여 결정됩니다. 이 과정은 질의와 청크 간의 상관성을 효율적으로 평가하고, 관련성이 높은 청크의 KV Cache만을 우선 로드함으로써 불필요한 메모리 접근을 최소화하고 효율성을 높이는 방식으로 설계되었습니다.

CacheBlend의 질의 기반 선택적 KV 로드 과정

1. 질의의 키워드 추출:
   • CacheBlend는 먼저 사용자 질의에서 중요한 키워드나 주제어를 추출합니다. 예를 들어, 사용자가 “인공지능의 발전과 미래”라는 질문을 했다면, “인공지능”, “발전”, “미래”와 같은 키워드가 추출될 수 있습니다.
   • 이를 위해 TF-IDF(term frequency-inverse document frequency) 또는 BERT와 같은 임베딩 기반의 문장 표현 모델을 사용하여 질의의 의미적 주제를 이해하고 주요 키워드를 뽑아냅니다.
2. 청크(Chunk)와 질의 간의 연관성 평가:
   • 추출된 키워드를 기준으로, SSD에 저장된 각 청크의 요약 정보 또는 메타데이터와의 연관성을 계산합니다. 이때 연관성 스코어가 높을수록 사용자 질의와 관련성이 큰 청크로 간주됩니다.
   • 각 청크의 요약 정보는 미리 계산된 키워드 리스트 또는 임베딩으로 관리될 수 있으며, 이를 통해 사용자가 입력한 질의와의 유사도를 빠르게 계산할 수 있습니다.
3. 우선 로드할 청크 결정 (연관성 스코어 기반):
   • 연관성 스코어가 높은 청크들만 우선적으로 로드하도록 설정합니다. 예를 들어, 상위 3개의 청크가 연관성 스코어가 높게 나왔다면, 이 3개의 청크에 대한 KV Cache를 우선 로드하고 나머지 청크들은 우선 로드 대상에서 제외됩니다.
   • 연관성 스코어 계산 방식은 다음과 같이 예시로 나타낼 수 있습니다: [ \text{Score}(Chunk_i) = \text{cosine similarity}(\text{Query embedding}, \text{Chunk}_i \text{ embedding}) ]
   • 이와 같은 방식으로 각 청크와 사용자 질의 간의 의미적 유사도를 계산하여 우선 로드 대상이 결정됩니다.
4. 동적 로드 관리 (Deferred Loading):
   • 만약 로드된 청크들만으로 충분한 응답을 생성하지 못할 경우, 그 다음으로 연관성 높은 청크를 추가로 로드할 수 있도록 동적 로드 방식을 사용합니다.
   • 이를 통해 초기에 필요한 청크들만 로드하여 응답을 생성하고, 추가 정보가 필요할 때만 선택적으로 청크를 로드해 메모리 효율성을 극대화합니다.

예시로 설명하기

• 사용자 질의: “인공지능의 역사와 최근 발전”
• 추출된 키워드: “인공지능”, “역사”, “발전”

SSD에 저장된 청크 예시:

1. Chunk A: “Artificial intelligence has its roots in simulating human cognition…”
2. Chunk B: “In the 21st century, deep learning has revolutionized AI…”
3. Chunk C: “AI ethics and legal implications are a topic of ongoing discussion…”

로드 과정 예시:

• Step 1: 각 청크의 요약 정보와 “인공지능”, “역사”, “발전” 간의 연관성 스코어 계산.
• Step 2: Chunk A와 Chunk B는 “역사”와 “발전” 키워드와 높은 연관성을 가지므로 우선적으로 로드.
• Step 3: Chunk C는 연관성 스코어가 낮아 우선 로드 대상에서 제외.

이와 같이 CacheBlend는 사용자 질의와 관련된 청크만 우선 로드하여, 효율적인 캐싱 및 메모리 사용을 가능하게 합니다.