TGI 아키텍처

TGI가 하나의 프로그램처럼 동작하는 것처럼 보여도, 실제로는 여러 구성 요소가 협력해서 요청을 처리해요. 클라이언트 요청을 받는 라우터, 모델 서버를 띄우는 런처, 그리고 실제 추론을 담당하는 모델 서버로 나뉘죠. 이 글에서는 이 세 구성 요소가 어떤 역할을 하고, 서로 어떻게 통신하며 한 번의 생성 요청이 어떤 흐름으로 처리되는지 함께 살펴볼게요.

출처: Hugging Face Text Generation Inference Architecture

구성 요소

TGI는 크게 세 가지 구성 요소로 나뉘어 동작해요.

라우터(The router), 다른 말로 webserver라고도 불러요. 클라이언트 요청을 받아 버퍼링하고 배치를 만들어 모델 서버로 보낼 gRPC 호출을 준비해요.
런처(The launcher): 모델이 샤딩된 경우 하나 또는 여러 모델 서버를 띄우고, 호환되는 인자로 라우터까지 실행해 주는 도우미예요.
모델 서버(The model server): gRPC 요청을 받아 모델에서 실제 추론을 처리해요. 모델이 여러 가속기(예: 여러 GPU)에 샤딩돼 있다면 NCCL 등을 통해 샤드들이 서로 동기화돼요.

TRTLLM 같은 다른 백엔드에서는 모델 서버와 런처가 백엔드 전용으로 따로 구성돼요. 또 라우터와 모델 서버는 반드시 같은 머신에 둘 필요 없이, 서로 다른 머신에 배포해도 돼요.

라우터 (The Router)

라우터는 Rust로 작성된 웹 서버 바이너리예요. 커스텀 HTTP API와 OpenAI의 Messages API 요청을 모두 받아요. 요청을 받으면 배칭(batching) 로직을 처리하는데, 요청과 응답 사이의 지연, 특히 디코딩 지연을 줄이기 위한 여러 전략을 사용해요. 큐(queue), 스케줄러, 블록 할당자(block allocator)를 이용해 요청을 배치로 묶은 뒤 모델 서버로 보내요.

라우터는 설정 파일이 아니라 커맨드라인 인자로 파라미터를 받아요. 주요 옵션들을 보면 다음과 같아요.

Text Generation Webserver

Usage: text-generation-router [OPTIONS]

Options:
      --max-concurrent-requests <MAX_CONCURRENT_REQUESTS>
          [env: MAX_CONCURRENT_REQUESTS=] [default: 128]
      --max-best-of <MAX_BEST_OF>
          [env: MAX_BEST_OF=] [default: 2]
      --max-stop-sequences <MAX_STOP_SEQUENCES>
          [env: MAX_STOP_SEQUENCES=] [default: 4]
      --max-top-n-tokens <MAX_TOP_N_TOKENS>
          [env: MAX_TOP_N_TOKENS=] [default: 5]
      --max-input-tokens <MAX_INPUT_TOKENS>
          [env: MAX_INPUT_TOKENS=] [default: 1024]
      --max-total-tokens <MAX_TOTAL_TOKENS>
          [env: MAX_TOTAL_TOKENS=] [default: 2048]
      --waiting-served-ratio <WAITING_SERVED_RATIO>
          [env: WAITING_SERVED_RATIO=] [default: 1.2]
      --max-batch-prefill-tokens <MAX_BATCH_PREFILL_TOKENS>
          [env: MAX_BATCH_PREFILL_TOKENS=] [default: 4096]
      --max-batch-total-tokens <MAX_BATCH_TOTAL_TOKENS>
          [env: MAX_BATCH_TOTAL_TOKENS=]
      --max-waiting-tokens <MAX_WAITING_TOKENS>
          [env: MAX_WAITING_TOKENS=] [default: 20]
      --max-batch-size <MAX_BATCH_SIZE>
          [env: MAX_BATCH_SIZE=]
      --hostname <HOSTNAME>
          [env: HOSTNAME=] [default: 0.0.0.0]
  -p, --port <PORT>
          [env: PORT=] [default: 3000]
      --master-shard-uds-path <MASTER_SHARD_UDS_PATH>
          [env: MASTER_SHARD_UDS_PATH=] [default: /tmp/text-generation-server-0]
      --tokenizer-name <TOKENIZER_NAME>
          [env: TOKENIZER_NAME=] [default: bigscience/bloom]
      --tokenizer-config-path <TOKENIZER_CONFIG_PATH>
          [env: TOKENIZER_CONFIG_PATH=]
      --revision <REVISION>
          [env: REVISION=]
      --validation-workers <VALIDATION_WORKERS>
          [env: VALIDATION_WORKERS=] [default: 2]
      --json-output
          [env: JSON_OUTPUT=]
      --otlp-endpoint <OTLP_ENDPOINT>
          [env: OTLP_ENDPOINT=]
      --otlp-service-name <OTLP_SERVICE_NAME>
          [env: OTLP_SERVICE_NAME=]
      --cors-allow-origin <CORS_ALLOW_ORIGIN>
          [env: CORS_ALLOW_ORIGIN=]
      --ngrok
          [env: NGROK=]
      --ngrok-authtoken <NGROK_AUTHTOKEN>
          [env: NGROK_AUTHTOKEN=]
      --ngrok-edge <NGROK_EDGE>
          [env: NGROK_EDGE=]
      --messages-api-enabled
          [env: MESSAGES_API_ENABLED=]
      --disable-grammar-support
          [env: DISABLE_GRAMMAR_SUPPORT=]
      --max-client-batch-size <MAX_CLIENT_BATCH_SIZE>
          [env: MAX_CLIENT_BATCH_SIZE=] [default: 4]
  -h, --help
          Print help
  -V, --version
          Print version

모델 서버 (The Model Server)

모델 서버는 Python 서버로, gRPC 요청을 기다리며 대기하고, 주어진 모델을 로드해 텐서 병렬(tensor parallelism)을 위한 샤딩을 수행한 뒤 새 요청을 기다리며 유지돼요. 주로 CUDA/ROCm에서 최적화된 Pytorch 모델을 지원해요.

모델 서버 변형

Hugging Face가 공식 지원하는 모델 서버 변형은 여러 가지예요.

기본값은 Nvidia GPU용 CUDA에 최적화된 서버로, 코드는 메인 TGI 저장소에 있어요.
AMD ROCm에 최적화된 버전도 메인 저장소에 있고, 일부 모델 기능이 달라요.
Intel GPU에 최적화된 버전도 메인 저장소에 있으며 일부 기능이 달라요.
Intel Gaudi용 버전은 포크(fork)된 저장소에서 유지보수되며, 메인 TGI 저장소와 자주 재동기화돼요.
Neuron(AWS Inferentia2)용 버전은 메인 저장소에서 유지보수되고 일부 기능이 달라요.
Google TPU용 버전은 Optimum TPU의 일부로 관리돼요.

하드웨어와 미들웨어가 제공하는 최적화가 다르기 때문에 모든 변형이 동일한 기능을 제공하는 건 아니에요.

커맨드라인 인터페이스

공식 CLI는 download-weights, quantize, serve 세 가지 서브커맨드를 지원해요.

download-weights: 허브에서 가중치를 내려받고, 일부 변형에서는 구현에 맞는 형식으로 변환해요.
quantize: qptq 패키지를 이용해 모델을 양자화해요. 모든 변형에서 지원되지는 않아요.
serve: 모델(또는 모델 샤드)을 로드하고 라우터의 gRPC 호출을 받아 추론을 수행한 뒤 응답을 반환하는 서버를 시작해요.

serve의 주요 인자는 아래와 같아요.

 Usage: cli.py serve [OPTIONS] MODEL_ID

╭─ Arguments ──────────────────────────────────────────────────────────────────────────────────────────────╮
│ *    model_id      TEXT  [default: None] [required]                                                      │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ Options ────────────────────────────────────────────────────────────────────────────────────────────────╮
│ --revision                                       TEXT                        [default: None]             │
│ --sharded              --no-sharded                                          [default: no-sharded]       │
│ --quantize                                       [bitsandbytes|bitsandbytes  [default: None]             │
│                                                  -nf4|bitsandbytes-fp4|gptq                              │
│                                                  |awq|eetq|exl2|fp8]                                     │
│ --speculate                                      INTEGER                     [default: None]             │
│ --dtype                                          [float16|bfloat16]          [default: None]             │
│ --trust-remote-code    --no-trust-remote-code                                [default:                   │
│                                                                              no-trust-remote-code]       │
│ --uds-path                                       PATH                        [default:                   │
│                                                                              /tmp/text-generation-serve… │
│ --logger-level                                   TEXT                        [default: INFO]             │
│ --json-output          --no-json-output                                      [default: no-json-output]   │
│ --otlp-endpoint                                  TEXT                        [default: None]             │
│ --otlp-service-name                              TEXT                        [default:                   │
│                                                                              text-generation-inference...│
│ --help                                                                       Show this message and exit. │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯

일부 변형은 다른 파라미터를 지원하거나 환경 변수를 통해 더 많은 옵션을 받을 수도 있어요.

호출 흐름 (Call Flow)

두 구성 요소가 초기화되고 가중치가 내려지며 모델 서버가 떠 있으면, 라우터와 모델 서버는 gRPC 호출로 데이터를 주고받아요. 현재 지원되는 스키마는 v2와 v3 두 가지인데, 거의 동일하고 '텍스트·이미지 데이터의 input chunks 지원'과 'paged attention 지원'에서만 차이가 있어요.

서버가 뜬 뒤 라우터와 모델 서버가 교환하는 흐름을 시퀀스 다이어그램으로 보면 아래와 같아요.

sequenceDiagram

    Router->>Model Server: service discovery
    Model Server-->>Router: urls for other shards

    Router->>Model Server: get model info
    Model Server-->>Router: shard info

    Router->>Model Server: health check
    Model Server-->>Router: health OK

    Router->>Model Server: warmup(max_input_tokens, max_batch_prefill_tokens, max_total_tokens, max_batch_size)
    Model Server-->>Router: warmup result

이 과정이 끝나면 라우터는 여러 클라이언트의 생성 요청을 받을 준비가 돼요. 여러 클라이언트가 동시에 스트리밍 생성을 요청하는 실제 흐름은 아래처럼 그려볼 수 있어요.

sequenceDiagram
    participant Client 1
    participant Client 2
    participant Client 3
    participant Router
    participant Model Server

    Client 1->>Router: generate_stream
    Router->>Model Server: prefill(batch1)
    Model Server-->>Router: generations, cached_batch1, timings
    Router-->>Client 1: token 1

    Router->>Model Server: decode(cached_batch1)
    Model Server-->>Router: generations, cached_batch1, timings
    Router-->>Client 1: token 2

    Router->>Model Server: decode(cached_batch1)
    Model Server-->>Router: generations, cached_batch1, timings
    Router-->>Client 1: token 3

    Client 2->>Router: generate_stream
    Router->>Model Server: prefill(batch2)
    Note right of Model Server: This stops previous batch, that is restarted
    Model Server-->>Router: generations, cached_batch2, timings
    Router-->>Client 2: token 1'

    Router->>Model Server: decode(cached_batch1, cached_batch2)
    Model Server-->>Router: generations, cached_batch1, timings
    Router-->>Client 1: token 4
    Router-->>Client 2: token 2'

    Note left of Client 1: Client 1 leaves
    Router->>Model Server: filter_batch(cached_batch1, request_ids_to_keep=batch2)
    Model Server-->>Router: filtered batch

    Router->>Model Server: decode(cached_batch2)
    Model Server-->>Router: generations, cached_batch2, timings
    Router-->>Client 2: token 3'

    Client 3->>Router: generate_stream
    Note right of Model Server: This stops previous batch, that is restarted
    Router->>Model Server: prefill(batch3)
    Note left of Client 1: Client 3 leaves without receiving any batch
    Router->>Model Server: clear_cache(batch3)
    Note right of Model Server: This stops previous batch, that is restarted

    Router->>Model Server: decode(cached_batch3)
    Note right of Model Server: Last token (stopping criteria)
    Model Server-->>Router: generations, cached_batch3, timings
    Router-->>Client 2: token 4'

배치가 계속 이어지는 동안 기존 배치가 잠시 멈췄다가 다시 재개되는 것, 클라이언트가 떠나면 필터링되고 캐시가 정리되는 것까지 한 번에 흐름이 이어져요. 이처럼 라우터가 배치 단위로 prefill과 decode를 오가며 여러 클라이언트의 토큰을 번갈아 내보내는 게 TGI의 핵심 동작 방식이에요.

TGI 아키텍처