LLM推論最適化完全ガイド: vLLM、TensorRT-LLM、Speculative Decoding

# LLM 추론 파이프라인의 두 단계를 개념적으로 표현
import torch
import time

def llm_inference_pipeline(model, tokenizer, prompt, max_new_tokens=128):
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)

    # 1단계: Prefill - 입력 프롬프트 전체를 병렬 처리
    prefill_start = time.time()
    with torch.no_grad():
        outputs = model(input_ids, use_cache=True)
        past_key_values = outputs.past_key_values  # KV Cache 생성
        next_token_logits = outputs.logits[:, -1, :]
    prefill_time = time.time() - prefill_start

    # 2단계: Decode - 토큰을 하나씩 자기회귀적으로 생성
    decode_start = time.time()
    generated_tokens = []
    for step in range(max_new_tokens):
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        generated_tokens.append(next_token.item())

        with torch.no_grad():
            outputs = model(
                next_token,
                past_key_values=past_key_values,  # 캐시 재사용
                use_cache=True,
            )
            past_key_values = outputs.past_key_values
            next_token_logits = outputs.logits[:, -1, :]

        if next_token.item() == tokenizer.eos_token_id:
            break

    decode_time = time.time() - decode_start
    tokens_per_sec = len(generated_tokens) / decode_time

    print(f"Prefill 시간: {prefill_time:.3f}s (입력 {input_ids.shape[1]} 토큰)")
    print(f"Decode 시간: {decode_time:.3f}s ({len(generated_tokens)} 토큰 생성)")
    print(f"Decode 속도: {tokens_per_sec:.1f} tokens/s")

    return tokenizer.decode(generated_tokens)

# FlashAttention과 기본 어텐션의 메모리 사용량 비교
import torch
from flash_attn import flash_attn_func

# 설정: batch=4, heads=32, seq_len=4096, head_dim=128
batch_size, num_heads, seq_len, head_dim = 4, 32, 4096, 128

q = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device="cuda")

# 기본 어텐션: O(N^2) 메모리 필요 (어텐션 행렬 전체 저장)
# 4096 * 4096 * 32 * 4 * 2 bytes = ~4 GB

# FlashAttention: O(N) 메모리만 필요 (타일링으로 분할 처리)
output = flash_attn_func(q, k, v, causal=True)
# 메모리 사용량이 시퀀스 길이에 선형적으로 증가
print(f"FlashAttention 출력 shape: {output.shape}")

# vLLM 설치 (CUDA 12.1 이상 필요)
pip install vllm

# OpenAI 호환 API 서버 실행
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --gpu-memory-utilization 0.90 \
    --max-model-len 8192 \
    --enable-prefix-caching \
    --dtype auto \
    --port 8000

# 요청 테스트
curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "meta-llama/Llama-3.1-70B-Instruct",
        "prompt": "LLM 추론 최적화의 핵심은",
        "max_tokens": 256,
        "temperature": 0.7
    }'

from vllm import LLM, SamplingParams

# 모델 로드 (자동으로 PagedAttention 적용)
llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    tensor_parallel_size=1,
    gpu_memory_utilization=0.90,
    max_model_len=4096,
    enable_prefix_caching=True,
    quantization="awq",            # AWQ 양자화 모델 사용 시
    # enforce_eager=True,           # CUDA Graph 비활성화 (디버깅용)
)

# 샘플링 파라미터 설정
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=512,
    repetition_penalty=1.1,
)

# 배치 추론 - 여러 프롬프트를 동시에 처리
prompts = [
    "Kubernetes에서 GPU 노드 스케줄링을 최적화하는 방법을 설명해주세요.",
    "Python에서 비동기 프로그래밍의 장단점은 무엇인가요?",
    "마이크로서비스 아키텍처의 장단점을 비교해주세요.",
]

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated = output.outputs[0].text
    tokens_count = len(output.outputs[0].token_ids)
    print(f"프롬프트: {prompt[:50]}...")
    print(f"생성 토큰 수: {tokens_count}")
    print(f"응답: {generated[:200]}...\n")

# vllm-deployment.yaml - Kubernetes 배포 예시
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-llama-70b
  labels:
    app: vllm-server
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-server
  template:
    metadata:
      labels:
        app: vllm-server
    spec:
      containers:
        - name: vllm
          image: vllm/vllm-openai:v0.17.1
          command:
            - python
            - -m
            - vllm.entrypoints.openai.api_server
          args:
            - --model
            - meta-llama/Llama-3.1-70B-Instruct
            - --tensor-parallel-size
            - '4'
            - --gpu-memory-utilization
            - '0.90'
            - --max-model-len
            - '8192'
            - --enable-prefix-caching
            - --max-num-seqs
            - '256'
          ports:
            - containerPort: 8000
          resources:
            limits:
              nvidia.com/gpu: '4'
            requests:
              nvidia.com/gpu: '4'
              memory: '64Gi'
              cpu: '16'
          readinessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 120
            periodSeconds: 10
          livenessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 180
            periodSeconds: 30
      nodeSelector:
        nvidia.com/gpu.product: A100-SXM4-80GB
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    app: vllm-server
  ports:
    - port: 80
      targetPort: 8000
  type: ClusterIP

# TensorRT-LLM 설치 (Docker 권장)
docker pull nvcr.io/nvidia/tensorrt-llm:latest

# 또는 pip 설치
pip install tensorrt-llm

# Llama 모델 체크포인트 변환 및 엔진 빌드
# 1단계: HuggingFace 모델을 TensorRT-LLM 형식으로 변환
python convert_checkpoint.py \
    --model_dir /models/Llama-3.1-70B-Instruct \
    --output_dir /engines/llama-70b-ckpt \
    --dtype float16 \
    --tp_size 4

# 2단계: TensorRT 엔진 빌드
trtllm-build \
    --checkpoint_dir /engines/llama-70b-ckpt \
    --output_dir /engines/llama-70b-engine \
    --gemm_plugin float16 \
    --max_batch_size 64 \
    --max_input_len 4096 \
    --max_seq_len 8192 \
    --paged_kv_cache enable \
    --use_paged_context_fmha enable \
    --multiple_profiles enable

# 3단계: Triton Inference Server로 서빙
docker run --gpus all -p 8000:8000 \
    -v /engines:/engines \
    nvcr.io/nvidia/tritonserver:latest \
    tritonserver --model-repository=/engines/model_repo

import tensorrt_llm
from tensorrt_llm import LLM, SamplingParams, KvCacheConfig

# KV Cache 설정
kv_cache_config = KvCacheConfig(
    free_gpu_memory_fraction=0.85,
    enable_block_reuse=True,
)

# 모델 로드 (HuggingFace 모델에서 직접 빌드 가능)
llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    tensor_parallel_size=1,
    kv_cache_config=kv_cache_config,
)

# 추론 실행
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=512,
)

prompts = [
    "LLM 추론 최적화에서 가장 중요한 요소를 설명해주세요.",
    "TensorRT-LLM의 장점과 한계를 비교해주세요.",
]

outputs = llm.generate(prompts, sampling_params=sampling_params)

for output in outputs:
    print(f"생성 결과: {output.outputs[0].text[:200]}")

# vLLM에서 Speculative Decoding 활성화
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --speculative-model meta-llama/Llama-3.1-8B-Instruct \
    --num-speculative-tokens 5 \
    --speculative-max-model-len 4096 \
    --use-v2-block-manager \
    --port 8000

# Speculative Decoding 효과 측정 스크립트
import time
import requests
import json
import statistics

API_URL = "http://localhost:8000/v1/completions"

def measure_latency(prompt, max_tokens=256, n_requests=10):
    """여러 요청의 지연 시간과 처리량 측정"""
    latencies = []
    total_tokens = []

    for i in range(n_requests):
        payload = {
            "model": "meta-llama/Llama-3.1-70B-Instruct",
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": 0.7,
        }

        start = time.time()
        response = requests.post(API_URL, json=payload)
        elapsed = time.time() - start

        result = response.json()
        completion_tokens = result["usage"]["completion_tokens"]

        latencies.append(elapsed)
        total_tokens.append(completion_tokens)

    avg_latency = statistics.mean(latencies)
    p99_latency = sorted(latencies)[int(0.99 * len(latencies))]
    avg_throughput = statistics.mean(total_tokens) / avg_latency

    print(f"평균 지연: {avg_latency:.3f}s")
    print(f"P99 지연: {p99_latency:.3f}s")
    print(f"평균 처리량: {avg_throughput:.1f} tokens/s")
    return avg_latency, p99_latency, avg_throughput

# 테스트 프롬프트
test_prompt = "다음 주제에 대해 상세히 설명해주세요: 마이크로서비스 아키텍처"

print("=== Speculative Decoding 벤치마크 ===")
measure_latency(test_prompt)

# AWQ 양자화 모델을 vLLM에서 사용하는 예시
from vllm import LLM, SamplingParams

# AWQ 양자화 모델 로드 (HuggingFace에서 AWQ 모델 직접 사용)
llm_awq = LLM(
    model="TheBloke/Llama-2-70B-Chat-AWQ",
    quantization="awq",
    tensor_parallel_size=2,        # 4비트이므로 GPU 2장이면 충분
    gpu_memory_utilization=0.90,
    max_model_len=4096,
)

# GPTQ 양자화 모델 로드
llm_gptq = LLM(
    model="TheBloke/Llama-2-70B-Chat-GPTQ",
    quantization="gptq",
    tensor_parallel_size=2,
    gpu_memory_utilization=0.90,
)

# FP8 양자화 (H100 이상 권장)
llm_fp8 = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    quantization="fp8",
    tensor_parallel_size=4,
    gpu_memory_utilization=0.90,
)

# 성능 비교 벤치마크
import time

sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
test_prompts = ["마이크로서비스 아키텍처의 장단점을 설명해주세요."] * 10

for name, model in [("AWQ", llm_awq), ("GPTQ", llm_gptq), ("FP8", llm_fp8)]:
    start = time.time()
    outputs = model.generate(test_prompts, sampling_params)
    elapsed = time.time() - start
    total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
    print(f"{name}: {elapsed:.2f}s, {total_tokens/elapsed:.1f} tokens/s")

# GPU 모니터링 - nvidia-smi 활용
watch -n 1 nvidia-smi

# vLLM 메트릭 확인 (Prometheus 포맷)
curl http://localhost:8000/metrics | grep -E "vllm_(num_requests|gpu_cache|avg_generation)"

# 핵심 모니터링 대상:
# - vllm:num_requests_running: 현재 처리 중인 요청 수
# - vllm:num_requests_waiting: 대기 중인 요청 수
# - vllm:gpu_cache_usage_perc: GPU KV 캐시 사용률
# - vllm:avg_generation_throughput_toks_per_s: 평균 토큰 생성 처리량