LLM 서빙 최적화 완전 가이드: KV Cache, PagedAttention, 양자화의 모든 것

LLM 추론의 두 단계: Prefill과 Decode
KV Cache: 어텐션의 메모리 딜레마
PagedAttention (vLLM): 가상 메모리가 LLM을 구하다
- 핵심 아이디어: OS 가상 메모리의 LLM 적용
- Prefix Caching: 공통 프롬프트 공유
연속 배칭(Continuous Batching): 처리량 극대화
- 정적 배칭의 문제
- 연속 배칭: 매 스텝마다 동적 스케줄링
양자화(Quantization): 정밀도를 트레이드오프하여 속도와 메모리 절약
추측 디코딩(Speculative Decoding): 무료 점심은 있다
- 아이디어: 작은 모델이 초안을 쓰고 큰 모델이 검증
텐서 병렬성과 파이프라인 병렬성
- Tensor Parallelism: 레이어를 GPU에 분산
- Pipeline Parallelism: 레이어를 GPU에 순서대로 배치
vLLM vs TGI vs TensorRT-LLM: 주요 프레임워크 비교
- 어떤 프레임워크를 선택해야 하는가?
실전: 프로덕션 LLM 서빙 스택

LLM 추론의 두 단계: Prefill과 Decode

LLM이 텍스트를 생성하는 과정은 근본적으로 다른 두 단계로 나뉜다. 이 차이를 이해하지 않으면 최적화가 불가능하다.

Phase 1: PREFILL (입력 처리)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
입력: "프랑스의 수도는 어디인가요?"
       └─ 15개 토큰 동시 처리

동작:
  - 모든 입력 토큰을 병렬로 처리 (행렬 곱셈!)
  - 각 토큰의 Q, K, V 계산
  - KV Cache 생성 (재사용할 K, V 저장)
  - 첫 번째 출력 토큰 생성

특성:
  - GPU 연산: COMPUTE-BOUND (행렬 × 행렬)
  - GPU 활용률: HIGH ✅
  - 대기 시간 결정 요인: TTFT (Time To First Token)

Phase 2: DECODE (토큰 생성)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
생성: "파리" → "는" → "입니다" → "." → ...

동작:
  - 한 번에 토큰 1개씩 순차 생성
  - 새 토큰의 Q 계산 후 캐시된 K, V와 어텐션
  - 모든 모델 가중치를 토큰마다 읽어야 함

특성:
  - GPU 연산: MEMORY-BOUND (행렬 × 벡터)
  - GPU 활용률: LOW (종종 5~20%!)
  - 처리량 결정 요인: TBT (Time Between Tokens)

이것이 LLM 서빙 최적화가 어려운 근본 이유:
두 단계가 서로 다른 병목을 가짐!

실제 수치로 확인해 보자:

import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer

def measure_llm_phases():
    model_name = "meta-llama/Llama-3.2-1B"
    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.float16, device_map="cuda"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    prompt = "Explain quantum computing in simple terms:"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    input_len = inputs["input_ids"].shape[1]

    # Prefill 시간 측정
    torch.cuda.synchronize()
    t_start = time.perf_counter()
    with torch.no_grad():
        outputs = model(**inputs)  # 입력 처리만
    torch.cuda.synchronize()
    t_prefill = time.perf_counter() - t_start

    # Decode 시간 측정 (토큰 하나씩)
    t_start = time.perf_counter()
    with torch.no_grad():
        generated = model.generate(
            inputs["input_ids"],
            max_new_tokens=50,
            do_sample=False
        )
    torch.cuda.synchronize()
    t_total = time.perf_counter() - t_start

    t_decode = t_total - t_prefill
    num_new_tokens = generated.shape[1] - input_len

    print(f"입력 토큰 수: {input_len}")
    print(f"Prefill 시간: {t_prefill*1000:.1f}ms (TTFT)")
    print(f"생성 토큰 수: {num_new_tokens}")
    print(f"Decode 시간: {t_decode*1000:.1f}ms")
    print(f"토큰당 시간: {t_decode/num_new_tokens*1000:.1f}ms/token")
    # Llama-1B on H100 기준:
    # Prefill: ~5ms (입력 길이와 선형)
    # Decode: ~3ms/token (모델 크기와 비례, 배치에 의존)

KV Cache: 어텐션의 메모리 딜레마

KV Cache 없으면 어떻게 되나?

KV Cache 없는 자기회귀 생성:

Step 1: 입력 ["안녕", "하세요"] → "저는" 생성
  - "안녕": Q1,K1,V1 계산
  - "하세요": Q2,K2,V2 계산
  - "저는": Q3,K3,V3 계산
  - 어텐션: Q3 × [K1,K2,K3]^T
  연산량: 3² = 9 번의 내적

Step 2: ["안녕","하세요","저는"] → "AI" 생성
  - 모든 토큰 처음부터 다시 계산!
  - "안녕": Q1,K1,V1 다시 계산 (낭비!)
  - "하세요": Q2,K2,V2 다시 계산 (낭비!)
  - "저는": Q3,K3,V3 다시 계산 (낭비!)
  연산량: 4² = 16 번의 내적

N번째 토큰: O(N²) 연산
전체 생성 (L 토큰): O(L³) 연산
→ 100 토큰 생성: 1백만 번의 내적 연산
→ 1000 토큰 생성: 10억 번의 내적 연산

KV Cache: 이전 계산 결과 재사용

import torch
import torch.nn as nn
import math

class AttentionWithKVCache(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

        # KV Cache 저장소
        self.k_cache = None  # (batch, n_heads, seq_len, d_k)
        self.v_cache = None  # (batch, n_heads, seq_len, d_k)

    def forward(self, x, use_cache=True):
        batch, seq, d = x.shape

        # 새 Q, K, V 계산
        Q = self.W_q(x).view(batch, seq, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq, self.n_heads, self.d_k).transpose(1, 2)

        if use_cache and self.k_cache is not None:
            # 캐시에 새 K, V 추가
            K = torch.cat([self.k_cache, K], dim=2)  # (batch, heads, past+new, d_k)
            V = torch.cat([self.v_cache, V], dim=2)

        if use_cache:
            self.k_cache = K.detach()  # 캐시 업데이트
            self.v_cache = V.detach()

        # Q는 현재 토큰만, K/V는 전체 시퀀스
        scale = math.sqrt(self.d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / scale
        weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(weights, V)

        return output.transpose(1, 2).contiguous().view(batch, seq, d)


# KV Cache 크기 계산
def calculate_kv_cache_size(
    seq_len, n_layers, n_heads, head_dim, batch_size, dtype_bytes=2
):
    """
    KV Cache 메모리 사용량 (bytes)
    2: K와 V 두 개
    """
    size = (2 * n_layers * n_heads * head_dim * seq_len * batch_size * dtype_bytes)
    return size


# Llama 3.1 70B 기준:
size = calculate_kv_cache_size(
    seq_len=4096,
    n_layers=80,
    n_heads=8,           # GQA로 줄어든 KV heads
    head_dim=128,
    batch_size=1,
    dtype_bytes=2        # FP16
)
print(f"KV Cache (Llama 70B, seq=4096, batch=1): {size/1e9:.1f} GB")
# 결과: ~6.7 GB per 요청!
# 배치 32: ~214 GB → H100 80GB에 들어가지 않음!
# → KV Cache 메모리 관리가 핵심 문제

KV Cache 메모리 단편화 문제

전통적인 KV Cache 할당 방식:

요청이 도착하면 max_seq_len만큼 미리 할당:
┌───────────────────────────────────────────────────────┐
│ Request 1: seq_len=100  (max=512 할당)                │
│ ████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░         │
│ [사용: 100] [낭비: 412 슬롯 = 80%!]                   │
├───────────────────────────────────────────────────────┤
│ Request 2: seq_len=50   (max=512 할당)                │
│ ██████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░         │
│ [사용: 50]  [낭비: 462 슬롯 = 90%!]                   │
├───────────────────────────────────────────────────────┤
│ Request 3: seq_len=300  (max=512 할당)                │
│ █████████████████████████████████████░░░░░░░░         │
│ [사용: 300] [낭비: 212 슬롯 = 41%!]                   │
└───────────────────────────────────────────────────────┘

총 할당: 3 × 512 = 1536 슬롯
총 사용: 450 슬롯
낭비: 1086 슬롯 = 70%!

+ 내부 단편화: 할당은 되었으나 사용 안 됨
+ 외부 단편화: 작은 요청들 사이의 빈 공간

PagedAttention (vLLM): 가상 메모리가 LLM을 구하다

핵심 아이디어: OS 가상 메모리의 LLM 적용

Kwon et al. (UC Berkeley, 2023)의 통찰: "OS가 수십 년 전에 메모리 단편화 문제를 해결했잖아. 그 방법을 KV Cache에 적용하자!"

OS 가상 메모리의 교훈:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
가상 주소 → 페이지 테이블 → 물리 주소
프로세스는 연속된 가상 주소를 봄
실제 물리 메모리는 불연속할 수 있음
→ 단편화 없음, 효율적 메모리 사용

PagedAttention 아이디어:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
가상 KV 슬롯 → 블록 테이블 → 물리 블록
시퀀스는 연속된 가상 슬롯을 봄
실제 GPU 메모리의 블록은 불연속할 수 있음
→ 단편화 없음, 효율적 GPU 메모리 사용

PagedAttention 메모리 레이아웃:

물리 메모리를 고정 크기 블록으로 분할 (기본: 16 토큰):
┌─────────────────────────────────────────────────────────┐
│                  물리 KV Cache 블록                     │
│ Block 0: [tok0~15]   Block 1: [tok16~31]               │
│ Block 2: [tok32~47]  Block 3: [tok48~63]               │
│ Block 4: [tok64~79]  Block 5: FREE                     │
│ Block 6: FREE        Block 7: FREE                     │
└─────────────────────────────────────────────────────────┘

블록 테이블 (페이지 테이블과 동일한 역할):
┌──────────┬──────────────────────────────────────────────┐
│ 요청 ID  │ 가상 블록 → 물리 블록 매핑                   │
├──────────┼──────────────────────────────────────────────┤
│ Req 1    │ virt[0]→phys[0], virt[1]→phys[2]            │
│          │ (토큰 0~15: 블록0, 토큰 32~47: 블록2)        │
├──────────┼──────────────────────────────────────────────┤
│ Req 2    │ virt[0]→phys[1], virt[1]→phys[3]            │
│          │ (토큰 0~15: 블록1, 토큰 16~31: 블록3)        │
└──────────┴──────────────────────────────────────────────┘

핵심:
- 블록은 필요할 때 ON-DEMAND로 할당!
- 시퀀스 길이가 늘어날수록 블록 추가
- 단편화 < 1 블록 = 최대 16 토큰 (사실상 없음)
- 블록 공유 가능! (프롬프트 공유 시)

# vLLM을 이용한 PagedAttention 서빙:
from vllm import LLM, SamplingParams
import time

def benchmark_vllm():
    # 모델 로드 (PagedAttention 자동 사용)
    llm = LLM(
        model="meta-llama/Llama-3.2-8B-Instruct",
        gpu_memory_utilization=0.9,    # GPU 메모리의 90% 사용
        max_model_len=8192,            # 최대 시퀀스 길이
        block_size=16,                 # PagedAttention 블록 크기
        max_num_seqs=256,              # 최대 동시 요청 수
    )

    # 다양한 길이의 요청 생성
    prompts = [
        "짧은 질문: 파이썬이란?",
        "중간 길이 질문: " + "딥러닝의 역사에 대해 자세히 설명해주세요. " * 5,
        "긴 질문: " + "트랜스포머 아키텍처를 처음부터 구현하는 방법은? " * 10,
    ] * 20  # 60개 요청

    params = SamplingParams(
        temperature=0.0,    # 그리디 디코딩
        max_tokens=100
    )

    start = time.perf_counter()
    outputs = llm.generate(prompts, params)
    elapsed = time.perf_counter() - start

    total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
    print(f"요청 수: {len(prompts)}")
    print(f"총 생성 토큰: {total_tokens}")
    print(f"처리 시간: {elapsed:.1f}s")
    print(f"처리량: {total_tokens/elapsed:.0f} tokens/s")


# vLLM의 메모리 효율:
# 전통 방식: GPU 메모리의 20-40%만 실제 KV Cache에 사용
# PagedAttention: GPU 메모리의 >95% 실제 KV Cache에 사용
# 결과: 동일 GPU에서 2-3배 더 많은 요청 처리 가능!

Prefix Caching: 공통 프롬프트 공유

PagedAttention의 보너스: 여러 요청이 같은 프롬프트를 공유하면 블록을 재사용할 수 있다!

# vLLM Prefix Caching 설정:
llm = LLM(
    model="meta-llama/Llama-3.2-8B-Instruct",
    enable_prefix_caching=True,   # 공통 프리픽스 KV 캐시 공유
)

# System 프롬프트가 같은 여러 요청:
system_prompt = "당신은 도움이 되는 AI 어시스턴트입니다. " * 50  # 긴 시스템 프롬프트

requests = [
    system_prompt + "질문1: 파이썬이란?",
    system_prompt + "질문2: 자바스크립트란?",
    system_prompt + "질문3: Go 언어란?",
]

# system_prompt 부분의 KV Cache가 3개 요청에서 공유됨!
# system_prompt 처리: 1번만 실행 (3번이 아님)
# Prefill 비용: 3 → 1 (3배 절약!)

연속 배칭(Continuous Batching): 처리량 극대화

정적 배칭의 문제

정적 배칭 (Static Batching):

GPU가 처리하는 배치:
Step 1:  [Req1: 진행중] [Req2: 진행중] [Req3: 진행중]
Step 2:  [Req1: 진행중] [Req2: 완료!] [Req3: 진행중]
Step 3:  [Req1: 진행중] [     대기    ] [Req3: 진행중]  ← GPU 낭비!
Step 4:  [Req1: 완료!] [     대기    ] [Req3: 진행중]  ← GPU 낭비!
Step 5:  [     대기    ] [     대기    ] [Req3: 완료!]  ← GPU 낭비!

새 요청은 배치 전체가 완료될 때까지 대기
→ GPU 낭비율: 50% 이상

연속 배칭: 매 스텝마다 동적 스케줄링

연속 배칭 (Continuous Batching / Iteration-level Scheduling):

GPU가 처리하는 배치 (매 토큰 생성 후 재조정):
Step 1:  [Req1] [Req2] [Req3]
Step 2:  [Req1] [Req2] [Req3]
Step 3:  [Req1] [Req4] [Req3]   ← Req2 완료 즉시 Req4 삽입!
Step 4:  [Req5] [Req4] [Req3]   ← Req1 완료 즉시 Req5 삽입!
Step 5:  [Req5] [Req4] [Req6]   ← Req3 완료 즉시 Req6 삽입!

GPU는 항상 최대 활용!
처리량 향상: 정적 배칭 대비 2-4배

# vLLM의 연속 배칭 내부 구조 (간략화):
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.engine.arg_utils import AsyncEngineArgs
import asyncio

async def serve_with_continuous_batching():
    engine = AsyncLLMEngine.from_engine_args(
        AsyncEngineArgs(
            model="meta-llama/Llama-3.2-8B-Instruct",
            max_num_seqs=256,          # 최대 동시 시퀀스
            max_num_batched_tokens=8192,  # 배치당 최대 토큰
        )
    )

    # 비동기로 여러 요청 동시 처리
    async def generate(prompt, request_id):
        async for output in engine.generate(prompt, request_id=request_id):
            if output.finished:
                return output.outputs[0].text

    # 요청들이 비동기로 처리됨 (연속 배칭 자동 적용)
    results = await asyncio.gather(
        generate("질문1: ...", "req_1"),
        generate("질문2: ...", "req_2"),
        generate("질문3: ...", "req_3"),
    )
    return results

양자화(Quantization): 정밀도를 트레이드오프하여 속도와 메모리 절약

왜 양자화인가?

LLM 메모리 사용량 (FP16 기준):
  Llama 3.1 8B:   16 GB
  Llama 3.1 70B:  140 GB
  Llama 3.1 405B: 810 GB

일반 GPU 메모리:
  RTX 4090:  24 GB  → 8B 모델도 빡빡
  A100 80GB:         → 70B 불가
  H100 80GB:         → 70B 불가 (단일 GPU)
  H100 8×80GB:       → 70B 가능, 405B 불가

양자화로 메모리 절약:
  FP16 (16비트): 기준선
  INT8  (8비트):  50% 절약, ~1% 성능 손실
  INT4  (4비트):  75% 절약, ~2-3% 성능 손실
  INT3  (3비트):  81% 절약, 성능 저하 주의
  INT2  (2비트):  88% 절약, 보통 사용 불가 수준

Post-Training Quantization (PTQ): INT8

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# INT8 양자화 (LLM.int8() - Dettmers et al., 2022):
config_int8 = BitsAndBytesConfig(
    load_in_8bit=True,
    # 선택적: 특정 레이어는 FP16 유지 (정확도 중요한 레이어)
    llm_int8_skip_modules=["lm_head"],
)

model_int8 = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B-Instruct",
    quantization_config=config_int8,
    device_map="auto",    # 자동으로 GPU에 분산
)

# INT8의 핵심 혁신 (LLM.int8()):
# 문제: activation에 이상값(outlier)이 있으면 양자화 품질 저하
# 해결: 이상값이 있는 열은 FP16으로 유지, 나머지만 INT8
# → "혼합 정밀도 분해(Mixed-precision decomposition)"

4비트 양자화: NF4와 GPTQ

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# NF4 양자화 (QLoRA에서 도입, Dettmers et al., 2023):
config_4bit = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,   # 계산은 FP16으로
    bnb_4bit_quant_type="nf4",              # NormalFloat4 형식
    bnb_4bit_use_double_quant=True,         # 양자화 스케일도 양자화!
)

# NF4란?
# 신경망 가중치는 대략 정규분포를 따름
# 정규분포에 최적화된 4비트 코드북
# 각 코드가 정규분포 상의 동일 확률 구간을 커버
# → 일반 INT4보다 1-2% 정확도 향상

# Double quantization:
# 가중치 양자화 스케일 (FP32): 4096 파라미터당 1개
# 이 스케일 자체도 INT8로 양자화
# → 추가 0.5비트 절약


# GPTQ (Frantar et al., 2022):
# 레이어별 최적화 기반 양자화
# Hessian 행렬을 이용해 양자화 오차 최소화
from auto_gptq import AutoGPTQForCausalLM

model_gptq = AutoGPTQForCausalLM.from_quantized(
    "TheBloke/Llama-2-70B-GPTQ",    # 이미 양자화된 모델
    device="cuda:0",
    use_triton=True,                  # Triton 커널 사용 (더 빠름)
)
# GPTQ는 PTQ 중 정확도가 가장 높은 편

AWQ: 활성화를 고려한 가중치 양자화

# AWQ (Lin et al., 2023) 핵심 통찰:
# 모든 가중치가 동등하게 중요하지 않다!
# 1%의 "핵심(salient)" 가중치가 큰 활성화를 생성함
# 이것들을 INT4로 양자화하면 품질이 급격히 저하

# 해결책:
# 1. 활성화 통계를 사용해 핵심 채널 찾기
# 2. 해당 채널의 가중치 스케일 조정 (per-channel scaling)
# 3. 모든 가중치를 INT4로 양자화하되, 스케일로 오차 보상

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "meta-llama/Llama-3.1-8B-Instruct"
quant_path = "llama-3.1-8b-awq"

# 양자화 실행 (calibration 데이터 필요):
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,           # 4비트
    "version": "GEMM"
}

model.quantize(tokenizer, quant_config=quant_config)
model.save_quantized(quant_path)


# AWQ vs GPTQ 비교:
# AWQ:  빠른 추론 (Triton/CUDA 커널 최적화)
#       메모리: FP16의 25%
#       정확도: GPTQ와 유사하거나 약간 낮음
# GPTQ: 더 높은 정확도 (Hessian 기반)
#       추론 속도: AWQ와 유사
#       메모리: AWQ와 동일

양자화 성능 비교표

Llama 3.1 70B 양자화 비교 (H100 단일 GPU):

┌──────────┬──────────┬──────────┬──────────┬──────────┐
│ 형식     │ 메모리   │ 처리량   │ MMLU     │ 실용성   │
├──────────┼──────────┼──────────┼──────────┼──────────┤
│ FP16     │ 140 GB   │ 기준선   │ 80.9%    │ 8×H100   │
│ BF16     │ 140 GB   │ +5%      │ 80.9%    │ 8×H100   │
│ INT8     │ 70 GB    │ +10%     │ 80.2%    │ 2×H100   │
│ GPTQ-4b  │ 36 GB    │ +30%     │ 79.8%    │ 1×H100   │
│ AWQ-4b   │ 36 GB    │ +35%     │ 79.5%    │ 1×H100   │
│ GGUF-Q4  │ 38 GB    │ CPU 가능 │ 79.1%    │ CPU/GPU  │
└──────────┴──────────┴──────────┴──────────┴──────────┘

추측 디코딩(Speculative Decoding): 무료 점심은 있다

아이디어: 작은 모델이 초안을 쓰고 큰 모델이 검증

표준 디코딩:
  70B 모델이 1 토큰 생성 = 1 forward pass = 10ms
  100 토큰 = 1000ms = 1초

추측 디코딩 (Leviathan et al., 2023):
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Step 1: 초안 모델(7B) → K개 토큰 빠르게 생성
         "파리" "는" "프랑스의" "수도" "입니다"
         5개 토큰, 5ms (7B 모델)

Step 2: 대형 모델(70B) → K+1개 토큰 한 번에 검증!
        ["파리" "는" "프랑스의" "수도" "입니다"] → 병렬 처리!
        10ms (70B 모델이지만 1번의 forward pass)

Step 3: 검증 결과:
        "파리" ✅ "는" ✅ "프랑스의" ✅ "수도" ✅ "입니다" ❌
        → 4개 토큰 수락, 1개 토큰 거부

Step 4: 거부된 위치 이후 대형 모델이 재생성
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
결과: 10ms로 4개 토큰 생성 (표준: 40ms)
속도 향상: 2.5-4배 (수락률에 따라 다름)
품질 손실: ZERO (대형 모델이 최종 검증)

# PyTorch를 이용한 추측 디코딩 구현 (개념적):
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def speculative_decode(
    target_model,    # 큰 모델 (70B)
    draft_model,     # 작은 모델 (7B)
    input_ids,
    max_new_tokens=100,
    K=5,             # 초안 토큰 수
    temperature=1.0
):
    """
    추측 디코딩 구현
    K: 초안 모델이 한 번에 생성할 토큰 수
    """
    generated = input_ids.clone()

    while generated.shape[1] < input_ids.shape[1] + max_new_tokens:
        # Phase 1: 초안 모델로 K개 토큰 생성
        draft_ids = []
        draft_probs = []

        draft_input = generated.clone()
        for _ in range(K):
            with torch.no_grad():
                draft_out = draft_model(draft_input)
                logits = draft_out.logits[:, -1, :] / temperature
                probs = torch.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, 1)
                draft_ids.append(next_token)
                draft_probs.append(probs[0, next_token[0, 0]])
                draft_input = torch.cat([draft_input, next_token], dim=1)

        # Phase 2: 대형 모델로 K+1개 위치 동시 검증
        candidate = torch.cat([generated] + draft_ids, dim=1)
        with torch.no_grad():
            target_out = target_model(candidate)
            # 각 위치에서의 확률
            target_logits = target_out.logits[:, len(generated[0])-1:-1, :]
            target_probs = torch.softmax(target_logits / temperature, dim=-1)

        # Phase 3: 수락/거부 결정
        accepted = 0
        for i, (draft_id, draft_prob, tgt_prob_dist) in enumerate(
            zip(draft_ids, draft_probs,
                target_probs[0])
        ):
            token_id = draft_id[0, 0].item()
            tgt_prob = tgt_prob_dist[token_id].item()

            # 수락 확률: min(1, p_target / p_draft)
            accept_prob = min(1.0, tgt_prob / max(draft_prob.item(), 1e-8))
            if torch.rand(1).item() < accept_prob:
                generated = torch.cat([generated, draft_id], dim=1)
                accepted += 1
            else:
                # 거부: 대형 모델의 확률로 재샘플링
                adjusted = torch.clamp(tgt_prob_dist - draft_probs[i], min=0)
                adjusted = adjusted / adjusted.sum()
                next_token = torch.multinomial(adjusted.unsqueeze(0), 1)
                generated = torch.cat([generated, next_token], dim=1)
                break

        if accepted == K:
            # 모두 수락: 대형 모델의 다음 토큰도 추가
            last_logits = target_out.logits[:, -1, :] / temperature
            last_probs = torch.softmax(last_logits, dim=-1)
            next_token = torch.multinomial(last_probs, 1)
            generated = torch.cat([generated, next_token], dim=1)

    return generated

텐서 병렬성과 파이프라인 병렬성

Tensor Parallelism: 레이어를 GPU에 분산

텐서 병렬성 (Tensor Parallelism, Shoeybi et al., 2019):

70B 모델, 8 GPU, 64 Attention Heads:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
GPU 0: Head 0~7   (W_q: [d_model, 8*d_k] 의 1/8)
GPU 1: Head 8~15
GPU 2: Head 16~23
GPU 3: Head 24~31
GPU 4: Head 32~39
GPU 5: Head 40~47
GPU 6: Head 48~55
GPU 7: Head 56~63

각 GPU가 독립적으로 해당 헤드 계산
마지막에 All-Reduce (각 GPU 결과 합산)

통신 비용:
- 어텐션 레이어당 1번의 All-Reduce
- FFN 레이어당 1번의 All-Reduce
- NVLink (H100): 900 GB/s (최신)
- PCIe: 64 GB/s (느림)
→ NVLink 있으면 텐서 병렬성 효율적

# Megatron-LM 스타일 텐서 병렬 어텐션:
import torch
import torch.distributed as dist

def tensor_parallel_attention(x, W_q, W_k, W_v, W_o,
                               n_heads, rank, world_size):
    """
    rank: 현재 GPU의 번호 (0 ~ world_size-1)
    world_size: 총 GPU 수
    """
    local_heads = n_heads // world_size
    local_head_dim = W_q.shape[1] // world_size

    # 각 GPU는 자신의 헤드만 계산
    # W_q는 이미 샤딩되어 있음 (각 GPU에 1/world_size 크기)
    Q_local = x @ W_q    # (batch, seq, local_heads * d_k)
    K_local = x @ W_k
    V_local = x @ W_v

    # 로컬 어텐션 계산
    Q = Q_local.view(*Q_local.shape[:-1], local_heads, local_head_dim)
    K = K_local.view(*K_local.shape[:-1], local_heads, local_head_dim)
    V = V_local.view(*V_local.shape[:-1], local_heads, local_head_dim)

    # 어텐션 계산 (FlashAttention 사용 시 더 빠름)
    attn_out = compute_attention(Q, K, V)

    # 출력 프로젝션 (W_o도 샤딩)
    out_local = attn_out.view(*attn_out.shape[:-2], -1) @ W_o

    # All-Reduce: 모든 GPU의 결과 합산
    dist.all_reduce(out_local, op=dist.ReduceOp.SUM)

    return out_local


def compute_attention(Q, K, V):
    """간단한 스케일드 내적 어텐션"""
    scale = Q.shape[-1] ** -0.5
    scores = torch.matmul(Q, K.transpose(-2, -1)) * scale
    weights = torch.softmax(scores, dim=-1)
    return torch.matmul(weights, V)

Pipeline Parallelism: 레이어를 GPU에 순서대로 배치

파이프라인 병렬성 (Pipeline Parallelism):

70B 모델, 80 레이어, 4 GPU:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
GPU 0: 레이어  0~19 (임베딩 + 첫 20 트랜스포머)
GPU 1: 레이어 20~39
GPU 2: 레이어 40~59
GPU 3: 레이어 60~79 + LM 헤드

순서:
  마이크로 배치 1: GPU0→GPU1→GPU2→GPU3
  마이크로 배치 2: GPU0→GPU1→GPU2→GPU3
  ...

파이프라인 버블 (비효율):
  │GPU0 [mb1]│GPU0 [mb2]│GPU0 [mb3]│GPU0 [mb4]│ GPU0 대기
  │    대기   │GPU1 [mb1]│GPU1 [mb2]│GPU1 [mb3]│ GPU1 대기
  │    대기   │    대기   │GPU2 [mb1]│GPU2 [mb2]│ GPU2 대기
  │    대기   │    대기   │    대기   │GPU3 [mb1]│ GPU3 대기

마이크로 배칭으로 파이프라인 버블 감소:
버블 비율 = (p-1)/(m+p-1)
p: 파이프라인 스테이지 수, m: 마이크로 배치 수
m을 크게 할수록 효율 향상

vLLM vs TGI vs TensorRT-LLM: 주요 프레임워크 비교

LLM 서빙 프레임워크 비교 (2025년 기준):

┌────────────────┬─────────────────────────────────────────────────────┐
│ 프레임워크     │ vLLM                                                │
├────────────────┼─────────────────────────────────────────────────────┤
│ 개발처         │ UC Berkeley / vLLM team                             │
│ 핵심 기술      │ PagedAttention, Continuous Batching                 │
│ 양자화 지원    │ AWQ, GPTQ, INT8, FP8                                │
│ 최대 처리량    │ ★★★★☆ (높음)                                       │
│ TTFT 지연      │ ★★★☆☆ (중간)                                       │
│ 사용 편의성    │ ★★★★★ (매우 쉬움)                                  │
│ 커스터마이즈   │ ★★★☆☆ (중간)                                       │
│ 라이선스       │ Apache 2.0                                          │
│ 특징           │ Python 친화적, 가장 활발한 커뮤니티                 │
└────────────────┴─────────────────────────────────────────────────────┘

┌────────────────┬─────────────────────────────────────────────────────┐
│ 프레임워크     │ TGI (Text Generation Inference)                     │
├────────────────┼─────────────────────────────────────────────────────┤
│ 개발처         │ Hugging Face                                        │
│ 핵심 기술      │ Continuous Batching, Flash Attention                │
│ 양자화 지원    │ GPTQ, AWQ, Bitsandbytes                             │
│ 최대 처리량    │ ★★★☆☆ (중간)                                       │
│ TTFT 지연      │ ★★★☆☆ (중간)                                       │
│ 사용 편의성    │ ★★★★☆ (쉬움)                                       │
│ 커스터마이즈   │ ★★★★☆ (높음)                                       │
│ 라이선스       │ HFOIL (상업용 주의)                                 │
│ 특징           │ HF 생태계 통합, Docker 기반 배포                    │
└────────────────┴─────────────────────────────────────────────────────┘

┌────────────────┬─────────────────────────────────────────────────────┐
│ 프레임워크     │ TensorRT-LLM                                        │
├────────────────┼─────────────────────────────────────────────────────┤
│ 개발처         │ NVIDIA                                              │
│ 핵심 기술      │ TensorRT 최적화, In-flight Batching                 │
│ 양자화 지원    │ INT8, INT4, FP8, SmoothQuant, AWQ                  │
│ 최대 처리량    │ ★★★★★ (가장 높음, NVIDIA GPU 전용)                 │
│ TTFT 지연      │ ★★★★★ (최저)                                       │
│ 사용 편의성    │ ★★☆☆☆ (복잡, C++ 기반)                             │
│ 커스터마이즈   │ ★★☆☆☆ (어려움)                                     │
│ 라이선스       │ Apache 2.0                                          │
│ 특징           │ NVIDIA GPU 최적화, 기업용 프로덕션                  │
└────────────────┴─────────────────────────────────────────────────────┘

┌────────────────┬─────────────────────────────────────────────────────┐
│ 프레임워크     │ llama.cpp / Ollama                                  │
├────────────────┼─────────────────────────────────────────────────────┤
│ 개발처         │ ggerganov / Ollama Inc.                             │
│ 핵심 기술      │ GGUF 양자화, CPU/GPU 하이브리드                     │
│ 양자화 지원    │ Q2~Q8 (GGUF 형식)                                   │
│ 최대 처리량    │ ★★☆☆☆ (낮음, CPU 사용 시)                          │
│ TTFT 지연      │ ★★☆☆☆ (높음)                                       │
│ 사용 편의성    │ ★★★★★ (매우 쉬움)                                  │
│ 커스터마이즈   │ ★★☆☆☆ (제한적)                                     │
│ 라이선스       │ MIT                                                 │
│ 특징           │ CPU 실행 가능, 개인/개발 환경                       │
└────────────────┴─────────────────────────────────────────────────────┘

어떤 프레임워크를 선택해야 하는가?

선택 가이드:

프로덕션 서빙, 높은 처리량이 필요:
  → vLLM (Python 팀) 또는 TensorRT-LLM (NVIDIA GPU 전용)

HuggingFace 생태계와 통합 필요:
  → TGI

개인 실험, 로컬 실행:
  → Ollama (가장 쉬움)

CPU에서 실행해야 함:
  → llama.cpp 또는 Ollama

최고 성능, NVIDIA GPU 전용:
  → TensorRT-LLM + Triton Inference Server

실전: 프로덕션 LLM 서빙 스택

프로덕션 LLM 서빙 아키텍처:

클라이언트
  │
  ▼
로드 밸런서 (Nginx / AWS ALB)
  │
  ▼
API 게이트웨이 (FastAPI / Kong)
  │ 레이트 리미팅, 인증, 로깅
  ▼
라우터 (모델 선택, 요청 큐)
  │
  ├─→ vLLM 서버 1 (8B 모델, 빠른 응답용)
  │
  ├─→ vLLM 서버 2 (70B 모델, 고품질용)
  │
  └─→ vLLM 서버 3 (특수 도메인 모델)
         │
         ▼
      모니터링 (Prometheus + Grafana)
      지표: TTFT, TBT, throughput, GPU 사용률

# 프로덕션 vLLM 서버 설정:
import subprocess

# vLLM OpenAI 호환 서버 시작:
cmd = [
    "python", "-m", "vllm.entrypoints.openai.api_server",
    "--model", "meta-llama/Llama-3.1-8B-Instruct",
    "--tensor-parallel-size", "2",         # 2 GPU 텐서 병렬
    "--gpu-memory-utilization", "0.9",
    "--max-model-len", "8192",
    "--max-num-seqs", "256",               # 최대 동시 요청
    "--max-num-batched-tokens", "8192",    # 배치당 최대 토큰
    "--quantization", "awq",               # AWQ 양자화
    "--enable-prefix-caching",            # 프리픽스 캐싱 활성화
    "--block-size", "16",                  # PagedAttention 블록 크기
    "--port", "8000",
]

# 클라이언트 측:
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")

response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[{"role": "user", "content": "파이썬의 GIL이란?"}],
    temperature=0.0,
    max_tokens=500,
)
print(response.choices[0].message.content)

LLM 서빙 최적화는 하드웨어, 알고리즘, 시스템 소프트웨어가 교차하는 복잡한 분야다. PagedAttention이 가상 메모리 개념을 LLM에 적용했듯, 앞으로도 컴퓨터 과학의 오래된 지혜가 새로운 형태로 재발견될 것이다. 각 기술이 풀고자 하는 근본 문제를 이해하면, 새로운 최적화가 나올 때마다 빠르게 습득하고 적용할 수 있다.