TGI 서빙 엔진: 대형 언어 모델 배포의 혁신과 핵심 기능 소개

TGI (Text Generation Inference)

• LLM을 효율적이고 빠르게 서빙하기 위한 툴

• 대규모 배포 및 실시간 애플리케이션에 적합

• Hugging Face의 Hugging Chat, Inference API, Inference Endpoints에서 프로덕션으로 사용 중

주요 특징

아키텍처 구조

Client → [Rust Router (HTTP)] → [gRPC] → [Python Model Server (GPU)]
              ↓                              ↓
    - 요청 스케줄링                   - 모델 추론
    - Continuous Batching            - Flash Attention
    - Token Streaming                - PagedAttention
    - 큐 관리                        - Tensor Parallelism

1. Router (Rust) 🦀: HTTP 요청을 받아 스케줄링하고, Continuous Batching을 관리합니다. Rust로 구현되어 매우 빠른 처리 속도를 제공합니다.

2. Model Server (Python) 🐍: 실제 GPU에서 모델 추론을 수행합니다. PyTorch 기반으로 Flash Attention, PagedAttention 등의 최적화 기법을 적용합니다.

3. gRPC 통신: Router와 Model Server 사이를 gRPC로 연결하여 효율적인 직렬화와 낮은 레이턴시를 보장합니다.

주요 구성 요소

Streaming

from huggingface_hub import InferenceClient

client = InferenceClient("http://127.0.0.1:8080")
for token in client.text_generation(
	"How do you make cheese?",
	max_new_tokens=12, 
	stream=True):
    print(token)
    
# To
# make
# cheese
#,
# you
# need
# to
# start
# with
# milk
#.

Quantization (양자화)

# 8bit 양자화
docker run --gpus all --shm-size 1g -p 8080:80 \
  -v $volume:/data ghcr.io/huggingface/text-generation-inference:latest \
  --model-id $model --quantize bitsandbytes

# 4bit 양자화
docker run --gpus all --shm-size 1g -p 8080:80 \
  -v $volume:/data ghcr.io/huggingface/text-generation-inference:latest \
  --model-id $model --quantize bitsandbytes-nf4

Tensor Parallelism

[행렬 A] × [행렬 B] = [결과 C]

GPU 0: [A의 절반] × [B의 절반] = [C의 부분 결과 0]
GPU 1: [A의 절반] × [B의 절반] = [C의 부분 결과 1]

AllReduce → [C의 부분 결과 0] + [C의 부분 결과 1] = [결과 C]

# 4개 GPU로 Tensor Parallelism 실행
docker run --gpus all --shm-size 1g -p 8080:80 \
  -v $volume:/data ghcr.io/huggingface/text-generation-inference:latest \
  --model-id $model --num-shard 4

PagedAttention

• 블록(Block): KV Cache를 고정 크기 블록으로 분할 (OS의 페이지에 해당)

• 토큰(Token): 각 블록 안의 개별 단위 (OS의 바이트에 해당)

• 시퀀스(Sequence): 하나의 요청 처리 과정 (OS의 프로세스에 해당)

• 블록 테이블(Block Table): 논리적 블록 → 물리적 블록 매핑 (OS의 페이지 테이블에 해당)

PagedAttention의 효과

Flash Attention

[기존 Attention]
HBM → Q, K 로드 → S = QK^T 계산 → S를 HBM에 저장
HBM → S 로드 → P = softmax(S) 계산 → P를 HBM에 저장
HBM → P, V 로드 → O = PV 계산 → O를 HBM에 저장
(HBM 읽기/쓰기 6회)

[Flash Attention]
HBM → Q, K, V의 블록(타일) 로드 → SRAM에서 S, P, O 한 번에 계산 → O를 HBM에 저장
(HBM 읽기/쓰기 2회)

1. Tiling (타일링): Q, K, V 행렬을 작은 블록(타일)으로 분할하여 GPU의 빠른 SRAM에 로드

2. 커널 퓨전: Softmax, MatMul 등 여러 연산을 하나의 GPU 커널로 합쳐서 중간 결과의 HBM 왕복을 제거

3. Online Softmax: 전체 행을 보지 않고도 softmax를 점진적으로 계산하는 알고리즘 사용

Continuous Batching

[Static Batching]
요청 A (10 토큰) ████████████░░░░░░░░ (끝났지만 대기)
요청 B (20 토큰) ████████████████████ (완료)
요청 C (15 토큰) ███████████████░░░░░ (끝났지만 대기)
                  → B가 끝날 때까지 모두 대기 → 다음 배치 시작

[Continuous Batching]
요청 A (10 토큰) ██████████
요청 D (새 요청)           ████████████████  ← A 완료 즉시 투입
요청 B (20 토큰) ████████████████████
요청 C (15 토큰) ███████████████
요청 E (새 요청)                ██████████  ← C 완료 즉시 투입
                  → 빈 슬롯에 즉시 새 요청 투입

TGI 빠르게 시작하기

# 기본 실행 (Llama 3.1 8B 예시)
model=meta-llama/Llama-3.1-8B-Instruct
volume=$PWD/data

docker run --gpus all --shm-size 1g -p 8080:80 \
  -v $volume:/data \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model-id $model

# API 호출 테스트
curl localhost:8080/generate \
  -X POST \
  -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
  -H 'Content-Type: application/json'

주요 실행 옵션

TGI vs vLLM vs SGLang 비교

정리

핵심 정리

참고 링크