zxb000/tensorrt_llm

镜像概述

本镜像基于NVIDIA TensorRT-LLM 2025年4月28日最新提交版本（commit id: ad15e45f07b0907）编译，编译时指定CUDA_ARCHS="90-real;100-real"，支持Hopper（如H200）和Blackwell（如B200）架构，适用于DeepSeek-V3/R1等大语言模型的高效推理部署。

核心功能

1. 支持FP8/FP4精度模型推理，优化GPU资源利用率
2. 提供trtllm-serve服务，快速搭建模型API接口
3. 兼容多GPU并行（TP/EP/PP），满足大模型部署需求
4. 支持CUDA Graph优化，提升推理性能和吞吐量

使用场景

1. 大语言模型生产环境的推理部署
2. 研究人员测试DeepSeek-V3/R1模型的性能
3. 搭建AI服务接口，提供自然语言生成、问答等能力

配置说明

硬件要求

根据模型精度和参数规模选择合适的GPU配置：

部署示例

启动容器

bash
docker run -itd --privileged --ipc=host --gpus all --network=host --ulimit memlock=-1 --ulimit stack=67108864 -v /deepseek:/root/.cache/huggingface --name tensor docker.xuanyuan.run/zxb000/tensorrt_llm:20250428-ad15e45-nv-sm190 bash

快速开始推理

进入容器后执行：

bash
cd examples/pytorch/
python3 quickstart_advanced.py --model_dir /root/.cache/huggingface/DeepSeek-R1-FP4 --tp_size 8

启动API服务

bash
cd /root/.cache/huggingface
# 创建配置文件
cat >./extra-llm-api-config.yml <<EOF
pytorch_backend_config:
    use_cuda_graph: true
    cuda_graph_padding_enabled: true
    cuda_graph_batch_sizes: [1,2,4,8,16,32,64,128,256,384]
    print_iter_log: true
    enable_overlap_scheduler: true
enable_attention_dp: true
EOF
# 启动服务
trtllm-serve \
  /root/.cache/huggingface/DeepSeek-R1-FP4 \
  --host localhost \
  --port 8000 \
  --backend pytorch \
  --max_batch_size 161 \
  --max_num_tokens 8192 \
  --tp_size 8 \
  --ep_size 8 \
  --pp_size 1 \
  --kv_cache_free_gpu_memory_fraction 0.7 \
  --extra_llm_api_options ./extra-llm-api-config.yml

关键参数说明

• --kv_cache_free_gpu_memory_fraction: 分配模型权重和缓冲区后，为KV缓存保留的可用GPU内存比例（建议0.7左右）
• --tp_size: 张量并行的GPU数量
• --ep_size: 专家并行的GPU数量
• --max_batch_size: 服务支持的最大批量大小
• --max_num_tokens: 模型生成的最大token数