Research Paper Review : CACHE BLEND, Fast Large Language Model Serving for RAG with Cached Knowledge Fusion

 오늘 이 게시물에서는 CACHE BLEND라는 논문에 대해서 설명해보도록 하겠다. 

https://arxiv.org/abs/2405.16444

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

 

Background 

 요즘 LLM에서는 inference 시, RAG와 같은 retrieval 기술을 사용해 query와 연관된 text chunk들을 input 앞에 붙여(concat) 주는 방식이 널리 활용된다. 이러한 retrieval-augmented 방식은 이미 여러 연구와 실전 서비스들을 통해 LLM의 inference 성능을 크게 향상시키는 정석적인 방법으로 받아들여지고 있다.

<LLM input of nowadays>

이때 여러 사용자들의 입력에서 동일한 Chunk가 반복적으로 등장하는 상황이 자주 발생할 수 있다. 예를 들어 회사 내규를 기반으로 질문에 답하는 Chatbot을 생각해보자. 이 Chatbot은 RAG를 통해 회사 내규 문서를 retrieval하고, 그 내용을 Query 앞에 붙여 답변을 생성할 것이다. 이런 환경에서는 특정 규정이나 조항에 해당하는 동일한 text Chunk가 여러 사용자 입력에서 반복적으로 사용될 가능성이 매우 높다. 따라서 해당 Chunk에 대한 KV cache를 한 번 미리 계산해두고 그대로 재사용할 수 있다면, 매번 재계산해야 하는 prefill overhead를 크게 줄일 수 있을 것이다.

 

 이처럼 pre-computed KV cache를 재사용하는 방법론 중 하나로 잘 알려진 것이 prefix caching이다. 이 방법은 RAG 데이터베이스에 포함된 모든 text chunk들에 대해 각각 KV cache를 미리 계산해 저장해둔다. 이때 각 chunk는 서로 연결된 상황을 고려하여 계산되는 것이 아니라, 각 chunk가 항상 input의 맨 앞(prefix)에 위치한다고 가정한 상태에서 독립적으로 KV를 생성한다. 그리고 inference 시 실제 input의 prefix 영역에 특정 chunk가 등장하면, 그 chunk의 pre-computed KV cache를 그대로 불러와 사용하는 방식이다.

하지만 이 접근법에는 명확한 한계가 있다. 최신 RAG 기반 시스템에서는 하나의 query에 여러 개의 chunk가 연속적으로 붙는 경우가 흔하다. 그럼에도 prefix caching은 맨 앞에 오는 첫 번째 chunk에 대해서만 KV 재사용이 가능하다. 뒤에 따라붙는 chunk들은 KV cache를 재사용할 수 없는데, 이는 미리 계산된 KV cache가 이전 chunk들과의 attention 정보가 전혀 반영되지 않은 상태이기 때문이다. 그 결과, prefix caching은 실제로는 전체 prefill 시간을 크게 줄이지 못한다는 문제가 존재한다.

 

 두 번째 방법은 full KV reuse이다. 이 방법은 prefix caching과 동일하게, 우선 각 Chunk의 KV cache를 서로 독립적으로 미리 계산해 둔다. 그리고 inference 단계에서는, 이 미리 계산해 둔 KV cache를 그대로 재사용하되, Positional Embedding(PE)만 실제 위치에 맞게 조정해서 사용하자는 아이디어다. 물론 첫 번째 Chunk가 아니라면 원래는 이전 Chunk들과 attention을 해야 하는데, full KV reuse에서는 이 부분이 완전히 무시되므로 성능 하락이 생길 수밖에 없다. 그럼에도 불구하고 prefill 비용을 큰 폭으로 줄일 수 있기 때문에 시도해볼 가치가 있다는 것이다.

이 방식이 가능한 이유는, 최근 LLM들이 거의 모두 RoPE(Rotary Positional Embedding) 을 사용한다는 데에 있다. RoPE는 token의 절대 위치가 아니라 상대적 위치(relative position) 만을 attention 계산에 반영한다. 즉, token들이 실제 input에서 어디에 있었든, Chunk 내부 token들끼리의 상대적 거리만 동일하다면 attention 구조는 완전히 동일하게 유지된다. 예를 들어 Chunk의 길이가 5라고 하면, 우리는 이 Chunk 안의 token들을 항상 [0,1,2,3,4] 위치에 있다고 가정하고 KV cache를 미리 계산할 수 있다. 실제로 이 Chunk가 input 상에서는 [11,12,13,14,15] 위치에 있었더라도, Chunk 내부 token들 간의 attention은 RoPE의 성질 때문에 전혀 바뀌지 않는다.

변화가 필요한 부분은 Chunk 밖의 token들과의 attention이다. Chunk가 실제 input에서 어떤 위치에 있었는지를 반영하려면, 미리 계산해 둔 Key 벡터들에 추가적인 회전(rotation) 을 적용하여 RoPE의 각도를 보정해주어야 한다. 이 회전은 “실제 시작 위치 − 우리가 가정한 시작 위치(0)”만큼 일정하게 적용하면 되기 때문에, Chunk 내부 구조를 망가뜨리지 않으면서도 외부 token들과의 상대적 위치 관계만 맞춰줄 수 있다.

정리하면, RoPE를 사용하기 때문에 우리는 Chunk의 내부 구조를 완전히 그대로 재사용할 수 있고, Chunk가 실제 input 상에서 위치한 절대 위치는 Key 벡터 전체를 한 번 회전시키는 것만으로 보정할 수 있다. 이런 구조적 특성 덕분에 full KV reuse는 원천적으로 가능해지며, 이는 prefill 단계를 크게 줄여 inference 속도를 높일 수 있는 중요한 아이디어가 된다.

<뒤에 Chunk들의 KV cache가 섞여 있는것은 chunk들끼리 attention을 했다는 것을 의미한다>

 

 

하지만 앞서 언급했듯이, full KV reuse는 모든 KV를 re-compute하는 것에 비해 명백한 성능 drop이 존재한다. 밑의 그림을 참조하자. 

 저자들은 여기서 한 가지 흥미로운 아이디어를 제안한다. full KV reuse를 기본 전략으로 사용하되, 실제로 이전 Chunk들과의 cross-attention이 필요한 token들에 대해서만 KV를 다시 계산(re-compute)하자는 것이다. 즉, 대부분의 token은 미리 계산해둔 KV cache를 그대로 쓰고, attention에서 중요한 역할을 하는 소수의 token만 선별적으로 재계산하는 방식이다.

또한 KV cache를 layer 단위로 GPU에 로드하는 상황을 고려하면, 재계산의 비용은 더욱 줄어든다. 만약 “다음 layer의 KV cache를 로드하는 시간”이 “현재 layer에서 필요한 소수의 token들의 KV를 재계산하는 시간”보다 더 길다면, 이 재계산은 사실상 추가적인 비용이 들지 않는 것과 동일하다. 이유는 간단하다. 어차피 로딩을 기다리는 동안 계산 자원은 놀고 있으므로, 그 유휴 자원을 이용해 필요한 KV만 병렬로 재계산하면 되기 때문이다.

요약하면, 이 방식은 full KV reuse의 장점을 유지하면서도, cross-attention이 중요한 token들만 selective recompute를 통해 보정함으로써 성능 저하를 최소화하고, 실제 overhead도 거의 발생시키지 않는 효율적인 절충안이라고 볼 수 있다.

Methods 

 그렇다면 저자들의 제안이 실제로 유효한 방식인지 확인하려면 두 가지 핵심 질문에 답해야 한다.
1) 이전 Chunk들과의 cross-attention이 필요한 token들은 정말 전체 token 중 극히 일부(소수)인가?
2) 그렇다면, 그 소수의 token들을 어떻게 효과적으로 골라낼 수 있는가?

먼저 첫 번째 질문을 살펴보자. 만약 어떤 token이 이전 Chunk들과의 cross-attention을 크게 활용해야 하는 token이라면, 그 token의 KV를 “미리 계산해둔 값”과 “실제로 재계산했을 때의 값”은 큰 차이를 보이게 된다. 저자들은 이 차이를 KV deviation이라고 정의한다. 따라서 특정 token의 KV deviation이 크다는 것은, 그 token이 실제 inference에서 앞선 Chunk들과 강하게 interaction해야 한다는 의미로 해석할 수 있다.

논문에서 제시된 그림을 보면, KV deviation의 CDF가 거의 선형적으로 증가하는 형태를 보인다. 이는 KV deviation이 큰 token들과 작은 token들이 전 구간에 걸쳐 고르게 분포한다는 뜻이다. 다시 말해, KV deviation이 특히 큰 token들만 선별해 re-compute한다면, 전체 token 중 극히 일부만 계산해도 충분히 품질을 보정할 수 있다는 의미가 된다. 즉, selective recompute 전략이 실제로 유효할 수 있음을 보여주는 근거가 된다.

 

 이제 두 번째 질문, 즉 “그 소수의 token들을 어떻게 효율적으로 골라낼 것인가?”에 대해 답해보자. 가장 단순한 방법은 모든 token에 대해 KV를 다시 계산한 뒤, 기존에 저장된 KV와 비교해 KV deviation이 큰 token을 골라내는 것이다. 그러나 이는 전혀 실용적이지 않은 접근이다. 어차피 전체를 re-compute할 것이라면, 굳이 deviation을 측정할 필요 없이 그냥 re-computed KV cache를 그대로 사용하면 되기 때문이다.

다행스럽게도 저자들은 중요한 관찰을 하나 제시한다. 인접한 layer들 사이에서는 KV deviation이 큰 token들이 거의 동일하게 유지되는 경향이 있다는 점이다. 즉, 어떤 token이 한 layer에서 “앞선 Chunk와 강한 cross-attention이 필요하다”고 판단되면, 다음 layer에서도 동일한 token이 여전히 중요한 역할을 한다는 것이다. 이 특성 덕분에, 매 layer마다 전체 token을 재검사할 필요 없이, 한 layer에서 얻은 정보를 다음 layer에서도 그대로 활용할 수 있다.

이 관찰을 기반으로 하면, 사실 layer 1에서만 KV deviation을 직접 계산해 토큰을 선별하고, 이후 layer에서는 단순히 layer 1에서 선정한 token들에 대해서만 re-compute를 수행하는 전략도 충분히 feasible하다고 볼 수 있다. 그러나 저자들은 더 높은 성능을 얻기 위해 한 단계 발전된 접근을 제안한다.

우선 layer 1에서 비교적 넉넉한 수의 token 후보군을 선정해 둔다. 그리고 layer가 깊어질수록, 그 후보군 중에서도 KV deviation이 특히 큰 token들만 남겨 추려가는 방식을 사용한다. 즉, layer 1에서는 넓게 뽑고, 이후 layer에서는 좁혀 나가는 점진적 필터링 전략이다. 이미 선별된 token들에 대해서는 re-compute를 수행하므로, 해당 token들의 KV deviation은 거의 추가 비용 없이 자연스럽게 계산할 수 있다. 이를 활용해 매 layer에서 더 정교하게 token 집합을 줄여가는 것이다.

이 방식은 처음 layer에서만 모든 것을 결정해버리는 deterministic한 전략보다, 인접 layer 간의 KV deviation 패턴을 더 잘 반영하여 더 효과적인 token selection이 가능하다는 장점이 있다.

 

System Design

 앞서 간단히 언급했듯이, 만약 layer (L–1)에서 선정된 token들의 재계산(re-compute)에 걸리는 시간이, 다음 layer의 KV cache를 로드하는 시간보다 짧다면, 우리는 거의 추가적인 overhead 없이 re-compute를 수행할 수 있다. 즉, 로딩을 기다리는 동안 idle 상태인 계산 자원을 활용하는 것이기 때문에, 이 구간에서 발생하는 re-compute는 사실상 “공짜”에 가깝다.

이를 실현하려면 각 layer마다 KV cache 로드 시간(load time) 을 알고 있어야 하고, 해당 시간 안에 끝낼 수 있는 token re-compute 비율(recompute ratio) 을 결정해야 한다. 하지만 recompute ratio를 무작정 낮추면 안 된다. 논문에서는 recompute ratio가 15% 미만으로 떨어지는 순간 성능 하락이 크게 발생한다고 보고한다. 따라서 저자들은 최종 recompute ratio를 다음과 같이 설정한다.

 

 recompute ratio = max (15%, ratio where re-compute time = load time)

 

즉, 성능을 보장하기 위한 최소 비율(15%)을 유지하되, 가능하다면 로드 시간과 동일한 수준까지 re-compute를 밀어 넣어 “idle 구간”을 최대한 활용하는 방식이다. 여기서 load time 은 KV cache가 저장된 하드웨어의 특성(대역폭, IO 속도 등)을 Loading Controller(추후 설명)가 이미 알고 있으므로, 전통적인 계산 공식들을 이용해 쉽게 추정할 수 있고, re-compute time 또한 모델 구조(LLM 종류), 선택된 token 수, hidden dimension 등을 알면 비교적 간단하게 계산할 수 있다. 결과적으로 이 두 값을 비교해 각 layer별로 적절한 recompute ratio를 동적으로 결정하는 것이 가능해진다.  

 

이 전략을 반영한 시스템 개요도는 다음과 같다. 

 

Loading Controller는 각 KV cache가 디스크·메모리·버퍼 등 어떤 스토리지 계층에 저장되어 있는지에 대한 정보를 모두 알고 있으며, 이를 바탕으로 layer별 loading time 과 re-compute time 을 계산할 수 있다. 이 두 값을 비교하여 가장 효율적인 re-compute ratio 를 결정한 뒤, 그 값을 KV cache Fusor에게 전달한다.

KV cache Fusor는 전달받은 recompute ratio에 따라 선택된 token들의 KV를 재계산하고, 이후 여러 Chunk에서 업데이트된 KV cache들을 하나의 일관된 KV 세트로 merge(fuse) 하여 저장한다. 결국 이 흐름을 통해, 시스템은 각 layer마다 로딩과 재계산 사이의 최적 균형점을 찾아 효율적으로 KV cache를 관리하게 된다.

 

Experiments 

 아래 그림을 보면, CacheBlend 방법이 Full KV reuse 수준으로 TTFT(Time To First Token)를 효과적으로 단축하면서도, 동시에 full-recompute와 거의 동일한 모델 성능을 달성하고 있음을 확인할 수 있다. 즉, CacheBlend는 두 접근법의 장점을 모두 취해, 지연(latency)과 품질 간의 균형을 매우 성공적으로 달성한 방법임을 실험 결과가 잘 보여준다.