Text Generation Inference 评测：HuggingFace 的推理服务器

┌─────────────────────────────────────────────┐
│        OpenAI Compatible REST API           │
│   (Streaming / Function Calling / Tools)    │
├─────────────────────────────────────────────┤
│          Router                             │
│   (Load Balancing + Queue Management)       │
├──────────────┬──────────────────────────────┤
│  Flash       │  Continuous Batching         │
│  Attention   │  Scheduler                   │
│  (v2/v3)     │                              │
├──────────────┴──────────────────────────────┤
│           CUDA / ROCm Workers               │
│   (Tensor Parallel + Quantization)          │
└─────────────────────────────────────────────┘

标准 Attention:
  Q, K, V → 写入 HBM → 读取计算 S=QK^T → 写入 HBM
  → 读取 S → 计算 P=softmax(S) → 写入 HBM
  → 读取 P, V → 计算 O=PV → 写入 HBM
  # 3次完整矩阵读写，O(N²) HBM 访问

Flash Attention:
  分块计算，在 SRAM 中完成全部中间计算
  # 仅需 O(N) HBM 访问，计算结果分块写入

# TGI 默认启用 Flash Attention，也可显式指定
docker run --gpus all \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --flash-attention \
  --max-input-tokens 4096 \
  --max-total-tokens 8192

// TGI Scheduler 核心逻辑（Rust 实现，简化）
pub struct Scheduler {
    waiting_queue: VecDeque<Request>,
    running_queue: VecDeque<Request>,
    max_batch_size: usize,
    max_waiting_tokens: usize,
}

impl Scheduler {
    pub fn schedule(&mut self) -> Batch {
        let mut batch = Batch::new();

        // 1. 优先处理 running 请求（decode step）
        for req in &self.running_queue {
            batch.add(req);
            if batch.size() >= self.max_batch_size {
                return batch;
            }
        }

        // 2. 从 waiting 队列拉取新请求（prefill）
        // TGI 策略：prefill 和 decode 不混合，避免 padding 浪费
        while !self.waiting_queue.is_empty() {
            let req = self.waiting_queue.front().unwrap();
            if batch.size() + 1 > self.max_batch_size {
                break;
            }
            let req = self.waiting_queue.pop_front().unwrap();
            batch.add(&req);
        }

        batch
    }
}

# GPTQ 量化模型
docker run --gpus all \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model TheBloke/Llama-2-13B-GPTQ \
  --quantize gptq \
  --max-total-tokens 4096

# AWQ 量化模型
docker run --gpus all \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model TheBloke/Mistral-7B-Instruct-v0.2-AWQ \
  --quantize awq \
  --max-total-tokens 8192

# EETQ 8-bit（推荐，精度损失最小）
docker run --gpus all \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --quantize eetq \
  --max-total-tokens 8192

# TGI 原生 API
curl http://localhost:8080/generate \
  -H "Content-Type: application/json" \
  -d '{
    "inputs": "Explain quantum computing",
    "parameters": {
      "max_new_tokens": 512,
      "temperature": 0.7,
      "top_p": 0.9,
      "repetition_penalty": 1.1,
      "do_sample": true
    },
    "stream": true
  }'

# OpenAI 兼容端点（/v1/chat/completions）
curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "messages": [
      {"role": "user", "content": "Hello"}
    ],
    "stream": true,
    "max_tokens": 512
  }'

import requests

response = requests.post("http://localhost:8080/v1/chat/completions", json={
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "messages": [
        {"role": "user", "content": "What's the weather in Tokyo?"}
    ],
    "tools": [{
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get weather for a city",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string"}
                }
            }
        }
    }],
    "tool_choice": "auto"
})

# TGI：Docker 优先，单命令启动
docker run --gpus all -p 8080:80 \
  -v $PWD/data:/data \
  ghcr.io/huggingface/text-generation-inference:latest \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --num-shard 1 \
  --max-total-tokens 8192

# vLLM：Python 优先，更灵活
pip install vllm
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --tensor-parallel-size 1 \
  --max-model-len 8192

version: "3.9"
services:
  tgi:
    image: ghcr.io/huggingface/text-generation-inference:latest
    container_name: tgi
    runtime: nvidia
    environment:
      - HF_TOKEN=${HF_TOKEN}
    volumes:
      - tgi_data:/data
    ports:
      - "8080:80"
    command:
      - --model-id=meta-llama/Llama-3.1-70B-Instruct
      - --num-shard=4
      - --quantize=eetq
      - --max-input-tokens=4096
      - --max-total-tokens=32768
      - --max-batch-size=128
      - --flash-attention
      - --trust-remote-code
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:80/health"]
      interval: 30s
      timeout: 10s
      retries: 5
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 4
              capabilities: [gpu]

volumes:
  tgi_data:

# Prometheus 指标端点
curl http://localhost:8080/metrics | grep tgi

# 关键指标
tgi_request_count               # 请求总数
tgi_request_duration_seconds    # 请求耗时
tgi_request_mean_time_seconds   # 平均耗时
tgi_batch_current_size          # 当前批大小
tgi_queue_size                  # 队列长度