[AI/LLM] KV Cache(Key-Value Cache)에 대해 자세히 알아보자! (정의, 원리, 장단점, 실습)

1. 들어가며

ChatGPT나 Claude 같은 생성형 AI를 사용할 때, 문장이 길어질수록 답변 속도가 미세하게 느려지거나 GPU 메모리가 급격히 차오르는 현상을 경험해 보셨을 겁니다.

LLM은 기본적으로 자기회귀(Auto-regressive) 모델입니다. 즉, "나는 학교에"라는 문장이 주어지면 "간다"를 예측하고, 다시 "나는 학교에 간다"를 입력으로 넣어 "."을 예측합니다.

문제는 이 과정에서 이미 계산했던 "나는 학교에" 부분의 연산을 매번 처음부터 다시 반복한다는 것입니다. 문장이 길어질수록 이 중복 연산은 기하급수적으로 늘어나, 치명적인 비효율을 초래합니다.

이 문제를 해결하기 위해 등장한 기술이 바로 KV Cache(Key-Value Cache)입니다. "계산(Compute)을 아끼기 위해 메모리(Memory)를 쓴다"는 철학을 가진 이 기술은 현대 LLM 서빙의 필수 요소가 되었습니다.

이번 포스팅에서는 KV Cache의 정의와 원리, 그리고 이를 직접 구현하여 성능을 최적화하는 방법까지 A to Z를 파헤쳐 보겠습니다.

---

2. KV Cache란 무엇인가?

 

2.1 정의

KV Cache는 트랜스포머(Transformer) 모델의 추론(Inference) 단계에서, 이전 토큰들에 대한 Key(K)와 Value(V) 벡터를 미리 저장(Caching)해 두고 재사용하는 기술입니다.

• Before (No Cache): 매 스텝마다 전체 문장의 $Q, K, V$를 처음부터 다시 계산 >>> $O(N^2)$ 연산 비용
• After (With Cache): 새로 들어온 토큰의 $Q, K, V$만 계산하고, 이전 $K, V$는 꺼내서 쓴다 >>> $O(N)$ 연산 비용

 

2.2 탄생 배경: 어텐션 메커니즘의 비효율

트랜스포머의 핵심인 Self-Attention 수식은 다음과 같습니다.

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

LLM이 다음 단어를 예측할 때, 현재 시점의 단어(Query)는 과거의 모든 단어(Key, Value)와 연관성을 계산해야 합니다. 하지만 과거의 단어들은 변하지 않으므로, 이들의 $K$와 $V$ 값도 변하지 않습니다. 변하지 않는 값을 매번 다시 계산하는 낭비를 없애는 것이 KV Cache의 핵심입니다.

 

---

3. KV Cache의 동작 원리 (Step-by-Step)

예를 들어, "Time flies" 다음에 "fast"를 생성하는 과정을 보겠습니다.

 

3.1 상황 1: KV Cache 미사용 (비효율적)

1. Step 1: 입력 ["Time"] >>> 연산 수행 >>> ["flies"] 생성.
2. Step 2: 입력 ["Time", "flies"] >>> "Time" 연산 다시 수행 + "flies" 연산 >>> ["fast"] 생성.
3. Step 3: 입력 ["Time", "flies", "fast"] >>> "Time", "flies" 연산 다시 수행...

 

3.2 상황 2: KV Cache 사용 (효율적)

1. Step 1 (Prefill): 입력 ["Time"] >>> $K_1, V_1$ 계산 및 저장 >>> ["flies"] 생성.
2. Step 2 (Decoding):
   • 입력으로 전체 문장이 아닌, 새로운 토큰 ["flies"]만 넣음.
   • "flies"의 $K_2, V_2$만 계산.
   • 저장해둔 $K_1, V_1$을 불러와 합침 ($[K_1, K_2]$, $[V_1, V_2]$).
   • Attention 수행 >>> ["fast"] 생성.
   • $K_2, V_2$를 캐시에 추가 저장.

결과: 문장 길이가 길어져도 연산량은 선형적(Linear)으로만 증가하여 속도가 획기적으로 빨라집니다.

 

---

4. 특징 및 장단점

 

4.1 ✅ 장점 (Pros)

1. 추론 속도(Latency) 향상: 중복 연산(Matrix Multiplication)을 제거하여, 특히 긴 문장을 생성할 때 속도가 비약적으로 빨라집니다.
2. 연산 비용(FLOPs) 절감: GPU가 수행해야 할 곱셈 연산 횟수가 줄어들어 전력 소모가 줄어듭니다.

4.2 ❌ 단점 (Cons)

1. GPU 메모리(VRAM) 사용량 폭증: 계산된 $K, V$ 벡터를 GPU 메모리에 계속 들고 있어야 합니다.
   • 배치 크기(Batch Size)나 문장 길이(Context Length)가 커지면 메모리 부족(OOM)이 발생하기 쉽습니다.
2. 구현 복잡도: 캐시를 관리하고 업데이트하는 로직이 추가되어야 합니다.

 

---

5. [Deep Dive] 메모리 용량 계산과 최적화 (vLLM)

KV Cache가 얼마나 많은 메모리를 먹는지 계산해보면 놀랍습니다.

 

5.1 메모리 용량 공식

한 토큰당 필요한 KV Cache 메모리 크기는 다음과 같습니다. (FP16 기준 2 bytes)

$$\text{Size} = 2 \times 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}}$$

• $2$: Key와 Value 두 개.
• $2$: FP16 (2 bytes).

 

예시: LLaMA-2 70B 모델 (Sequence Length 128K일 때)

배치 사이즈 32로 설정하면, KV Cache만 무려 1TB가 넘는 메모리가 필요할 수 있습니다. 이는 A100 GPU 여러 장을 합쳐도 감당하기 힘든 수준입니다.

 

5.2 해결책: PagedAttention (vLLM)

이 메모리 문제를 해결하기 위해 등장한 것이 vLLM 라이브러리의 PagedAttention입니다.

• 문제: 기존에는 메모리를 미리 연속적으로 할당해두어, 안 쓰는 공간(Fragmentation) 낭비가 심했습니다.
• 해결: 운영체제(OS)의 가상 메모리 페이징 기법을 도입하여, 메모리를 블록 단위로 쪼개서 필요할 때마다 동적으로 할당합니다.
• 효과: 메모리 낭비를 줄여 동일한 GPU에서 더 많은 요청(Batch)을 처리할 수 있게 되었습니다.

 

---

6. [실습] Python & PyTorch로 KV Cache 직접 구현하기

실제 코드로 KV Cache가 어떻게 작동하는지 구현해 보겠습니다. 이해를 돕기 위해 간단한 GPT 스타일의 모델을 가정합니다.

 

6.1 기본 설정 및 Attention 클래스

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head, max_len):
        super().__init__()
        self.n_head = n_head
        self.head_dim = d_model // n_head
        
        # Q, K, V Projection
        self.c_attn = nn.Linear(d_model, 3 * d_model)
        self.c_proj = nn.Linear(d_model, d_model)
        
        # 캐시 저장소 (Buffer로 등록하여 모델 state에 포함되지 않게 함 - 선택사항)
        # 실제로는 forward 시에 외부에서 주입받는 방식을 많이 씀
        self.register_buffer("bias", torch.tril(torch.ones(max_len, max_len))
                                     .view(1, 1, max_len, max_len))

    def forward(self, x, past_kv=None, use_cache=False):
        B, T, C = x.size() # Batch, Time(Length), Channel
        
        # 1. Q, K, V 추출
        qkv = self.c_attn(x)
        q, k, v = qkv.split(C, dim=2)
        
        # Head 분리 (B, n_head, T, head_dim)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # 2. KV Cache 로직 적용
        if use_cache:
            if past_kv is not None:
                past_k, past_v = past_kv
                # 기존 캐시에 현재 스텝의 k, v를 붙임 (Concatenate)
                k = torch.cat([past_k, k], dim=2)
                v = torch.cat([past_v, v], dim=2)
            
            # 현재까지의 k, v를 반환 (다음 스텝을 위해)
            current_kv = (k, v)
        else:
            current_kv = None

        # 3. Attention Score 계산
        # (B, n_head, T_q, T_k) -> 쿼리 길이와 키 길이는 다를 수 있음 (캐시 사용 시)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        
        # 마스킹 (Causal Masking)
        # 캐시를 쓸 때는 쿼리가 1개(마지막 토큰)이므로 마스킹 처리가 조금 다름
        # 여기서는 단순화를 위해 전체 시퀀스 길이에 맞춰 마스킹 적용
        total_len = k.size(2)
        att = att.masked_fill(self.bias[:, :, T-1:T, :total_len] == 0, float('-inf'))
        
        att = F.softmax(att, dim=-1)
        y = att @ v # (B, n_head, T, head_dim)
        
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y), current_kv

 

6.2 추론 루프 (Inference Loop) 비교

Case A: KV Cache 미사용

매번 문장이 길어지며 전체를 다시 계산합니다.

# 초기 입력
input_ids = torch.tensor([[1, 2, 3]]) # "I am a"

for _ in range(5):
    # 모델에 전체 시퀀스 입력
    logits, _ = model(input_ids, use_cache=False)
    
    # 마지막 토큰 예측
    next_token = logits[:, -1, :].argmax(dim=-1).unsqueeze(0)
    
    # 입력에 추가 (Concatenate)
    input_ids = torch.cat([input_ids, next_token], dim=1)

 

Case B: KV Cache 사용

첫 번째만 전체 계산(Prefill), 이후엔 마지막 토큰만 입력(Decoding)합니다.

# 초기 입력
input_ids = torch.tensor([[1, 2, 3]]) # "I am a"
past_kv = None # 캐시 초기화

# 1. Prefill 단계 (최초 1회)
logits, past_kv = model(input_ids, past_kv=None, use_cache=True)
next_token = logits[:, -1, :].argmax(dim=-1).unsqueeze(0)
generated_ids = [next_token]

# 2. Decoding 단계 (반복)
for _ in range(4):
    # ★ 핵심: 전체 input_ids가 아니라, 방금 만든 'next_token'만 넣음
    # past_kv에 이전 정보가 다 들어있기 때문
    logits, past_kv = model(next_token, past_kv=past_kv, use_cache=True)
    
    next_token = logits[:, -1, :].argmax(dim=-1).unsqueeze(0)
    generated_ids.append(next_token)

print("생성 완료")

 

---

7. 한 눈에 비교: KV Cache 사용 전후

비교 항목	No Cache (미사용)	With KV Cache (사용)
입력 데이터	전체 시퀀스 (점점 길어짐)	마지막 토큰 1개
Attention 연산량	$O(N^2)$ (이차 함수 증가)	$O(N)$ (선형 증가)
GPU 메모리 (VRAM)	적게 사용 (활성값만 저장)	많이 사용 (과거 KV 모두 저장)
주요 병목	Compute Bound (연산 속도)	Memory Bound (메모리 대역폭)
생성 속도 (TPS)	문장이 길어질수록 급격히 느려짐	일정하게 빠름

 

---

8. 마치며

KV Cache는 LLM이 실시간 서비스로 거듭나기 위한 '필수 불가결한 타협'입니다. 연산 속도를 얻는 대신 메모리라는 비용을 지불하는 셈이죠.

하지만 최근에는 이 메모리 비용조차 줄이기 위해 다음과 같은 기술들이 발전하고 있습니다.

• MQA (Multi-Query Attention) / GQA (Grouped-Query Attention): Key와 Value 헤드(Head)의 개수를 줄여 캐시 용량을 1/8 수준으로 압축 (LLaMA-2, 3 등이 채택).
• PagedAttention (vLLM): 메모리 파편화를 막아 실질적인 배치 처리량을 극대화.

오늘 소개한 KV Cache의 원리를 이해한다면, vLLM이나 Hugging Face의 use_cache=True 옵션이 내부적으로 어떻게 마법을 부리는지 명확히 아실 수 있을 것입니다.

다음 프로젝트에서는 KV Cache를 적극 활용하여, 번개처럼 빠른 AI 서비스를 구축해 보세요! 🚀