KV Cache

KV Cache는 Transformer 추론 시 이전 토큰의 Key·Value 텐서를 저장·재사용하는 메모리 최적화 기법

• Self-Attention의 중복 연산을 제거하여 연산량을 O(n²)에서 O(n)으로 줄이는 캐싱 메커니즘
• 각 Transformer 블록마다 독립적으로 유지되는 레이어별 Key-Value 저장소

해당 개념이 필요한 이유

• Autoregressive 모델은 토큰을 하나씩 생성할 때마다 이전 모든 토큰에 대해 Self-Attention을 재계산
• KV Cache 없이는 시퀀스 길이가 길어질수록 연산량이 기하급수적으로 증가
• KV Cache를 사용하면 동일 하드웨어에서 약 5배 빠른 추론 가능 (T4 GPU 벤치마크 기준)

Tensor란?

데이터를 담는 다차원 컨테이너. 차원 수에 따라 이름이 달라진다:

차원	이름	예시
0차원	Scalar	42 — 숫자 하나
1차원	Vector	[1, 2, 3] — 숫자의 나열
2차원	Matrix	행과 열로 구성된 표 (엑셀 시트)
3차원+	Tensor	행렬을 여러 장 쌓은 것

Q) “KV Cache 문서에서 말하는 텐서가 뭔가요?”

• KV Cache에서 Key, Value 텐서는 4차원 텐서: [batch_size, num_heads, seq_len, head_dim]
• 직관적으로 — “배치 × 어텐션 헤드 수 × 시퀀스 길이 × 헤드 차원”의 4겹 숫자 덩어리
• 예: [1, 32, 512, 128] → 배치 1개, 헤드 32개, 토큰 512개, 각 헤드 크기 128

Self-Attention이란?

문장 속 각 단어가 다른 모든 단어와의 관계를 계산하여 문맥을 파악하는 메커니즘.

도서관 비유:

• Query (Q): “나는 이런 정보를 찾고 있어” — 현재 단어가 원하는 것
• Key (K): “나는 이런 정보를 갖고 있어” — 각 단어가 제공할 수 있는 것
• Value (V): “내 실제 내용은 이거야” — 각 단어의 실제 정보

sequenceDiagram
    autonumber
    participant Token as 입력 토큰
    participant QKV as Q, K, V 생성
    participant Score as 유사도 계산
    participant Soft as Softmax
    participant Out as 출력

    Token->>QKV: 각 토큰 → Q, K, V 벡터 변환
    Note over QKV: 입력 × 가중치 행렬 3개<br/>(Wq, Wk, Wv)
    QKV->>Score: Q × Kᵀ (내적)
    Note over Score: "나(Q)와 비슷한<br/>정보(K)를 가진 단어는?"
    Score->>Soft: 점수 → 확률 분포 변환
    Note over Soft: 합이 1이 되도록 정규화
    Soft->>Out: 확률 × V = 가중 평균
    Note over Out: 관련도 높은 단어의<br/>정보(V)를 더 많이 반영

핵심 수식: Attention(Q, K, V) = softmax(Q × Kᵀ / √d) × V

KV Cache와의 연결: Self-Attention에서 Q는 매번 새로 계산하지만, K와 V는 이전 토큰의 것을 재활용할 수 있다. 이것이 바로 KV Cache의 원리.

AS-IS

sequenceDiagram
    autonumber
    participant Input as 입력 시퀀스
    participant Attn as Self-Attention
    participant Out as 출력

    Note over Input: "The red cat was"
    Input->>Attn: Token 1~4 → Q, K, V 전체 계산
    Attn->>Out: "sitting" 생성

    Note over Input: "The red cat was sitting"
    Input->>Attn: Token 1~5 → Q, K, V 전체 재계산
    Note over Attn: Token 1~4의 K,V를<br/>처음부터 다시 계산 (낭비)
    Attn->>Out: "on" 생성

    Note over Input: "The red cat was sitting on"
    Input->>Attn: Token 1~6 → Q, K, V 전체 재계산
    Note over Attn: 매 스텝마다<br/>전체 재계산 반복
    Attn->>Out: "the" 생성

TO-BE

sequenceDiagram
    autonumber
    participant Input as 새 토큰
    participant Cache as KV Cache
    participant Attn as Self-Attention
    participant Out as 출력

    Note over Cache: Prefill Phase
    Input->>Cache: Token 1~4 → K₁₋₄, V₁₋₄ 저장
    Input->>Attn: Q₁₋₄ + K₁₋₄ + V₁₋₄
    Attn->>Out: "sitting" 생성

    Note over Cache: Decode Phase
    Input->>Cache: Token 5 → K₅, V₅ 추가 저장
    Note over Cache: Cache: [K₁..K₅], [V₁..V₅]
    Input->>Attn: Q₅ + Cache(K₁₋₅, V₁₋₅)
    Note over Attn: 새 토큰의 Q만 계산
    Attn->>Out: "on" 생성

    Input->>Cache: Token 6 → K₆, V₆ 추가 저장
    Input->>Attn: Q₆ + Cache(K₁₋₆, V₁₋₆)
    Attn->>Out: "the" 생성

Prefill과 Decode — 두 단계의 추론 (상세)

Prefill Phase (채우기 단계)

사용자가 입력한 프롬프트 전체를 한 번에 병렬 처리하는 단계.

입력: "오늘 날씨 어때?" (5토큰)
→ 5개 토큰을 GPU에서 동시에 처리
→ 5개 토큰 각각의 K, V를 한꺼번에 Cache에 저장
→ 첫 번째 출력 토큰 생성

• 특징: GPU 연산(compute) 집약적 — 행렬곱이 많아 GPU 코어를 최대한 활용
• 병목: 연산량 (FLOPS)
• 비유: 도서관에서 참고 도서 10권을 한꺼번에 책상 위에 펼쳐놓는 것

Decode Phase (생성 단계)

토큰을 하나씩 순차적으로 생성하는 단계.

1단계: Q("맑") × Cache(K₁₋₅, V₁₋₅) → "고" 생성, K₆V₆ 캐시 추가
2단계: Q("고") × Cache(K₁₋₆, V₁₋₆) → "화창" 생성, K₇V₇ 캐시 추가
3단계: Q("창") × Cache(K₁₋₇, V₁₋₇) → "합니다" 생성, K₈V₈ 캐시 추가

• 특징: 메모리 대역폭(memory bandwidth) 의존적 — 캐시에서 K,V를 읽어오는 속도가 병목
• 병목: 메모리 읽기 속도 (GB/s)
• 비유: 펼쳐놓은 참고 도서(Cache)에서 한 줄씩 읽으며 답안을 적는 것

두 단계의 핵심 차이

	Prefill	Decode
처리 방식	병렬 (전체 토큰 동시)	순차 (토큰 1개씩)
병목	GPU 연산량 (compute-bound)	메모리 대역폭 (memory-bound)
KV Cache	캐시 생성	캐시 읽기 + 추가
실행 횟수	1번	출력 토큰 수만큼 반복

메모리 비용

각 Transformer 블록에서 토큰당 저장하는 텐서:

• K 텐서: [batch_size, num_heads, seq_len, head_dim]
• V 텐서: [batch_size, num_heads, seq_len, head_dim]

예시: LLaMA-2 13B 기준 토큰당 ~1MB, 4K 컨텍스트에서 시퀀스당 ~4GB

성능 비교 (T4 GPU 벤치마크)

방식	소요 시간
KV Cache 미사용	1분 1초
KV Cache 사용	11.7초
속도 향상	~5.2배

참고 문서

• KV Caching Explained - HuggingFace Blog
• Transformers Key-Value Caching - Neptune.ai
• Tensor 개념 - GeeksforGeeks
• Self-Attention Q,K,V 설명 - Dive into Deep Learning