GPT-4나 Claude 3와 같은 거대언어모델(LLM)과 대화할 때, 우리는 50페이지짜리 PDF를 분석하는 복잡한 작업에서도 거의 즉각적인 반응을 기대합니다. 하지만 대화가 길어질수록 응답 속도가 느려지거나, 로컬 환경에서 모델을 구동할 때 “메모리 부족(Out of Memory, OOM)” 오류가 발생하는 것을 경험해 보셨을 겁니다.
이러한 현상의 원인이자 동시에 해결책이 되는 것이 바로 KV 캐시(Key-Value Cache)입니다. KV 캐시는 자기회귀(autoregressive) 모델의 추론 속도를 최적화하는 가장 중요한 요소이지만, 그만큼 막대한 메모리 비용을 요구합니다. LLM 배포를 최적화하거나, RAG 파이프라인을 구축하거나, 혹은 왜 긴 문맥(Long-context)을 처리하는 모델이 거대한 GPU를 필요로 하는지 이해하고 싶다면 KV 캐시를 반드시 알아야 합니다.
핵심 개념: 왜 캐시가 필요한가?
KV 캐시를 이해하려면 LLM이 텍스트를 생성하는 방식을 살펴봐야 합니다. LLM은 자기회귀적(autoregressive)입니다. 즉, 한 번에 하나의 토큰(단어의 일부)을 생성하며, 새로운 토큰을 생성할 때마다 문맥을 유지하기 위해 이전에 생성된 모든 토큰을 참조합니다.
어텐션(Attention)의 문제점
트랜스포머(Transformer) 아키텍처의 “셀프 어텐션(Self-Attention)” 메커니즘은 모델이 현재 토큰의 의미를 파악하기 위해 과거의 토큰들을 되돌아볼 수 있게 해줍니다.
수학적으로 볼 때, 모델은 각 토큰에 대해 세 가지 벡터를 계산합니다.
- Query ($Q$): 현재 토큰이 찾고자 하는 정보.
- Key ($K$): 현재 토큰이 자신을 정의하는 정보(일종의 신분증).
- Value ($V$): 토큰이 가진 실제 콘텐츠/정보 값.
캐싱이 없다면, 100번째 토큰을 생성하기 위해 모델은 1번부터 99번 토큰까지의 $K$와 $V$ 행렬을 전부 다시 계산해야 합니다. 101번째 토큰을 생성할 때는 1번부터 100번까지 다시 계산해야겠죠. 이는 연산량을 이차함수적(quadratic)으로 증가시키며 엄청난 컴퓨팅 자원 낭비를 초래합니다.
해결책
과거 토큰에 대한 Key와 Value 벡터는 한 번 계산되면 변하지 않습니다. 따라서 우리는 이것을 GPU 메모리에 저장(캐시)해 둘 수 있습니다. 새로운 토큰을 생성할 때, 모델은 새로운 토큰에 대한 $Q, K, V$만 계산하고, 과거의 데이터는 캐시에서 $K$와 $V$를 불러와 사용하면 됩니다.
프로세스 시각화
KV 캐시를 적용했을 때 추론 과정이 어떻게 변하는지 보여주는 흐름도입니다.
graph TD
A["사용자 프롬프트 (입력)"] --> B["프리필(Prefill) 단계"]
B --> C["모든 입력 토큰의 K, V 계산"]
C --> D["KV 캐시(VRAM)에 저장"]
D --> E{"생성 루프 (디코딩)"}
E --> F["새로운 토큰 입력"]
F --> G["'새 토큰'의 Q, K, V만 계산"]
G --> H["캐시에서 과거 K, V 조회"]
H --> I["어텐션 수행 (새 Q vs 캐시된 K, V)"]
I --> J["다음 토큰 출력"]
J --> K["새 K, V를 캐시에 추가"]
K --> E
코드 구현: 기본적인 KV 캐시
딥러닝 라이브러리들이 이 과정을 자동으로 처리해주지만, 파이썬 코드로 로직을 살펴보면 이해가 훨씬 빠릅니다. 아래는 순전파(forward pass) 중에 캐시가 어떻게 초기화되고 업데이트되는지 보여주는 간단한 예시입니다.
import torch
class SimpleKVCache:
def __init__(self, max_seq_len, hidden_dim):
# 빈 캐시 텐서 초기화
# 형태: [배치 크기, 시퀀스 길이, 히든 디멘션]
self.k_cache = torch.zeros(1, max_seq_len, hidden_dim)
self.v_cache = torch.zeros(1, max_seq_len, hidden_dim)
self.current_pos = 0
def update(self, key_state, value_state):
"""
새로운 토큰의 Key와 Value 상태로 캐시를 업데이트합니다.
"""
seq_len = key_state.shape[1]
# 미리 할당된 버퍼에 새로운 K와 V를 삽입
self.k_cache[:, self.current_pos : self.current_pos + seq_len, :] = key_state
self.v_cache[:, self.current_pos : self.current_pos + seq_len, :] = value_state
self.current_pos += seq_len
# 현재 위치(current_pos)까지의 유효한 데이터만 슬라이싱하여 반환
return (
self.k_cache[:, :self.current_pos, :],
self.v_cache[:, :self.current_pos, :]
)
# 사용 예시
# 모델이 최신 토큰에 대해 계산한 new_k와 new_v가 있다고 가정
# cache_engine.update(new_k, new_v)
단계별 분석: 추론의 라이프사이클
- 프리필(Prefill) 단계:
프롬프트를 전송하면 모델은 모든 입력 토큰을 병렬로 처리합니다. 전체 프롬프트에 대한 $K$와 $V$ 행렬을 계산하여 캐시에 저장합니다. 이 단계는 보통 연산 제한(Compute-bound) 상태입니다(GPU가 수학 연산을 얼마나 빨리 처리하느냐가 관건). - 디코딩(Decoding) 단계 (토큰 생성):
모델이 한 번에 하나의 토큰을 생성하는 모드로 전환됩니다. 캐시된 데이터를 가져와서 단일 신규 토큰만 처리합니다.- 병목의 이동: 이 단계는 주로 메모리 대역폭 제한(Memory-bandwidth bound) 상태가 됩니다. GPU는 실제 연산을 수행하는 시간보다 거대한 KV 캐시 데이터를 VRAM에서 연산 코어로 이동시키는 데 더 많은 시간을 소비합니다.
- 메모리 폭증:
컨텍스트 길이가 길어질수록 KV 캐시는 선형적으로 증가합니다.- 참고: 숫자 ‘2’는 K와 V 두 가지를 저장하기 때문입니다.
- 컨텍스트 윈도우 한계:
결국 KV 캐시가 가용 VRAM을 가득 채우게 됩니다. 이렇게 되면 모델은 더 이상 토큰을 처리할 수 없게 되어, 컨텍스트가 잘리거나(truncation) 프로세스가 중단(crash)됩니다.
메모리 영향 분석
다음 표는 왜 긴 컨텍스트(예: 128k)를 가진 모델을 구동할 때, 모델 가중치(Weights) 외에 KV 캐시만으로도 막대한 VRAM이 필요한지 보여줍니다.
가정: Llama-3-70B 모델, Float16 정밀도 (요소당 2바이트) 기준 대략적 수치.
| 컨텍스트 길이 (토큰) | KV 캐시 크기 (약) | 소비자 하드웨어에 미치는 영향 |
|---|---|---|
| 4,096 | ~0.6 GB | 미미함. 대부분의 카드에서 구동 가능. |
| 32,000 | ~5.0 GB | 상당함. 12GB 이상의 VRAM 카드 필요. |
| 128,000 | ~20.0 GB | 치명적. 모델 가중치까지 합치면 RTX 4090(24GB) 용량 초과. |
전문가 팁: KV 캐시 최적화 방법
LLM을 배포하거나 애플리케이션을 구축하고 있다면 기본 캐싱만으로는 충분하지 않습니다. 다음과 같은 고급 기술을 활용하여 메모리 사용량을 관리하세요.
- PagedAttention (vLLM 사용): 전통적인 캐싱은 연속된 메모리 블록을 예약하므로 단편화(fragmentation)와 낭비가 심합니다. vLLM은 운영체제의 가상 메모리에서 착안한 PagedAttention을 사용하여 키와 값을 불연속적인 메모리 블록에 저장함으로써 처리량(throughput)을 최대 24배까지 높입니다.
- GQA (Grouped Query Attention): Llama 3와 같은 최신 모델들은 GQA를 사용합니다. 모든 쿼리(Query) 헤드마다 고유한 Key/Value 헤드를 갖는 대신, 여러 쿼리 헤드가 KV 헤드를 공유하게 합니다. 이는 성능 저하를 최소화하면서 캐시 크기를 획기적으로(보통 8배) 줄여줍니다.
- KV 캐시 양자화(Quantization): KV 캐시 자체를 압축할 수 있습니다 (예: FP16에서 INT8 또는 FP8로 변환). 이렇게 하면 동일한 VRAM 용량으로 최대 컨텍스트 길이를 사실상 두 배로 늘릴 수 있습니다.
- FlashAttention: 백엔드에서 FlashAttention-2를 활용하고 있는지 확인하세요. 이는 어텐션 메커니즘의 읽기/쓰기 작업을 최적화하여 디코딩 단계의 메모리 대역폭 병목 현상을 줄여줍니다.
KV 캐시는 최신 LLM 효율성의 핵심 엔진으로, VRAM 공간을 내어주고 연산 속도를 얻는 트레이드오프(trade-off)의 산물입니다. 우리가 즐기는 매끄러운 대화 경험을 가능하게 해주지만, 동시에 긴 컨텍스트 애플리케이션의 주된 제약 사항이기도 합니다. PagedAttention이나 양자화 같은 기술을 적절히 활용한다면, 제한된 하드웨어에서도 더 긴 문맥을 효율적으로 다룰 수 있습니다.