GPU 메모리 관리 & LLM 추론 최적화: vLLM, PagedAttention, GPTQ, TensorRT-LLM까지

Arithmetic Intensity (AI) = FLOP 수 / 메모리 접근량 (bytes)

성능 상한 = min(Peak FLOPS, Peak Memory BW × AI)

def calc_model_memory_gb(
    num_params: int,       # 파라미터 수 (예: 70e9)
    dtype_bytes: int = 2,  # FP16=2, FP32=4, INT8=1, INT4=0.5
) -> float:
    """모델 가중치 메모리 계산"""
    return (num_params * dtype_bytes) / (1024 ** 3)

# 주요 모델 메모리 (FP16 기준)
models = {
    "Llama-3.1-8B":   {"params": 8e9,   "bytes": 2},
    "Llama-3.1-70B":  {"params": 70e9,  "bytes": 2},
    "Llama-3.1-405B": {"params": 405e9, "bytes": 2},
    "Mistral-7B":     {"params": 7e9,   "bytes": 2},
    "Qwen2-72B":      {"params": 72e9,  "bytes": 2},
}

for name, cfg in models.items():
    mem_gb = calc_model_memory_gb(cfg["params"], cfg["bytes"])
    print(f"{name}: {mem_gb:.1f} GB")

def calc_kv_cache_memory_gb(
    num_layers: int,
    num_heads: int,
    head_dim: int,
    seq_len: int,
    batch_size: int,
    dtype_bytes: int = 2,  # FP16
) -> float:
    """
    KV 캐시 메모리 계산
    각 레이어: 2 (K, V) × num_heads × head_dim × seq_len × batch_size
    """
    kv_per_layer = 2 * num_heads * head_dim * seq_len * batch_size
    total_bytes = kv_per_layer * num_layers * dtype_bytes
    return total_bytes / (1024 ** 3)

# Llama-3.1-70B 예시
# layers=80, heads=64 (GQA: kv_heads=8), head_dim=128
kv_mem = calc_kv_cache_memory_gb(
    num_layers=80,
    num_heads=8,      # GQA의 경우 kv_heads 사용
    head_dim=128,
    seq_len=4096,
    batch_size=1,
    dtype_bytes=2,
)
print(f"KV 캐시 (seq=4096, bs=1): {kv_mem:.2f} GB")
# 출력: KV 캐시 (seq=4096, bs=1): 0.50 GB

# 배치 크기별 KV 캐시
for bs in [1, 4, 8, 16, 32]:
    mem = calc_kv_cache_memory_gb(80, 8, 128, 4096, bs, 2)
    print(f"  batch_size={bs:2d}: {mem:.2f} GB")

def calc_activation_memory_gb(
    hidden_size: int,
    seq_len: int,
    batch_size: int,
    num_layers: int,
    dtype_bytes: int = 2,
) -> float:
    """추론 시 활성화 메모리 근사 계산"""
    # 각 레이어: attention + FFN 활성화
    # 근사치: 2 × hidden_size × seq_len × batch_size per layer
    bytes_per_layer = 2 * hidden_size * seq_len * batch_size * dtype_bytes
    return (bytes_per_layer * num_layers) / (1024 ** 3)

OS 가상 메모리 → PagedAttention
────────────────────────────────
가상 페이지     → 논리 블록 (logical block)
물리 프레임     → 물리 블록 (physical block)
페이지 테이블   → 블록 테이블 (block table)
페이지 폴트     → 블록 할당

요청 A: [Block 0] → [Block 1] → [Block 2]
                                      ↕ 물리 블록 공유 (공통 프롬프트)
요청 B: [Block 0] → [Block 1] → [Block 3]

# vLLM 설치
pip install vllm

# 서버 시작 (단일 GPU)
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --dtype bfloat16 \
    --max-model-len 8192 \
    --gpu-memory-utilization 0.90

# OpenAI 호환 API 호출
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "messages": [{"role": "user", "content": "GPU 메모리 최적화를 설명해줘"}],
    "max_tokens": 512,
    "temperature": 0.7
  }'

# Python 클라이언트로 vLLM API 호출
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": "system", "content": "당신은 GPU 최적화 전문가입니다."},
        {"role": "user", "content": "PagedAttention의 작동 원리를 설명해주세요."},
    ],
    max_tokens=1024,
    stream=True,
)

for chunk in response:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# 4-bit NF4 양자화 설정
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,  # 계산은 BF16으로
    bnb_4bit_quant_type="nf4",              # NormalFloat4 양자화
    bnb_4bit_use_double_quant=True,         # 이중 양자화로 추가 압축
)

model_id = "meta-llama/Llama-3.1-70B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",      # 자동 멀티-GPU 분산
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# 메모리 사용량 확인
print(f"GPU 메모리: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer
import torch

model_id = "meta-llama/Llama-3.1-8B-Instruct"

quantize_config = BaseQuantizeConfig(
    bits=4,              # 4-bit 양자화
    group_size=128,      # 그룹 크기 (작을수록 정확하지만 메모리 증가)
    desc_act=False,      # 활성화 순서 기술
    damp_percent=0.01,   # Hessian 댐핑 계수
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# 캘리브레이션 데이터 준비 (대표 텍스트 샘플)
calibration_data = [
    tokenizer("The GPU accelerates machine learning by...", return_tensors="pt").input_ids,
    tokenizer("Quantization reduces model size while...", return_tensors="pt").input_ids,
    # 실제로는 1024개 이상의 샘플 사용 권장
]

# 모델 로드 및 양자화
model = AutoGPTQForCausalLM.from_pretrained(
    model_id,
    quantize_config=quantize_config,
    torch_dtype=torch.float16,
)

model.quantize(calibration_data)
model.save_quantized("llama-3.1-8b-gptq-4bit")
print("GPTQ 양자화 완료!")

# 양자화된 모델 로드
quantized_model = AutoGPTQForCausalLM.from_quantized(
    "llama-3.1-8b-gptq-4bit",
    use_safetensors=True,
    device="cuda:0",
)

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

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

# AWQ 양자화 설정
quant_config = {
    "zero_point": True,   # 제로 포인트 양자화
    "q_group_size": 128,  # 그룹 크기
    "w_bit": 4,           # 4-bit
    "version": "GEMM",    # 행렬 곱셈 커널
}

model = AutoAWQForCausalLM.from_pretrained(
    model_id,
    low_cpu_mem_usage=True,
    use_cache=False,
)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

model.quantize(tokenizer, quant_config=quant_config)
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
print("AWQ 양자화 완료!")

Static Batching 예시 (batch_size=3):

시간 →
[요청 A: ████████████░░░░░░░░]  (토큰 12개 생성)
[요청 B: ████░░░░░░░░░░░░░░░░]  (토큰 4개 생성)
[요청 C: ████████░░░░░░░░░░░░]  (토큰 8개 생성)
                └─ B, C가 끝나도 A를 기다려야 함 (GPU 낭비)

Continuous Batching:

Step 1: [A1][B1][C1]  ← 3개 동시 처리
Step 2: [A2][B2][C2]
Step 3: [A3][B3][C3]  ← B 완료, 새 요청 D 추가
Step 4: [A4][C4][D1]  ← 즉시 빈 슬롯 채움
Step 5: [A5][C5][D2]
Step 6: [A6][C6][D3]  ← C 완료, 새 요청 E 추가
...

from vllm import LLM, SamplingParams

# 텐서 병렬로 4 GPU에 분산
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,       # GPU 4개에 텐서 병렬 분산
    dtype="bfloat16",
    max_model_len=8192,
    gpu_memory_utilization=0.90,
    enforce_eager=False,           # CUDA 그래프 최적화 사용
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=512,
    stop=["</s>", "[INST]"],
)

prompts = [
    "GPU 메모리 계층을 설명해줘",
    "PagedAttention의 장점은?",
    "양자화 기법을 비교해줘",
]

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated = output.outputs[0].text
    print(f"프롬프트: {prompt[:50]}...")
    print(f"생성: {generated[:100]}...")
    print()

Attention 헤드 분산 (4-way tensor parallel):

GPU 0: Head 0~15
GPU 1: Head 16~31
GPU 2: Head 32~47
GPU 3: Head 48~63

각 GPU가 독립적으로 계산 후 AllReduce로 결과 합산

Llama-3.1-70B (80 레이어) → 4-way pipeline:

GPU 0: Layer 0~19
GPU 1: Layer 20~39
GPU 2: Layer 40~59
GPU 3: Layer 60~79

레이어 순서대로 처리, GPU 간 activation 전달

import torch

def profile_gpu_memory(func, *args, **kwargs):
    """GPU 메모리 사용량 프로파일링"""
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()

    before = torch.cuda.memory_allocated()
    result = func(*args, **kwargs)
    torch.cuda.synchronize()

    after = torch.cuda.memory_allocated()
    peak = torch.cuda.max_memory_allocated()

    print(f"메모리 증가: {(after - before) / 1e9:.3f} GB")
    print(f"피크 메모리: {peak / 1e9:.3f} GB")
    print()
    print(torch.cuda.memory_summary())
    return result

# 메모리 통계 출력 예시
def load_and_infer():
    from transformers import pipeline
    pipe = pipeline(
        "text-generation",
        model="microsoft/phi-2",
        torch_dtype=torch.float16,
        device_map="auto",
    )
    return pipe("GPU memory management is", max_new_tokens=50)

profile_gpu_memory(load_and_infer)

# 최적화된 vLLM 서버 설정 예시
vllm_config = {
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "dtype": "bfloat16",
    "tensor_parallel_size": 1,
    "gpu_memory_utilization": 0.90,   # 90% GPU 메모리 사용
    "max_model_len": 8192,
    "max_num_batched_tokens": 8192,   # 배치당 최대 토큰 수
    "max_num_seqs": 256,              # 동시 처리 시퀀스 수
    "enable_chunked_prefill": True,   # Chunked prefill 활성화
    "block_size": 16,                 # KV 캐시 블록 크기 (PagedAttention)
    "swap_space": 4,                  # CPU swap 공간 (GB)
    "enforce_eager": False,           # CUDA 그래프 사용
    "disable_log_stats": False,
}