Skip to content

TensorRT-LLM 与 TGI

TensorRT-LLM 和 TGI (Text Generation Inference) 是两个重要的 LLM 推理引擎。TensorRT-LLM 由 NVIDIA 开发，追求极致的推理性能；TGI 由 Hugging Face 开发，强调易用性和生态集成。

关于 vLLM 推理引擎，参见 vLLM。

TensorRT-LLM

概述

TensorRT-LLM 是 NVIDIA 基于 TensorRT 深度学习推理优化器构建的 LLM 专用推理框架。它将 LLM 编译为高度优化的 TensorRT 引擎，充分利用 NVIDIA GPU 的硬件特性。

核心优化技术

图优化与算子融合 (Kernel Fusion)

TensorRT-LLM 在编译阶段对计算图进行深度优化：

  • 层融合：将多个连续操作（如 LayerNorm + Linear + Activation）融合为单个 CUDA Kernel，减少显存访问和 Kernel 启动开销
  • 常量折叠：在编译时预计算静态表达式
  • 精度混合：自动决定哪些层使用 FP16、哪些使用 FP8/INT8，在精度和速度间取平衡

自定义 CUDA Kernel

NVIDIA 为 Transformer 的关键操作编写了高度优化的 CUDA Kernel：

  • Flash Attention 的定制化实现（针对不同 GPU 架构）
  • 融合的 MHA/GQA/MQA Kernel
  • 优化的 GEMM（矩阵乘法）Kernel，支持 FP8 Tensor Core

张量并行与流水线并行

原生支持多种并行策略：

  • Tensor Parallelism (TP)：将模型的权重矩阵切分到多张 GPU 上
  • Pipeline Parallelism (PP)：将模型的不同层分配到不同 GPU 上
  • 支持 TP + PP 的混合并行

In-flight Batching

类似于 vLLM 的 Continuous Batching，TensorRT-LLM 实现了 In-flight Batching：

  • 在每次迭代时动态加入/移除请求
  • 支持不同请求使用不同的采样参数
  • 与 KV Cache 管理紧密集成

使用流程

TensorRT-LLM 的使用分为编译和部署两个阶段：

# 1. 模型转换：将 HuggingFace 模型转换为 TensorRT-LLM checkpoint
python convert_checkpoint.py \
    --model_dir meta-llama/Llama-2-7b-hf \
    --output_dir ./checkpoint \
    --dtype float16

# 2. 引擎编译：将 checkpoint 编译为 TensorRT 引擎
trtllm-build \
    --checkpoint_dir ./checkpoint \
    --output_dir ./engine \
    --gemm_plugin float16 \
    --max_batch_size 64 \
    --max_input_len 2048 \
    --max_seq_len 4096

# 3. 部署：使用 Triton Server 提供服务
python launch_triton_server.py --model_repo ./triton_repo

注意： 引擎在编译时需要指定最大 batch size 和序列长度，编译后不可更改。不同硬件需要分别编译。

适用场景

  • 生产环境的极致性能：延迟要求严格的在线服务
  • NVIDIA GPU 集群：充分利用 Tensor Core、NVLink 等硬件特性
  • 大规模部署：通过 Triton Inference Server 实现弹性伸缩

TGI (Text Generation Inference)

概述

TGI 是 Hugging Face 开发的 LLM 推理与服务工具，用 Rust 和 Python 编写。它以 Docker 容器的形式提供开箱即用的部署体验，与 Hugging Face 生态深度集成。

核心特性

推理优化：

  • Continuous Batching：动态批处理
  • Flash Attention 2：高效 Attention 计算
  • Tensor Parallelism：多 GPU 推理
  • 量化支持：GPTQ、AWQ、BitsAndBytes、EETQ、FP8

服务化特性：

  • 流式输出（Server-Sent Events）
  • Token 级流式传输
  • 并发请求管理与队列
  • 结构化输出（JSON mode / Grammar）
  • OpenAI 兼容的 API 接口

运维友好：

  • Docker 一键部署
  • 内置 Prometheus 监控指标
  • 支持分布式追踪（OpenTelemetry）
  • HuggingFace Hub 直接拉取模型

使用方式

Docker 部署（推荐）：

# 最简单的部署方式
docker run --gpus all --shm-size 1g -p 8080:80 \
    ghcr.io/huggingface/text-generation-inference:latest \
    --model-id meta-llama/Llama-2-7b-chat-hf \
    --max-input-tokens 2048 \
    --max-total-tokens 4096

# 多GPU张量并行
docker run --gpus all --shm-size 1g -p 8080:80 \
    ghcr.io/huggingface/text-generation-inference:latest \
    --model-id meta-llama/Llama-2-70b-chat-hf \
    --num-shard 4 \
    --quantize gptq

客户端调用：

# 使用 huggingface_hub 客户端
from huggingface_hub import InferenceClient

client = InferenceClient("http://localhost:8080")

# 流式生成
for token in client.text_generation(
    "What is deep learning?",
    max_new_tokens=100,
    stream=True
):
    print(token, end="")

# 或使用 OpenAI 兼容 API
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8080/v1", api_key="-")
response = client.chat.completions.create(
    model="tgi",
    messages=[{"role": "user", "content": "Hello!"}]
)

适用场景

  • 快速原型与验证：Docker 一键启动，无需编译
  • HuggingFace 生态集成：直接使用 Hub 上的模型
  • 中等规模服务：满足大多数在线服务的性能需求

三大推理引擎对比

维度 TensorRT-LLM TGI vLLM
开发方 NVIDIA Hugging Face UC Berkeley
语言 C++ / Python Rust / Python Python
核心优势 极致性能（Kernel 级优化） 易用性与生态 高吞吐（PagedAttention）
部署复杂度 高（需编译引擎） 低（Docker） 低（pip install）
量化支持 FP8, INT8, INT4 (原生) GPTQ, AWQ, BnB, FP8 GPTQ, AWQ, FP8
并行策略 TP + PP TP TP + PP
延迟 最低 中等 中等
吞吐量 最高 高 很高
硬件要求 NVIDIA GPU（特定架构） NVIDIA GPU NVIDIA / AMD GPU
模型支持 需逐个适配 HF Hub 大量支持 HF Hub 大量支持
学习曲线 陡峭 平缓 平缓

选型决策

需要极致延迟/吞吐？
├── 是 → 有工程资源投入编译优化？
│   ├── 是 → TensorRT-LLM
│   └── 否 → vLLM
└── 否 → 需要 HuggingFace 深度集成？
    ├── 是 → TGI
    └── 否 → vLLM（通用推荐）

参考

评论 #