KV Cache

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

# 모델 로드
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

prompt = "인공지능의 미래에 대해 설명해주세요:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# KV Cache 비활성화 (비교용 - 매우 느림)
print("=== KV Cache 비활성화 ===")
start = time.time()
with torch.no_grad():
    output_no_cache = model.generate(
        **inputs,
        max_new_tokens=100,
        use_cache=False,  # KV Cache 비활성화
        do_sample=False
    )
time_no_cache = time.time() - start
print(f"생성 시간: {time_no_cache:.2f}초")

# KV Cache 활성화 (기본값 - 빠름)
print("\n=== KV Cache 활성화 ===")
start = time.time()
with torch.no_grad():
    output_with_cache = model.generate(
        **inputs,
        max_new_tokens=100,
        use_cache=True,   # KV Cache 활성화 (기본값)
        do_sample=False
    )
time_with_cache = time.time() - start
print(f"생성 시간: {time_with_cache:.2f}초")

speedup = time_no_cache / time_with_cache
print(f"\n속도 향상: {speedup:.1f}배")

# KV Cache 직접 관리 (streaming 용도)
past_key_values = None
generated_tokens = []

for _ in range(50):
    with torch.no_grad():
        outputs = model(
            **inputs if past_key_values is None else {"input_ids": next_token_id},
            past_key_values=past_key_values,
            use_cache=True
        )
    past_key_values = outputs.past_key_values  # KV Cache 저장
    next_token_id = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    generated_tokens.append(next_token_id.item())

    # KV Cache 크기 확인 (첫 레이어 기준)
    if len(generated_tokens) == 1:
        kv_shape = past_key_values[0][0].shape
        print(f"KV Cache shape: {kv_shape} (layer, batch, heads, seq, dim)")