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

Arithmetic Intensity (AI) = Number of FLOPs / Memory accessed (bytes)

Performance ceiling = min(Peak FLOPS, Peak Memory BW × AI)

def calc_model_memory_gb(
    num_params: int,       # Number of parameters (e.g., 70e9)
    dtype_bytes: int = 2,  # FP16=2, FP32=4, INT8=1, INT4=0.5
) -> float:
    """Calculate model weight memory"""
    return (num_params * dtype_bytes) / (1024 ** 3)

# Major model memory (FP16 baseline)
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 cache memory calculation.
    Per layer: 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 example
# 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 uses kv_heads
    head_dim=128,
    seq_len=4096,
    batch_size=1,
    dtype_bytes=2,
)
print(f"KV cache (seq=4096, bs=1): {kv_mem:.2f} GB")
# Output: KV cache (seq=4096, bs=1): 0.50 GB

# KV cache by batch size
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:
    """Approximate activation memory during inference"""
    # Per layer: attention + FFN activations
    # Approximation: 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 Virtual Memory    → PagedAttention
─────────────────────────────────────
Virtual page         → Logical block
Physical frame       → Physical block
Page table           → Block table
Page fault           → Block allocation

Request A: [Block 0] → [Block 1] → [Block 2]
                                        ↕ Physical block sharing (common prompt)
Request B: [Block 0] → [Block 1] → [Block 3]

# Install vLLM
pip install vllm

# Start server (single 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-compatible API call
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": "Explain GPU memory optimization"}],
    "max_tokens": 512,
    "temperature": 0.7
  }'

# Python client calling 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": "You are a GPU optimization expert."},
        {"role": "user", "content": "Explain how PagedAttention works."},
    ],
    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 quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,  # Compute in BF16
    bnb_4bit_quant_type="nf4",              # NormalFloat4 quantization
    bnb_4bit_use_double_quant=True,         # Double quantization for extra compression
)

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

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",      # Automatic multi-GPU distribution
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# Check memory usage
print(f"GPU memory: {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 quantization
    group_size=128,      # Group size (smaller = more accurate but more memory)
    desc_act=False,      # Describe activation ordering
    damp_percent=0.01,   # Hessian damping coefficient
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# Prepare calibration data (representative text samples)
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,
    # Recommend 1024+ samples in practice
]

# Load model and quantize
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 quantization complete!")

# Load quantized model
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 quantization configuration
quant_config = {
    "zero_point": True,   # Zero-point quantization
    "q_group_size": 128,  # Group size
    "w_bit": 4,           # 4-bit
    "version": "GEMM",    # Matrix multiplication kernel
}

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 quantization complete!")

Static Batching (batch_size=3):

Time →
[Request A: ████████████░░░░░░░░]  (12 tokens generated)
[Request B: ████░░░░░░░░░░░░░░░░]  (4 tokens generated)
[Request C: ████████░░░░░░░░░░░░]  (8 tokens generated)
                └─ B and C must wait for A to finish (GPU wasted)

Continuous Batching:

Step 1: [A1][B1][C1]  ← 3 processed simultaneously
Step 2: [A2][B2][C2]
Step 3: [A3][B3][C3]  ← B completes, new request D added
Step 4: [A4][C4][D1]  ← Empty slot immediately filled
Step 5: [A5][C5][D2]
Step 6: [A6][C6][D3]  ← C completes, new request E added
...

from vllm import LLM, SamplingParams

# Tensor parallel across 4 GPUs
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,        # Distribute across 4 GPUs
    dtype="bfloat16",
    max_model_len=8192,
    gpu_memory_utilization=0.90,
    enforce_eager=False,           # Use CUDA graph optimization
)

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

prompts = [
    "Explain the GPU memory hierarchy",
    "What are the benefits of PagedAttention?",
    "Compare quantization methods",
]

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated = output.outputs[0].text
    print(f"Prompt: {prompt[:50]}...")
    print(f"Generated: {generated[:100]}...")
    print()

Attention head distribution (4-way tensor parallel):

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

Each GPU computes independently, then AllReduce aggregates results

Llama-3.1-70B (80 layers) → 4-way pipeline:

GPU 0: Layers 0-19
GPU 1: Layers 20-39
GPU 2: Layers 40-59
GPU 3: Layers 60-79

Sequential layer processing, activations forwarded between GPUs

import torch

def profile_gpu_memory(func, *args, **kwargs):
    """Profile GPU memory usage"""
    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"Memory increase: {(after - before) / 1e9:.3f} GB")
    print(f"Peak memory:     {peak / 1e9:.3f} GB")
    print()
    print(torch.cuda.memory_summary())
    return result

# Example memory stats output
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)

# Optimized vLLM server configuration example
vllm_config = {
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "dtype": "bfloat16",
    "tensor_parallel_size": 1,
    "gpu_memory_utilization": 0.90,    # Use 90% of GPU memory
    "max_model_len": 8192,
    "max_num_batched_tokens": 8192,    # Max tokens per batch
    "max_num_seqs": 256,               # Max concurrent sequences
    "enable_chunked_prefill": True,    # Enable chunked prefill
    "block_size": 16,                  # KV cache block size (PagedAttention)
    "swap_space": 4,                   # CPU swap space in GB
    "enforce_eager": False,            # Use CUDA graphs
    "disable_log_stats": False,
}