KV Cache 완전 정복 — LLM이 메모리를 먹는 진짜 이유

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

[KV Cache 없이]
토큰 1 생성: K,V 계산 → 4개 토큰
토큰 2 생성: K,V 계산 → 5개 토큰 (전부 재계산)
토큰 3 생성: K,V 계산 → 6개 토큰 (전부 재계산)
→ 총 연산: 4 + 5 + 6 = 15

[KV Cache 사용]
토큰 1 생성: K,V 계산 → 4개 토큰 (캐시에 저장)
토큰 2 생성: 새 토큰 1개만 K,V 계산 + 캐시 재사용
토큰 3 생성: 새 토큰 1개만 K,V 계산 + 캐시 재사용
→ 총 연산: 4 + 1 + 1 = 6

\text{KV Cache} = 2 \times L \times n_{kv} \times d_{head} \times S \times b

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "Qwen/Qwen2.5-7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

# 프롬프트 토큰화
prompt = "Explain the KV cache in transformers."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# use_cache=True로 KV Cache 반환
with torch.no_grad():
    outputs = model(**inputs, use_cache=True)

# KV Cache 구조 확인
past_kv = outputs.past_key_values
print(f"레이어 수: {len(past_kv)}")
print(f"Key shape: {past_kv[0][0].shape}")  # (batch, kv_heads, seq_len, head_dim)
print(f"Value shape: {past_kv[0][1].shape}")

# 전체 KV Cache 메모리 계산
total_bytes = sum(
    k.element_size() * k.nelement() + v.element_size() * v.nelement()
    for k, v in past_kv
)
print(f"KV Cache 메모리: {total_bytes / 1024**2:.1f} MB")

레이어 수: 28
Key shape: torch.Size([1, 4, 8, 128])
Value shape: torch.Size([1, 4, 8, 128])
KV Cache 메모리: 0.4 MB

import torch
import matplotlib.pyplot as plt

def measure_kv_cache_memory(model, tokenizer, seq_lengths):
    """시퀀스 길이별 KV Cache 메모리 측정"""
    results = []

    for length in seq_lengths:
        # 원하는 길이의 입력 생성
        input_ids = torch.randint(
            100, 30000, (1, length), device=model.device
        )

        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()

        mem_before = torch.cuda.memory_allocated()

        with torch.no_grad():
            outputs = model(input_ids, use_cache=True)

        # KV Cache 메모리만 측정
        past_kv = outputs.past_key_values
        kv_bytes = sum(
            k.element_size() * k.nelement() + v.element_size() * v.nelement()
            for k, v in past_kv
        )

        results.append({
            'seq_len': length,
            'kv_cache_mb': kv_bytes / 1024**2,
            'total_gpu_mb': torch.cuda.max_memory_allocated() / 1024**2,
        })

        # 메모리 정리
        del outputs, past_kv
        torch.cuda.empty_cache()

    return results

# 측정
seq_lengths = [512, 1024, 2048, 4096, 8192, 16384]
results = measure_kv_cache_memory(model, tokenizer, seq_lengths)

# 시각화
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

lengths = [r['seq_len'] for r in results]
kv_sizes = [r['kv_cache_mb'] for r in results]
total_sizes = [r['total_gpu_mb'] for r in results]

ax1.plot(lengths, kv_sizes, 'o-', color='#4FC3F7', linewidth=2, markersize=8)
ax1.set_xlabel('Sequence Length')
ax1.set_ylabel('KV Cache (MB)')
ax1.set_title('KV Cache Memory vs Sequence Length')
ax1.grid(True, alpha=0.3)

# KV Cache가 전체 GPU 메모리에서 차지하는 비율
kv_ratio = [kv / total * 100 for kv, total in zip(kv_sizes, total_sizes)]
ax2.bar(range(len(lengths)), kv_ratio, color='#FFB74D', alpha=0.8)
ax2.set_xticks(range(len(lengths)))
ax2.set_xticklabels([f'{l//1024}K' for l in lengths])
ax2.set_ylabel('KV Cache / Total GPU Memory (%)')
ax2.set_title('KV Cache Ratio in Total GPU Usage')
ax2.grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.savefig('kv_cache_profile.png', dpi=150, bbox_inches='tight')
plt.show()

Q Heads: [H0] [H1] [H2] [H3] [H4] [H5] [H6] [H7]
K Heads: [H0] [H1] [H2] [H3] [H4] [H5] [H6] [H7]
V Heads: [H0] [H1] [H2] [H3] [H4] [H5] [H6] [H7]
→ KV Cache: 8 × head_dim × 2 per token per layer

Q Heads: [H0] [H1] [H2] [H3] [H4] [H5] [H6] [H7]
K Heads: [K0      ] [K1      ] [K2      ] [K3      ]
V Heads: [V0      ] [V1      ] [V2      ] [V3      ]
→ KV Cache: 4 × head_dim × 2 per token per layer (절반!)

Q Heads: [H0] [H1] [H2] [H3] [H4] [H5] [H6] [H7]
K Heads: [K0                                      ]
V Heads: [V0                                      ]
→ KV Cache: 1 × head_dim × 2 per token per layer (1/8!)

from transformers import AutoConfig

models = [
    "meta-llama/Llama-3.1-8B-Instruct",
    "Qwen/Qwen2.5-7B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "google/gemma-2-9b-it",
]

print(f"{'Model':<40} {'Layers':>7} {'Q Heads':>8} {'KV Heads':>9} {'Head Dim':>9} {'Type':>5}")
print("-" * 80)

for model_id in models:
    config = AutoConfig.from_pretrained(model_id)
    q_heads = config.num_attention_heads
    kv_heads = getattr(config, 'num_key_value_heads', q_heads)
    head_dim = config.hidden_size // q_heads

    if kv_heads == q_heads:
        attn_type = "MHA"
    elif kv_heads == 1:
        attn_type = "MQA"
    else:
        attn_type = "GQA"

    print(f"{model_id:<40} {config.num_hidden_layers:>7} {q_heads:>8} {kv_heads:>9} {head_dim:>9} {attn_type:>5}")

def calculate_vram_budget(
    model_params_b: float,  # 모델 파라미터 수 (billions)
    num_layers: int,
    kv_heads: int,
    head_dim: int,
    max_seq_len: int,
    batch_size: int = 1,
    model_dtype_bytes: int = 2,  # fp16=2, int4=0.5
    kv_dtype_bytes: int = 2,     # fp16=2, int8=1
):
    """VRAM 사용량 계산 (GB)"""

    # 1. 모델 가중치
    model_memory_gb = model_params_b * 1e9 * model_dtype_bytes / 1024**3

    # 2. KV Cache
    kv_cache_bytes = (
        2  # K + V
        * num_layers
        * kv_heads
        * head_dim
        * max_seq_len
        * batch_size
        * kv_dtype_bytes
    )
    kv_cache_gb = kv_cache_bytes / 1024**3

    # 3. 오버헤드 (activation, CUDA context 등) — 약 10~20%
    overhead_gb = (model_memory_gb + kv_cache_gb) * 0.15

    total_gb = model_memory_gb + kv_cache_gb + overhead_gb

    return {
        'model_weights_gb': round(model_memory_gb, 1),
        'kv_cache_gb': round(kv_cache_gb, 1),
        'overhead_gb': round(overhead_gb, 1),
        'total_gb': round(total_gb, 1),
    }

# 예시: Llama 3 8B (fp16) + 8K 컨텍스트
result = calculate_vram_budget(
    model_params_b=8,
    num_layers=32,
    kv_heads=8,
    head_dim=128,
    max_seq_len=8192,
)
print("Llama 3 8B @ 8K context (fp16):")
for k, v in result.items():
    print(f"  {k}: {v} GB")

# 예시: Llama 3 8B (int4 quantized) + 32K 컨텍스트
result_q = calculate_vram_budget(
    model_params_b=8,
    num_layers=32,
    kv_heads=8,
    head_dim=128,
    max_seq_len=32768,
    model_dtype_bytes=0.5,  # int4
    kv_dtype_bytes=2,       # KV는 여전히 fp16
)
print("\nLlama 3 8B (int4) @ 32K context:")
for k, v in result_q.items():
    print(f"  {k}: {v} GB")

Llama 3 8B @ 8K context (fp16):
  model_weights_gb: 14.9
  kv_cache_gb: 1.0
  overhead_gb: 2.4
  total_gb: 18.3

Llama 3 8B (int4) @ 32K context:
  model_weights_gb: 3.7
  kv_cache_gb: 4.0
  overhead_gb: 1.2
  total_gb: 8.9