Skip to content

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

|

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

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: FREEBlock 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     │ 실용성   │
├──────────┼──────────┼──────────┼──────────┼──────────┤
FP16140 GB   │ 기준선   │ 80.9%8×H100BF16140 GB+5%80.9%8×H100INT870 GB+10%80.2%2×H100GPTQ-4b  │ 36 GB+30%79.8%1×H100AWQ-4b   │ 36 GB+35%79.5%1×H100GGUF-Q438 GBCPU 가능 │ 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: GPU0GPU1GPU2GPU3
  마이크로 배치 2: GPU0GPU1GPU2GPU3
  ...

파이프라인 버블 (비효율):
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에 적용했듯, 앞으로도 컴퓨터 과학의 오래된 지혜가 새로운 형태로 재발견될 것이다. 각 기술이 풀고자 하는 근본 문제를 이해하면, 새로운 최적화가 나올 때마다 빠르게 습득하고 적용할 수 있다.

Complete LLM Serving Optimization Guide: KV Cache, PagedAttention, and Quantization

The Two Phases of LLM Inference: Prefill and Decode

LLM text generation splits into two fundamentally different phases. Without understanding this split, optimization is impossible.

Phase 1: PREFILL (process the input)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Input: "What is the capital of France?"
       └─ all 9 tokens processed at once

What happens:
  - All input tokens are processed in parallel (big matrix multiply!)
  - Q, K, V are computed for each input token
  - KV cache is created (saves K, V for later reuse)
  - First output token is generated

Characteristics:
  - GPU operation: COMPUTE-BOUND (matrix × matrix)
  - GPU utilization: HIGH  - Latency metric: TTFT (Time To First Token)

Phase 2: DECODE (generate tokens one-by-one)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Generates: "Paris""is""the""capital"...

What happens:
  - Generate exactly one token per forward pass
  - Compute Q for the new token, attend over cached K, V
  - Must read ALL model weights for every single token

Characteristics:
  - GPU operation: MEMORY-BOUND (matrix × vector)
  - GPU utilization: LOW (often 520%!)
  - Throughput metric: TBT (Time Between Tokens)

This is why LLM serving is hard to optimize:
the two phases have completely different bottlenecks!

Measuring it in practice:

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 the transformer architecture in detail:"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    input_len = inputs["input_ids"].shape[1]

    # Measure prefill time
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        _ = model(**inputs)   # forward pass on input only
    torch.cuda.synchronize()
    t_prefill = time.perf_counter() - t0

    # Measure decode time
    t0 = 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() - t0

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

    print(f"Input tokens:          {input_len}")
    print(f"Prefill time (TTFT):   {t_prefill*1000:.1f}ms")
    print(f"Generated tokens:      {n_new}")
    print(f"Decode time:           {t_decode*1000:.1f}ms")
    print(f"Per-token decode time: {t_decode/n_new*1000:.1f}ms/token")
    # Llama-1B on H100 (~):
    # Prefill: ~5ms (linear in input length)
    # Decode:  ~3ms/token (proportional to model size, batch-dependent)

KV Cache: The Memory Dilemma of Attention

What Happens Without a KV Cache?

Autoregressive generation WITHOUT KV cache:

Step 1: [token_1] → generate token_2
  - Compute Q1,K1,V1 for token_1
  - Compute Q2,K2,V2 for token_2 (partial)
  - Attention: Q2 × [K1,K2]^T
  - Ops: 2^2 = 4 dot products

Step 2: [token_1, token_2] → generate token_3
  - Re-compute K1,V1 (wasted work!)
  - Re-compute K2,V2 (wasted work!)
  - Compute Q3,K3,V3
  - Attention: Q3 × [K1,K2,K3]^T
  - Ops: 3^2 = 9 dot products

Step N: O(N^2) operations per token
Total for L tokens: O(L^3) total compute
  100 tokens:  1,000,000 dot products
  1000 tokens: 1,000,000,000 dot products

KV Cache: Reuse Previous Computation

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, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        # KV cache storage
        self.k_cache = None  # (batch, heads, past_len, d_k)
        self.v_cache = None

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

        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:
            # Append new K, V to the cache
            K = torch.cat([self.k_cache, K], dim=2)
            V = torch.cat([self.v_cache, V], dim=2)

        if use_cache:
            self.k_cache = K.detach()
            self.v_cache = V.detach()

        # Q is only the current token(s); K, V are the full sequence
        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)


def compute_kv_cache_bytes(seq_len, n_layers, n_kv_heads, head_dim,
                            batch_size, dtype_bytes=2):
    """
    KV cache memory in bytes.
    Factor of 2: one tensor for K, one for V.
    """
    return 2 * n_layers * n_kv_heads * head_dim * seq_len * batch_size * dtype_bytes


# Llama 3.1 70B (uses GQA: 8 KV heads, 64 Q heads):
size = compute_kv_cache_bytes(
    seq_len=4096, n_layers=80, n_kv_heads=8,
    head_dim=128, batch_size=1, dtype_bytes=2
)
print(f"KV cache (Llama-70B, seq=4096, batch=1): {size/1e9:.1f} GB")
# Result: ~6.7 GB per request
# batch=32: ~214 GB → won't fit in one H100 (80 GB)!

The Memory Fragmentation Problem

Traditional KV cache allocation (pre-allocate max_seq_len per request):

┌───────────────────────────────────────────────────────┐
Request 1: current_len=100, reserved=512│ ████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░         │
[used: 100]  [wasted: 412 slots = 80%!]├───────────────────────────────────────────────────────┤
Request 2: current_len=50,  reserved=512│ ██████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░         │
[used: 50]   [wasted: 462 slots = 90%!]├───────────────────────────────────────────────────────┤
Request 3: current_len=300, reserved=512│ █████████████████████████████████████░░░░░░░░         │
[used: 300]  [wasted: 212 slots = 41%!]└───────────────────────────────────────────────────────┘

Total allocated: 3 × 512 = 1536 slots
Total used:      450 slots
Wasted:          1086 slots = 71%

Internal fragmentation (allocated but unused) +
External fragmentation (gaps between allocations)
Typical real-world GPU memory utilization: 2040%

PagedAttention (vLLM): Virtual Memory Saves LLMs

The Core Insight: OS Virtual Memory Applied to KV Cache

Kwon et al. (UC Berkeley, 2023): "Operating systems solved memory fragmentation decades ago. Apply the same idea to KV cache."

The OS Virtual Memory Lesson:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
virtual address → page table → physical address
Process sees contiguous virtual space
Physical memory can be non-contiguous
No fragmentation, efficient RAM usage

PagedAttention Analogy:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
virtual KV slot → block table → physical block
Sequence sees contiguous virtual slots
Physical GPU blocks can be non-contiguous
Near-zero fragmentation, efficient GPU VRAM usage
PagedAttention Memory Layout:

GPU memory split into fixed-size blocks (default: 16 tokens each):
┌─────────────────────────────────────────────────────────┐
Physical KV Cache BlocksBlock 0: [tok0–15]    Block 1: [tok16–31]Block 2: [tok32–47]   Block 3: [tok48–63]Block 4: [tok64–79]   Block 5: FREEBlock 6: FREE         Block 7: FREE└─────────────────────────────────────────────────────────┘

Block Table (same role as OS page table):
┌──────────┬──────────────────────────────────────────────┐
RequestVirtual block → Physical block mapping       │
├──────────┼──────────────────────────────────────────────┤
Req 1    │ virt[0]→phys[0], virt[1]→phys[2] (tokens 015: block 0; tokens 3247: block 2)├──────────┼──────────────────────────────────────────────┤
Req 2    │ virt[0]→phys[1], virt[1]→phys[3] (tokens 015: block 1; tokens 1631: block 3)└──────────┴──────────────────────────────────────────────┘

Key properties:
- Blocks are allocated ON-DEMAND as the sequence grows
- Internal fragmentation < 1 block = at most 15 wasted slots
- Blocks can be SHARED across requests (same prefix)!
# vLLM with PagedAttention:
from vllm import LLM, SamplingParams
import time

def benchmark_vllm():
    llm = LLM(
        model="meta-llama/Llama-3.2-8B-Instruct",
        gpu_memory_utilization=0.9,
        max_model_len=8192,
        block_size=16,
        max_num_seqs=256,
    )

    prompts = [
        "Short question: What is Python?",
        "Medium question: " + "Explain the history of machine learning. " * 5,
        "Long question: " + "How do you build a transformer from scratch? " * 10,
    ] * 20  # 60 requests of varying lengths

    params = SamplingParams(temperature=0.0, max_tokens=100)

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

    total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
    print(f"Requests:         {len(prompts)}")
    print(f"Tokens generated: {total_tokens}")
    print(f"Elapsed:          {elapsed:.1f}s")
    print(f"Throughput:       {total_tokens/elapsed:.0f} tokens/s")


# Memory efficiency improvement:
# Traditional:     20–40% of GPU memory used for actual KV cache
# PagedAttention:  >95% of GPU memory used for actual KV cache
# Result: 2–3× more concurrent requests on the same GPU

Prefix Caching: Share Common Prompts

# Enable prefix caching in vLLM:
llm = LLM(
    model="meta-llama/Llama-3.2-8B-Instruct",
    enable_prefix_caching=True,
)

# Many requests sharing a long system prompt:
system = "You are an expert software engineer. " * 100  # long system prompt

requests = [
    system + "User: What is a binary search tree?",
    system + "User: How does garbage collection work?",
    system + "User: Explain ACID properties in databases.",
]

# The KV cache for `system` is computed ONCE and shared across all 3 requests.
# Prefill cost: computed 1 time instead of 3 times (3× savings on prefill!)
# This matters hugely for RAG pipelines where context is repeated.

Continuous Batching: Maximizing Throughput

The Problem with Static Batching

Static (request-level) batching:

GPU batch at each step:
Step 1:  [Req1: running] [Req2: running] [Req3: running]
Step 2:  [Req1: running] [Req2: DONE   ] [Req3: running]
Step 3:  [Req1: running] [  idle/wait  ] [Req3: running]GPU waste!
Step 4:  [Req1: DONE   ] [  idle/wait  ] [Req3: running]GPU waste!
Step 5:  [  idle/wait  ] [  idle/wait  ] [Req3: DONE   ]GPU waste!

New requests must wait until the ENTIRE batch finishes.
GPU waste rate: often 50%+

Continuous Batching: Dynamic Scheduling Every Token Step

Continuous (iteration-level) batching:

Step 1:  [Req1] [Req2] [Req3]
Step 2:  [Req1] [Req2] [Req3]
Step 3:  [Req1] [Req4] [Req3]Req2 done → Req4 inserted immediately!
Step 4:  [Req5] [Req4] [Req3]Req1 done → Req5 inserted!
Step 5:  [Req5] [Req4] [Req6]Req3 done → Req6 inserted!

GPU is at maximum utilization at every step.
Throughput improvement over static batching: 24×
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.engine.arg_utils import AsyncEngineArgs
import asyncio

async def run_continuous_batching_server():
    engine_args = AsyncEngineArgs(
        model="meta-llama/Llama-3.2-8B-Instruct",
        max_num_seqs=256,               # max concurrent sequences
        max_num_batched_tokens=8192,    # max tokens per batch step
    )
    engine = AsyncLLMEngine.from_engine_args(engine_args)

    async def generate_one(prompt, req_id):
        from vllm import SamplingParams
        params = SamplingParams(temperature=0.7, max_tokens=200)
        async for output in engine.generate(prompt, params, request_id=req_id):
            if output.finished:
                return output.outputs[0].text

    # Requests submitted concurrently — engine handles continuous batching
    results = await asyncio.gather(
        generate_one("Explain quantum entanglement.", "r1"),
        generate_one("Write a Python quicksort.", "r2"),
        generate_one("Summarize the French Revolution.", "r3"),
    )
    for r in results:
        print(r[:80])

Quantization: Trade Precision for Speed and Memory

Why Quantization?

LLM memory footprint (FP16):
  Llama 3.1 8B:    16 GB
  Llama 3.1 70B:   140 GB
  Llama 3.1 405B:  810 GB

Common GPU memory:
  RTX 4090:     24 GB  → tight even for 8B
  A100 80GB:           → 70B impossible on one card
  H100 80GB:           → 70B impossible on one card
  H100 ×8 (640 GB):    → 70B fine, 405B barely

Memory savings with quantization:
  FP16  (16-bit): baseline
  INT8   (8-bit): 50% saved, ~1%  accuracy loss
  INT4   (4-bit): 75% saved, ~23% accuracy loss
  INT3   (3-bit): 81% saved, use cautiously
  INT2   (2-bit): 88% saved, usually unacceptable

Post-Training Quantization: INT8 (LLM.int8())

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# INT8 quantization — Dettmers et al., 2022 (bitsandbytes):
config_int8 = BitsAndBytesConfig(
    load_in_8bit=True,
    # Optionally keep certain layers in FP16 (e.g., output head)
    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",
)
# 70B model: 140 GB (FP16) → 70 GB (INT8), ~1% accuracy loss

# The key insight behind LLM.int8():
# Problem: activation outliers in certain channels ruin naive INT8 quality
# Solution: "Mixed-precision decomposition"
#   - Detect outlier channels (top ~1% by magnitude)
#   - Keep those channels in FP16
#   - Quantize all other channels to INT8
#   → Near-lossless quality with ~50% memory savings

4-bit Quantization: NF4 and GPTQ

# NF4 quantization (QLoRA paper, Dettmers et al. 2023):
config_4bit = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,   # compute in FP16
    bnb_4bit_quant_type="nf4",              # NormalFloat4
    bnb_4bit_use_double_quant=True,         # quantize the scale too!
)

model_4bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B-Instruct",
    quantization_config=config_4bit,
    device_map="auto",
)
# 70B: 140 GB → 35 GB, ~2–3% accuracy loss

# Why NF4?
# Neural network weights are approximately normally distributed.
# NF4 uses 16 codepoints placed at equal-probability quantiles
# of a standard normal distribution.
# Each codepoint covers an equal probability mass → minimal quantization error
# vs. uniform INT4 which distributes points evenly on the number line.

# Double quantization:
# Quantization scale factors are FP32: 1 per group of 64 weights
# Double-quant quantizes those scale factors to INT8 too
# Net savings: ~0.5 additional bits per weight


# GPTQ (Frantar et al., 2022) — layer-wise optimal quantization:
from auto_gptq import AutoGPTQForCausalLM

model_gptq = AutoGPTQForCausalLM.from_quantized(
    "TheBloke/Llama-2-70B-GPTQ",
    device="cuda:0",
    use_triton=True,     # Triton kernels for faster inference
)
# GPTQ uses the Hessian of each layer's loss to minimize quantization error.
# Generally highest accuracy among INT4 methods.

AWQ: Activation-Aware Weight Quantization

# AWQ (Lin et al., 2023) key insight:
# Not all weights are equally important!
# ~1% of channels produce large activations — these are "salient"
# Naively quantizing them to INT4 crushes quality

# AWQ solution:
# 1. Run calibration data, record activation magnitudes per channel
# 2. Scale up salient channels in the weight matrix (per-channel scaling)
# 3. Quantize everything to INT4 — the scaling absorbs the error for salient channels

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

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

model     = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

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

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

# AWQ vs GPTQ:
# AWQ:  faster inference (hand-tuned CUDA/Triton kernels)
#       ~25% of FP16 memory
#       accuracy: slightly below GPTQ
# GPTQ: higher accuracy (Hessian-based error minimization)
#       similar inference speed
#       same memory as AWQ

Quantization Comparison Table

Llama 3.1 70B quantization comparison (single H100):

┌──────────┬──────────┬────────────┬──────────┬───────────────────┐
FormatMemoryThroughputMMLUHardware needed   │
├──────────┼──────────┼────────────┼──────────┼───────────────────┤
FP16140 GB   │ baseline   │ 80.9%8× H100BF16140 GB+5%80.9%8× H100INT870 GB+10%80.2%2× H100GPTQ-4b  │  36 GB+30%79.8%1× H100AWQ-4b   │  36 GB+35%79.5%1× H100GGUF-Q438 GBCPU ok     │ 79.1%CPU or 1× H100└──────────┴──────────┴────────────┴──────────┴───────────────────┘

Speculative Decoding: Free Lunch Exists

The Idea: Draft Fast, Verify in Parallel

Standard decode:
  70B model generates 1 token = 1 forward pass = ~10ms
  100 tokens = ~1000ms = 1 second

Speculative decoding (Leviathan et al., 2023):
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Step 1: Draft model (7B) generates K tokens quickly
        ["Paris"] ["is"] ["the"] ["capital"]
        4 tokens in ~2ms (7B model)

Step 2: Target model (70B) verifies all K tokens in ONE forward pass!
        Process all 4 draft tokens in parallel → ~10ms
        (same cost as generating just 1 token normally)

Step 3: Verify each draft token:
        "Paris""is""the""capital"        → accept 3 tokens, reject from position 4

Step 4: Resample from target model distribution at rejection point
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Result: 3 accepted tokens in ~12ms (vs 30ms standard decode)
Speedup: 2.5× (varies with acceptance rate ~7090%)
Quality: ZERO degradation (target model is the arbiter)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def speculative_decode(
    target_model,
    draft_model,
    input_ids,
    max_new_tokens=100,
    K=4,              # number of draft tokens per speculation
    temperature=1.0,
):
    """
    Speculative decoding: draft model proposes K tokens,
    target model verifies them all in one forward pass.
    Guarantees exactly the same distribution as target-only decoding.
    """
    generated = input_ids.clone()

    while generated.shape[1] < input_ids.shape[1] + max_new_tokens:
        # --- Phase 1: Draft model generates K candidates ---
        draft_ids   = []
        draft_probs = []

        ctx = generated.clone()
        for _ in range(K):
            with torch.no_grad():
                out    = draft_model(ctx)
                logits = out.logits[:, -1, :] / max(temperature, 1e-5)
                probs  = torch.softmax(logits, dim=-1)
                tok    = torch.multinomial(probs, 1)
                draft_ids.append(tok)
                draft_probs.append(probs[0, tok[0, 0]])
                ctx = torch.cat([ctx, tok], dim=1)

        # --- Phase 2: Target model verifies K positions simultaneously ---
        candidate = torch.cat([generated] + draft_ids, dim=1)
        with torch.no_grad():
            tgt_out    = target_model(candidate)
            # logits for positions where draft tokens are placed
            tgt_logits = tgt_out.logits[:, len(generated[0])-1:-1, :]
            tgt_probs  = torch.softmax(tgt_logits / max(temperature, 1e-5), dim=-1)

        # --- Phase 3: Accept/reject each draft token ---
        n_accepted = 0
        for i in range(K):
            token_id  = draft_ids[i][0, 0].item()
            p_target  = tgt_probs[0, i, token_id].item()
            p_draft   = draft_probs[i].item()

            # Acceptance probability: min(1, p_target / p_draft)
            accept_p = min(1.0, p_target / max(p_draft, 1e-8))
            if torch.rand(1).item() < accept_p:
                generated = torch.cat([generated, draft_ids[i]], dim=1)
                n_accepted += 1
            else:
                # Reject: resample from adjusted target distribution
                adjusted = torch.clamp(tgt_probs[0, i] - tgt_probs[0, i], min=0)
                # Correct adjusted distribution (residual of target minus draft)
                diff = tgt_probs[0, i].clone()
                diff[token_id] = max(0.0, diff[token_id] - p_draft)
                diff = diff / diff.sum().clamp(min=1e-8)
                new_tok = torch.multinomial(diff.unsqueeze(0), 1)
                generated = torch.cat([generated, new_tok], dim=1)
                break

        if n_accepted == K:
            # All accepted: also take the target model's bonus token
            bonus_logits = tgt_out.logits[:, -1, :] / max(temperature, 1e-5)
            bonus_probs  = torch.softmax(bonus_logits, dim=-1)
            bonus_tok    = torch.multinomial(bonus_probs, 1)
            generated    = torch.cat([generated, bonus_tok], dim=1)

    return generated


# Real speedups observed (A100, Llama-2 70B + Llama-2 7B draft):
# K=4: 2.3× speedup, acceptance rate ~80%
# K=8: 2.7× speedup, acceptance rate ~75%
# Optimal K depends on draft/target quality ratio

Tensor Parallelism and Pipeline Parallelism

Tensor Parallelism: Split Layers Across GPUs

Tensor Parallelism (Shoeybi et al., 2019Megatron-LM):

70B model, 8 GPUs, 64 attention heads:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
GPU 0: heads 0–7    (W_q slice: 1/8 of full matrix)
GPU 1: heads 815
GPU 2: heads 1623
GPU 3: heads 2431
GPU 4: heads 3239
GPU 5: heads 4047
GPU 6: heads 4855
GPU 7: heads 5663

Each GPU computes its heads independently,
then All-Reduce merges results.

Communication cost:
  1 All-Reduce per attention layer
  1 All-Reduce per FFN layer
  NVLink (H100): 900 GB/s bidirectional → viable
  PCIe:           64 GB/s              → too slow for TP>2
import torch
import torch.distributed as dist

def tensor_parallel_linear(x, W_local, rank, world_size):
    """
    Column-parallel linear (W split along output dimension).
    x:       (batch, seq, d_model)   -- replicated on all GPUs
    W_local: (d_model, d_out//world_size) -- each GPU holds a shard
    """
    # Each GPU computes its output shard
    out_local = x @ W_local    # (batch, seq, d_out//world_size)

    # All-Gather to reconstruct full output on every GPU
    out_list = [torch.zeros_like(out_local) for _ in range(world_size)]
    dist.all_gather(out_list, out_local)
    out_full = torch.cat(out_list, dim=-1)   # (batch, seq, d_out)

    return out_full


# For row-parallel (W split along input dimension):
def tensor_parallel_linear_row(x_local, W_local, rank, world_size):
    """
    Row-parallel linear: x is already sharded across GPUs.
    x_local: (batch, seq, d_in//world_size)
    W_local: (d_in//world_size, d_out)
    """
    partial = x_local @ W_local    # (batch, seq, d_out) — partial sum
    dist.all_reduce(partial, op=dist.ReduceOp.SUM)   # sum partial results
    return partial

Pipeline Parallelism: Split Layers Sequentially

Pipeline Parallelism:

70B model, 80 layers, 4 GPUs:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
GPU 0: layers  0–19   (embedding + first 20 transformer layers)
GPU 1: layers 2039
GPU 2: layers 4059
GPU 3: layers 6079  + LM head

With micro-batching to hide pipeline bubbles:

           | mb1 | mb2 | mb3 | mb4 | mb5 |
GPU 0 →→→ [f1 ] [f2 ] [f3 ] [f4 ] [f5 ] [b5 ] [b4 ] [b3 ] [b2 ] [b1 ]
GPU 1      [   ] [f1 ] [f2 ] [f3 ] [f4 ] [f5 ] [b5 ] [b4 ] [b3 ] [b2 ] [b1 ]
GPU 2            [   ] [f1 ] [f2 ] [f3 ] [f4 ] [f5 ]            [b1 ]
GPU 3                  [   ] [f1 ] [f2 ] [f3 ] [f4 ] [f5 ] [b5 ]

Pipeline bubble ratio = (p - 1) / (m + p - 1)
  p = number of pipeline stages
  m = number of micro-batches
Larger m = smaller bubble = better efficiency

vLLM vs TGI vs TensorRT-LLM: Framework Comparison

LLM serving framework comparison (as of early 2026):

┌───────────────────┬────────────────────────────────────────────────┐
Framework         │ vLLM                                           │
├───────────────────┼────────────────────────────────────────────────┤
DeveloperUC Berkeley / vLLM community                   │
Key innovations   │ PagedAttention, continuous batching            │
QuantizationAWQ, GPTQ, INT8, FP8Throughput        │ ★★★★☆  HighTTFT latency      │ ★★★☆☆  MediumEase of use       │ ★★★★★  Very easy (Python-native)Customizability   │ ★★★☆☆  MediumLicenseApache 2.0NotesMost active OSS community, OpenAI-compat API└───────────────────┴────────────────────────────────────────────────┘

┌───────────────────┬────────────────────────────────────────────────┐
FrameworkTGI (Text Generation Inference)├───────────────────┼────────────────────────────────────────────────┤
DeveloperHugging FaceKey innovations   │ Continuous batching, FlashAttentionQuantizationGPTQ, AWQ, bitsandbytes                        │
Throughput        │ ★★★☆☆  MediumTTFT latency      │ ★★★☆☆  MediumEase of use       │ ★★★★☆  Easy (Docker-first)Customizability   │ ★★★★☆  HighLicenseHFOIL (check commercial terms)NotesNative HF ecosystem integration                │
└───────────────────┴────────────────────────────────────────────────┘

┌───────────────────┬────────────────────────────────────────────────┐
FrameworkTensorRT-LLM├───────────────────┼────────────────────────────────────────────────┤
DeveloperNVIDIAKey innovations   │ TensorRT graph optimization, in-flight batching│
QuantizationINT8, INT4, FP8, SmoothQuant, AWQThroughput        │ ★★★★★  Highest (NVIDIA GPUs only)TTFT latency      │ ★★★★★  LowestEase of use       │ ★★☆☆☆  Complex (C++ heavy)Customizability   │ ★★☆☆☆  DifficultLicenseApache 2.0NotesBest raw performance; use via Triton Server└───────────────────┴────────────────────────────────────────────────┘

┌───────────────────┬────────────────────────────────────────────────┐
Framework         │ llama.cpp / Ollama├───────────────────┼────────────────────────────────────────────────┤
Developer         │ ggerganov / Ollama Inc.                        
Key innovations   │ GGUF quantization, CPU+GPU hybrid              │
QuantizationQ2Q8 (GGUF format)Throughput        │ ★★☆☆☆  Low (on CPU)TTFT latency      │ ★★☆☆☆  HighEase of use       │ ★★★★★  Simplest possible                      │
Customizability   │ ★★☆☆☆  LimitedLicenseMITNotesIdeal for local dev, CPU inference, demos      │
└───────────────────┴────────────────────────────────────────────────┘

Decision Guide

Choose vLLM if:
  - Production serving, Python team, open source preferred
  - Need OpenAI-compatible API drop-in replacement
  - Want the best community support and newest features fastest

Choose TGI if:
  - Deep HuggingFace ecosystem integration
  - Docker-first deployment culture
  - Need robust SSE streaming out-of-the-box

Choose TensorRT-LLM if:
  - Maximum raw performance on NVIDIA hardware
  - Have a team comfortable with C++/CUDA tooling
  - Enterprise production with dedicated MLOps

Choose Ollama / llama.cpp if:
  - Local development, prototyping
  - CPU inference required
  - Simplicity over performance

Production LLM Serving Stack

Production LLM serving architecture:

Clients
Load Balancer (Nginx / AWS ALB / Cloudflare)
API Gateway (FastAPI / Kong)
Rate limiting, auth, logging, request validation
Router (model selection, priority queue)
  ├──→ vLLM server A: 8B model   (fast/cheap requests)
  ├──→ vLLM server B: 70B model  (high-quality requests)
  └──→ vLLM server C: domain-specific fine-tune
      Observability (Prometheus + Grafana)
      Key metrics:
        - TTFT p50/p95/p99
        - TBT  p50/p95/p99
        - Throughput (tokens/s)
        - GPU utilization %
        - KV cache utilization %
        - Request queue depth
# Production vLLM server launch command:
import subprocess

cmd = [
    "python", "-m", "vllm.entrypoints.openai.api_server",
    "--model", "meta-llama/Llama-3.1-8B-Instruct",
    "--tensor-parallel-size", "2",          # 2-GPU tensor parallelism
    "--gpu-memory-utilization", "0.9",
    "--max-model-len", "8192",
    "--max-num-seqs", "256",
    "--max-num-batched-tokens", "8192",
    "--quantization", "awq",
    "--enable-prefix-caching",
    "--block-size", "16",
    "--port", "8000",
    "--disable-log-requests",               # reduce logging overhead
]

# Calling the server from a client:
from openai import OpenAI

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

response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[
        {"role": "system", "content": "You are a helpful AI assistant."},
        {"role": "user",   "content": "What is the GIL in Python?"},
    ],
    temperature=0.0,
    max_tokens=500,
    stream=True,    # streaming response
)

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

LLM serving optimization sits at the intersection of hardware, algorithms, and systems software. PagedAttention borrowed from operating system design. FlashAttention rediscovered the principle of tiling from numerical linear algebra. Speculative decoding revived draft-and-verify ideas from branch prediction. The engineers who will build the next generation of LLM serving systems are those who understand not just the current tools but the first-principles reasoning behind them.