Transformers 로컬 LLM KV 캐시로 TPS 2배 튜닝

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-hf"  # 예시

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
)
model.eval()

# 캐시 설정 확인
model.config.use_cache = True

prompt = "Explain KV cache in one paragraph."
inputs = tok(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    out = model(**inputs, use_cache=True)

print("has past_key_values:", out.past_key_values is not None)
if out.past_key_values is not None:
    # 레이어 수 확인
    print("layers:", len(out.past_key_values))

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-hf"

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
)
model.eval()
model.config.use_cache = True

def benchmark(prompt: str, max_new_tokens: int = 128):
    inputs = tok(prompt, return_tensors="pt").to("cuda")

    # 프리필만 측정: forward 1회
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        out = model(**inputs, use_cache=True)
    torch.cuda.synchronize()
    t1 = time.perf_counter()

    # 디코드 측정: generate
    torch.cuda.synchronize()
    t2 = time.perf_counter()
    with torch.no_grad():
        gen = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
        )
    torch.cuda.synchronize()
    t3 = time.perf_counter()

    prefill_s = t1 - t0
    total_s = t3 - t2

    # 생성 토큰 수 계산(프롬프트 길이 제외)
    prompt_len = inputs["input_ids"].shape[1]
    gen_len = gen.shape[1]
    new_tokens = gen_len - prompt_len

    # 전체 generate 시간에는 프리필이 포함되므로, 대략적인 참고치로만 사용
    approx_tps = new_tokens / total_s

    return {
        "prefill_s": prefill_s,
        "generate_s": total_s,
        "new_tokens": new_tokens,
        "approx_tps": approx_tps,
    }

print(benchmark("Write a short story about caching.", 128))

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-hf"

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
)
model.eval()
model.config.use_cache = True

@torch.no_grad()
def decode_with_kv_cache(prompt: str, max_new_tokens: int = 128):
    inputs = tok(prompt, return_tensors="pt").to("cuda")
    input_ids = inputs["input_ids"]

    # 1) prefill
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    out = model(input_ids=input_ids, use_cache=True)
    past = out.past_key_values
    torch.cuda.synchronize()
    t1 = time.perf_counter()

    # 2) decode loop
    generated = []
    next_token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)

    torch.cuda.synchronize()
    t2 = time.perf_counter()
    for _ in range(max_new_tokens):
        out = model(input_ids=next_token, past_key_values=past, use_cache=True)
        past = out.past_key_values
        next_token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
        generated.append(next_token)
    torch.cuda.synchronize()
    t3 = time.perf_counter()

    prefill_s = t1 - t0
    decode_s = t3 - t2
    tps = max_new_tokens / decode_s

    gen_ids = torch.cat([input_ids] + generated, dim=1)
    text = tok.decode(gen_ids[0], skip_special_tokens=True)

    return {
        "prefill_s": prefill_s,
        "decode_s": decode_s,
        "decode_tps": tps,
        "text": text,
    }

res = decode_with_kv_cache("Explain transformers KV cache briefly.", 128)
print({k: v for k, v in res.items() if k != "text"})

model.eval()
with torch.inference_mode():
    out = model(**inputs, use_cache=True)

import torch

torch.backends.cuda.matmul.allow_tf32 = True
# Ampere 이상에서 TF32가 matmul 처리량을 올릴 수 있음

# SDPA 커널 선택은 버전/모델에 따라 효과가 다름
# flash 우선이 항상 정답은 아니니 측정 기반으로 선택

tok.padding_side = "left"
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

from dataclasses import dataclass
import torch

@dataclass
class SessionState:
    past_key_values: object | None
    last_token: torch.Tensor | None

sessions: dict[str, SessionState] = {}

@torch.no_grad()
def chat_step(session_id: str, user_text: str, model, tok, max_new_tokens: int = 64):
    state = sessions.get(session_id) or SessionState(None, None)

    # 새 입력 토큰화
    # 실제 서비스에서는 역할 토큰/템플릿을 포함해야 함
    inp = tok(user_text, return_tensors="pt").to("cuda")
    input_ids = inp["input_ids"]

    if state.past_key_values is None:
        out = model(input_ids=input_ids, use_cache=True)
    else:
        out = model(input_ids=input_ids, past_key_values=state.past_key_values, use_cache=True)

    past = out.past_key_values
    next_token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)

    # 간단 디코드
    gen = []
    for _ in range(max_new_tokens):
        out = model(input_ids=next_token, past_key_values=past, use_cache=True)
        past = out.past_key_values
        next_token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
        gen.append(next_token)

    sessions[session_id] = SessionState(past, next_token)

    gen_ids = torch.cat(gen, dim=1)
    return tok.decode(gen_ids[0], skip_special_tokens=True)