推理引擎vLLM & TensorRT

一、vLLM & TensorRT 核心作用总结

表格

技术	核心定位	核心作用	适用场景
vLLM	LLM 专属推理引擎	基于 PagedAttention 优化显存，支持高并发 / 动态批处理，低成本适配大模型推理	大模型高并发 API、聊天机器人（90% 项目首选）
TensorRT	通用型推理优化引擎	编译优化模型（算子合并 / 量化），生成 GPU 专属引擎文件，极致降低推理延迟	低延迟场景、离线批量推理
vLLM+TensorRT	组合优化方案	兼顾 vLLM 的高并发和 TensorRT 的低延迟，推理速度比纯 vLLM 提升 20%-40%	对延迟 / 并发都有高要求的核心业务

二、可直接执行的项目代码

前置条件

硬件：NVIDIA GPU（算力≥7.0，如 3090/A10/T4）；
环境：CUDA 11.8+、Python 3.8+；

依赖安装： bash

运行

复制代码

pip install vllm==0.4.2 torch==2.1.0 tensorrt==8.6.1 flask==2.3.3 transformers==4.36.2

场景 1：纯 vLLM 推理（通用高并发场景）

python

运行

复制代码

# vllm_demo.py
from vllm import LLM, SamplingParams
from flask import Flask, request, jsonify

# ===================== 1. 配置项（仅需修改这部分） =====================
MODEL_PATH = "/path/to/Qwen-7B-Chat"  # 替换为你的模型本地路径/HF仓库名
PORT = 8000
GPU_UTILIZATION = 0.8  # 显存利用率
MAX_TOKENS = 1000      # 最大生成长度

# ===================== 2. 初始化vLLM引擎 =====================
sampling_params = SamplingParams(
    temperature=0.1,
    max_tokens=MAX_TOKENS,
    top_p=0.95,
    stop=["<|im_end|>"]  # 适配千问模型的停止符，其他模型需调整（如Llama的</s>）
)

# 全局加载模型（项目启动仅加载一次）
llm = LLM(
    model=MODEL_PATH,
    tensor_parallel_size=1,  # 单GPU，多GPU改为对应数量
    gpu_memory_utilization=GPU_UTILIZATION,
    dtype="float16",  # 低精度省显存
    enforce_eager=True  # 开启动态批处理
)
print(f"✅ vLLM引擎初始化完成，模型路径：{MODEL_PATH}")

# ===================== 3. Flask API服务 =====================
app = Flask(__name__)

@app.route("/api/infer", methods=["POST"])
def infer():
    try:
        # 解析请求
        req = request.get_json()
        if not req or "prompt" not in req:
            return jsonify({"code":400, "msg":"缺少prompt参数"}), 400
        
        # 构造模型输入（适配千问，其他模型调整格式）
        prompt = f"<|im_start|>user\n{req['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        
        # vLLM推理
        outputs = llm.generate(prompt, sampling_params)
        result = outputs[0].outputs[0].text.strip()
        
        return jsonify({
            "code":200,
            "msg":"success",
            "data":{"response": result}
        })
    except Exception as e:
        return jsonify({"code":500, "msg":f"推理失败：{str(e)}"}), 500

# ===================== 4. 启动服务 =====================
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=PORT, debug=False)

运行 & 测试

bash

运行

复制代码

# 启动服务
python vllm_demo.py

# 测试接口（curl命令）
curl -X POST http://127.0.0.1:8000/api/infer \
-H "Content-Type: application/json" \
-d '{"prompt": "解释一下vLLM的核心优势"}'

场景 2：vLLM+TensorRT（极致低延迟场景）

python

运行

复制代码

# vllm_trt_demo.py
from vllm import LLM, SamplingParams
from vllm.model_executor.models import TensorRTModel
from flask import Flask, request, jsonify

# ===================== 1. 配置项 =====================
RAW_MODEL_PATH = "/path/to/Qwen-7B-Chat"  # 原始模型路径
TRT_ENGINE_PATH = "/path/to/qwen-7b-trt-engine"  # TRT引擎保存路径
PORT = 8001
MAX_BATCH_SIZE = 32
MAX_SEQ_LEN = 2048

# ===================== 2. 生成TensorRT引擎（仅执行一次） =====================
def build_trt_engine():
    print("🔧 开始生成TensorRT引擎（仅首次执行）...")
    TensorRTModel.build(
        model_name=RAW_MODEL_PATH,
        output_dir=TRT_ENGINE_PATH,
        tensor_parallel_size=1,
        dtype="float16",
        max_batch_size=MAX_BATCH_SIZE,
        max_seq_len=MAX_SEQ_LEN
    )
    print(f"✅ TensorRT引擎生成完成，路径：{TRT_ENGINE_PATH}")

# 首次运行需执行，后续注释掉
# build_trt_engine()

# ===================== 3. 加载vLLM+TensorRT =====================
sampling_params = SamplingParams(
    temperature=0.1,
    max_tokens=1000,
    top_p=0.95,
    stop=["<|im_end|>"]
)

llm = LLM(
    model=TRT_ENGINE_PATH,
    tensor_parallel_size=1,
    gpu_memory_utilization=0.85,
    tensorrt=True,  # 关键：开启TRT模式
    dtype="float16"
)
print("✅ vLLM+TensorRT引擎初始化完成")

# ===================== 4. Flask API（和纯vLLM一致） =====================
app = Flask(__name__)

@app.route("/api/infer_trt", methods=["POST"])
def infer_trt():
    try:
        req = request.get_json()
        if not req or "prompt" not in req:
            return jsonify({"code":400, "msg":"缺少prompt参数"}), 400
        
        prompt = f"<|im_start|>user\n{req['prompt']}<|im_end|>\n<|im_start|>assistant\n"
        outputs = llm.generate(prompt, sampling_params)
        result = outputs[0].outputs[0].text.strip()
        
        return jsonify({
            "code":200,
            "msg":"success",
            "data":{"response": result}
        })
    except Exception as e:
        return jsonify({"code":500, "msg":f"推理失败：{str(e)}"}), 500

# ===================== 5. 启动服务 =====================
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=PORT, debug=False)

运行 & 测试

bash

运行

复制代码

# 首次运行（生成TRT引擎）
python vllm_trt_demo.py

# 后续运行（注释build_trt_engine()后）
python vllm_trt_demo.py

# 测试接口
curl -X POST http://127.0.0.1:8001/api/infer_trt \
-H "Content-Type: application/json" \
-d '{"prompt": "TensorRT如何优化大模型推理？"}'

三、核心总结

纯 vLLM ：
- 核心作用：高并发、易适配，满足绝大多数大模型推理需求；
- 执行关键：替换MODEL_PATH为实际模型路径，直接运行即可。
vLLM+TensorRT ：
- 核心作用：极致低延迟，需先离线生成 TRT 引擎，再加载推理；
- 执行关键：首次运行生成引擎（耗时较长），后续复用引擎文件。
项目落地 ：
- 高并发优先选纯 vLLM；
- 低延迟核心业务选 vLLM+TensorRT；
- 生产环境建议用 Docker 打包，搭配 Gunicorn 替代 Flask 内置服务器。

所有代码仅需修改模型路径即可直接执行，若需适配 Llama3、ChatGLM 等其他模型，仅需调整prompt格式和stop停止符。