대규모 언어 모델의 KV 캐시: 설계, 최적화 및 추론 가속

대규모 언어 모델의 KV 캐시: 설계, 최적화 및 추론 가속

Transformer는 대규모 언어 모델(LLM)의 핵심 구조이지만, 긴 컨텍스트를 처리하는 추론 과정은 많은 계산 (시퀀스 길이의 제곱) 비용이 듭니다. KV 캐시—Key/Value Cache는 자기회귀(autoregressive) 디코딩 과정에서 이전에 계산된 Key, Value의 값을 캐싱하여 재사용함으로써 추론 속도를 선형으로 줄여주는 핵심 최적화 기법 중 하나입니다. 이 글에서는 Efficiently Scaling Transformer Inference 논문을 바탕으로 KV 캐시의 구조, 동작 원리, 구현 방식, 최적화 전략에 대해 심층 분석합니다.

** You can find the English verion of this content at this page (https://markbyun.blogspot.com/2025/05/kv-cache-in-large-language-models.html)

---

1. KV 캐시란 무엇인가?

Transformer의 어텐션 레이어는 입력 시퀀스 전체에 대해 어텐션 가중치를 계산합니다. 추론 과정에서 토큰을 하나씩 생성할 때마다 이 연산이 반복되어 시간 복잡도 $O(n^2)$가 발생합니다. KV 캐시는 이전 토큰들의 Key, Value 텐서를 저장해 두고 이를 재활용함으로써 $O(n)$ 시간 복잡도로 줄입니다. 이는 실시간 채팅봇, 코드 어시스턴트와 같은 응용에서 필수적인 기술입니다.

---

2. 핵심 설계 원칙

• 시간적 재사용: 이전 단계에서 계산한 Key, Value를 이후 단계에서 재활용
• 메모리 효율성: 메모리 파편화를 방지하며 GPU 메모리 자원 절약
• 하드웨어 정렬: 최신 GPU/TPU의 메모리 접근 패턴에 맞는 캐시 레이아웃 설계
• 확장성: 다수 사용자 세션에서 다양한 길이의 입력 처리 가능

---

3. KV 캐시의 동작 원리

추론 중 각 Transformer 레이어에서 t번째 토큰에 대해 아래 연산이 수행됩니다:

$Q_t, K_t, V_t = Linear(H_t)$

$Attentiont = Softmax(Q_t \times [K_1 ... K_t]^T) \times [V_1 ... V_t]$

이 때 K₁~Kt-1과 V₁~Vt-1을 매번 다시 계산하지 않고, 별도의 캐시에 저장해두고 재사용합니다. 캐시는 보통 다음과 같은 텐서 형태로 유지됩니다:

[batch_size, num_heads, max_seq_len, head_dim]

매 토큰 생성 시 Kt, Vt만 새롭게 계산되어 캐시에 추가됩니다. 이후 어텐션은 전체 캐시를 대상으로 수행됩니다.

---

4. 캐시 저장 및 갱신 전략

• 사전 할당: 전체 최대 길이 기준으로 캐시를 미리 할당해 메모리 파편화 방지
• 페이징 기법 (vLLM): KV 블록을 고정 크기 페이지로 나눠 재사용 및 회수
• 초기 재사용 (TensorRT-LLM): 프롬프트 부분은 고정으로 한 번만 캐시
• 스트리밍 대응: 실시간으로 토큰이 추가될 때 캐시도 실시간으로 확장

---

5. 구현 시 고려 사항

Multiquery attention has lower memory cost to load the KV cache when sharded over batch

• 캐시 무효화: 시퀀스 종료 시 적절히 캐시 제거
• 동시성 처리: 멀티 유저 세션의 캐시 격리 및 병렬 처리 (vLLM은 가상 블록과 페이지 방식으로 해결)
• Multi-query Attention: 검색 기반 생성에서 공유 encoder로부터 K/V 처리
• 메모리 레이아웃 최적화: 텐서코어 및 GPU 워프 정렬을 고려한 배치 구성

---

6. 성능 최적화 기법

컨텍스트 길이가 길어지면 저장해야할 KV캐시 또한 증가하므로 캐시 메모리를 효율적으로 관리하는 것이 매우 중요한 문제가 됩니다. 따라서 아래와 같은 기법을 적용하여 캐시메모리 및 성능 최적화 작업을 진행합니다.

• 양자화 캐시: FP8 또는 INT8로 K/V를 저장하여 메모리 절약
• FlashAttention + 캐시: 롱컨텍스트 대응을 위한 조합 전략
• 페이징 KV 캐시 (vLLM): 메모리 페이지 기반 할당 및 LRU 교체 전략
• 분리 캐시 구조: 프롬프트와 사용자 입력을 분리하여 재사용 효율 증대

---

7. 실험 결과 및 성능 지표

• 첫 토큰 생성 시간 최대 5배 가속 (TensorRT-LLM 기준)
• 메모리 사용량 절감 (양자화 및 블록 재사용 적용 시)
• 1000개 이상의 동시 시퀀스 처리 가능 (vLLM 기준, A100 단일 GPU)

---

8. 오픈소스 구현 사례

• vLLM: 블록 페이징 기반의 고속 캐시 시스템
• TensorRT-LLM: 초기 캐시 재사용, 커스텀 CUDA 커널 최적화
• llama.cpp: CPU 및 GPU 온디바이스 경량 캐시

---

9. 결론

KV 캐시는 대규모 언어 모델의 실시간 추론을 가능하게 해주는 핵심 기술입니다. 복잡도를 줄이고, 다중 사용자 응답 처리 성능을 극적으로 향상시키며, 다양한 하드웨어 환경에서도 추론 속도를 유지할 수 있도록 도와줍니다. vLLM, TensorRT-LLM과 같은 시스템은 KV 캐시를 중심으로 혁신적인 최적화를 도입하여 지연 시간 최소화, 처리량 극대화, 메모리 효율 확보를 모두 달성하고 있습니다.

---

참고자료

• Efficiently Scaling Transformer Inference 논문
• vLLM: 고속 LLM 서빙 시스템
• HuggingFace: 추론 최적화 블로그
• PyTorch: 캐시 효율화 관련 소스 코드