2장: KV 캐시 메커니즘 심층 분석

이 장에서 배우는 것

• Attention 연산에서 KV 캐시가 필요한 이유
• KV 캐시 메모리 사용량 계산법
• 60-80% 메모리 낭비가 발생하는 원인
• Multi-Query Attention(MQA)과 Grouped-Query Attention(GQA)
• KV 캐시 압축 기법: 양자화와 프루닝

Attention과 KV 캐시의 관계

1장에서 Decode 단계가 토큰을 하나씩 순차적으로 생성한다고 설명했습니다. 핵심 질문은 이것입니다: 새 토큰을 생성할 때, 이전 토큰들의 정보를 어떻게 참조하는가?

Self-Attention 연산에서 각 토큰은 세 가지 벡터로 변환됩니다.

**Query(Q)**는 "무엇을 찾고 있는가", **Key(K)**는 "내가 무엇을 제공하는가", **Value(V)**는 "실제 제공할 정보"를 나타냅니다. Attention 스코어는 Q와 K의 내적으로 계산되고, 이 스코어를 가중치로 V를 합산합니다.

Decode 단계에서 t번째 토큰을 생성할 때, 이 토큰의 Q 벡터는 이전 1번째부터 t-1번째까지 모든 토큰의 K, V 벡터와 연산해야 합니다. 만약 매번 이전 토큰들의 K, V를 처음부터 다시 계산한다면, 토큰을 생성할 때마다 중복 연산이 기하급수적으로 늘어납니다.

KV 캐시의 동작 과정

Prefill 단계에서 입력 프롬프트의 모든 토큰에 대한 K, V 텐서가 한 번에 계산되어 캐시에 저장됩니다. 이후 Decode 단계에서는 새 토큰 하나의 K, V만 추가로 계산해 캐시에 덧붙입니다.

Prefill:  [K1, K2, K3, K4, K5]  ← 입력 5토큰의 K 벡터 일괄 계산
Decode 1: [K1, K2, K3, K4, K5, K6]  ← K6 추가
Decode 2: [K1, K2, K3, K4, K5, K6, K7]  ← K7 추가
...
Decode N: [K1, K2, ..., K_{5+N}]  ← 계속 성장

각 트랜스포머 레이어마다 독립적인 KV 캐시가 유지됩니다. 레이어가 80개인 70B 모델이라면, 80개의 KV 캐시 쌍이 메모리에 존재하게 됩니다.

메모리 사용량 계산

KV 캐시의 메모리 사용량을 구체적으로 계산해 봅시다. 계산에 필요한 파라미터는 다음과 같습니다.

KV 캐시 메모리 공식은 다음과 같습니다.

2를 곱하는 이유는 Key와 Value 두 가지를 모두 저장하기 때문입니다.

def calc_kv_cache_gb(
    layers: int,
    kv_heads: int,
    head_dim: int,
    seq_len: int,
    bytes_per_element: int = 2,  # FP16
) -> float:
    """KV 캐시 메모리 사용량 (GB) 계산"""
    total_bytes = 2 * layers * kv_heads * head_dim * seq_len * bytes_per_element
    return total_bytes / (1024 ** 3)
 
# Llama 2 70B (GQA, 8 KV heads)
cache_70b = calc_kv_cache_gb(
    layers=80, kv_heads=8, head_dim=128, seq_len=8192
)
print(f"Llama 2 70B (8K): {cache_70b:.2f} GB")
# 출력: Llama 2 70B (8K): 2.50 GB
 
# 만약 MHA (64 KV heads)였다면
cache_70b_mha = calc_kv_cache_gb(
    layers=80, kv_heads=64, head_dim=128, seq_len=8192
)
print(f"Llama 2 70B MHA (8K): {cache_70b_mha:.2f} GB")
# 출력: Llama 2 70B MHA (8K): 20.00 GB

GQA를 사용하는 Llama 2 70B는 KV 헤드를 8개로 줄여 약 2.5GB로 캐시 크기를 감소시켰지만, 컨텍스트 길이가 늘어나면 상황은 다시 악화됩니다.

60-80% 메모리 낭비 문제

KV 캐시 크기 자체도 문제이지만, 더 심각한 것은 **메모리 단편화(Fragmentation)**에 의한 낭비입니다.

기존 추론 시스템은 요청이 들어오면 최대 시퀀스 길이만큼의 KV 캐시 메모리를 연속된 블록으로 미리 할당합니다. 예를 들어 최대 길이가 8K인 시스템에서 실제로 2K 토큰만 생성하는 요청이 들어오면, 나머지 6K 토큰에 해당하는 메모리가 통째로 낭비됩니다.

이 문제가 발생하는 근본 원인은 세 가지입니다.

1. 예약 낭비(Reservation Waste): 최대 길이를 미리 예약하지만, 실제 생성 길이는 그보다 짧은 경우가 대부분입니다
2. 내부 단편화(Internal Fragmentation): 연속 메모리 할당 시 블록 내부에 사용되지 않는 공간이 생깁니다
3. 외부 단편화(External Fragmentation): 요청의 할당과 해제가 반복되면서 사용 가능하지만 연속되지 않은 작은 빈 공간들이 흩어집니다

실제 측정에서 기존 시스템은 할당된 KV 캐시 메모리의 60-80%가 비효율적으로 낭비되는 것으로 나타났습니다. 이것은 3장에서 다룰 PagedAttention이 해결하는 핵심 문제입니다.

Multi-Query Attention과 Grouped-Query Attention

KV 캐시의 근본적인 크기를 줄이기 위한 아키텍처 수준의 접근법이 있습니다.

Multi-Head Attention (MHA)

기존 트랜스포머의 표준 구조입니다. 모든 Attention 헤드가 독립적인 K, V 프로젝션을 가집니다. 예를 들어 64개의 헤드가 있다면, K와 V도 각각 64개씩 존재합니다.

Multi-Query Attention (MQA)

**MQA(Multi-Query Attention)**는 모든 Query 헤드가 하나의 Key, Value 쌍을 공유합니다. KV 캐시 크기가 헤드 수 분의 1로 줄어들어, 64헤드 모델에서 64배의 캐시 절감 효과를 얻습니다.

MHA (64 heads):
  Q heads: 64개  |  K heads: 64개  |  V heads: 64개
  
MQA (64 heads):
  Q heads: 64개  |  K heads: 1개   |  V heads: 1개

하지만 MQA는 모델 품질이 다소 저하될 수 있다는 단점이 있습니다. 모든 Query 헤드가 동일한 KV를 참조하므로, 다양한 패턴의 Attention을 학습하기 어려워집니다.

Grouped-Query Attention (GQA)

**GQA(Grouped-Query Attention)**는 MHA와 MQA의 절충안입니다. Query 헤드를 그룹으로 묶고, 각 그룹이 하나의 KV 쌍을 공유합니다.

그룹 1: Q1~Q8   → K1, V1
그룹 2: Q9~Q16  → K2, V2
...
그룹 8: Q57~Q64 → K8, V8

Llama 2 70B, Llama 3 시리즈, Mistral 등 최신 모델들이 GQA를 채택하고 있으며, 이는 캐시 효율과 모델 품질 사이의 가장 실용적인 균형점으로 자리 잡았습니다.

KV 캐시 압축 기법

아키텍처 변경 없이 KV 캐시를 사후적으로 줄이는 압축 기법들도 있습니다.

KV 캐시 양자화

저장된 KV 캐시의 정밀도를 낮추는 방법입니다. FP16으로 계산된 KV 텐서를 INT8이나 FP8, 심지어 INT4로 양자화하면 메모리 사용량을 2-4배 줄일 수 있습니다.

FP16 KV 캐시:  2.50 GB (Llama 2 70B, 8K)
FP8 KV 캐시:   1.25 GB  (50% 절감)
INT4 KV 캐시:  0.625 GB (75% 절감)

NVIDIA의 B200(Blackwell) GPU는 NVFP4 형식의 KV 캐시 압축을 하드웨어 수준에서 지원합니다. 이에 대해서는 7장에서 더 자세히 다루겠습니다.

KV 캐시 프루닝 (Eviction)

모든 토큰의 KV를 유지할 필요가 있는가라는 질문에서 출발합니다. Attention 스코어가 낮은, 즉 거의 참조되지 않는 토큰의 KV를 삭제(Evict)하는 방법입니다.

대표적인 접근법들은 다음과 같습니다.

• H2O (Heavy-Hitter Oracle): Attention 스코어 누적값이 높은 "중요 토큰"만 유지
• Sliding Window + Sink Token: 최근 N개 토큰과 처음 몇 개의 "sink" 토큰만 유지
• Pyramid KV: 하위 레이어는 넓게, 상위 레이어는 좁게 캐시를 유지

KV 캐시가 시스템 설계에 미치는 영향

KV 캐시 관리는 단순한 메모리 문제를 넘어, 추론 시스템 전체의 설계를 좌우합니다.

• 동시 처리 요청 수: 가용 GPU 메모리에서 모델 가중치를 뺀 나머지가 KV 캐시 예산이며, 이것이 동시 처리 가능한 요청 수를 결정합니다
• 최대 컨텍스트 길이: KV 캐시 크기가 컨텍스트 길이에 선형 비례하므로, 긴 컨텍스트 지원에 직접적인 제약이 됩니다
• 배칭 전략: KV 캐시를 효율적으로 공유하고 관리하는 것이 4장의 Continuous Batching과 3장의 PagedAttention의 핵심입니다

정리

이 장에서는 KV 캐시의 메커니즘과 메모리 영향을 심층 분석했습니다.

• KV 캐시는 Decode 단계의 중복 연산을 제거하지만, 메모리 사용량이 컨텍스트 길이에 선형 비례합니다
• MHA 기반 70B 모델은 요청당 약 20GB의 KV 캐시가 필요하며, GQA로 이를 8분의 1 수준으로 줄일 수 있습니다
• 기존 시스템의 사전 할당 방식은 60-80%의 메모리 낭비를 유발합니다
• 양자화, 프루닝 등의 사후 압축 기법으로 추가적인 절감이 가능합니다

다음 장에서는 이 메모리 낭비 문제를 근본적으로 해결한 PagedAttention과 이를 구현한 vLLM을 다룹니다. OS의 가상 메모리 관리에서 영감을 받은 이 기법이 어떻게 메모리 낭비를 4% 미만으로 줄이고 처리량을 2-4배 향상시키는지 알아보겠습니다.