VLLM和TensorRT

前提

使用CUDA

Python 代码 (vLLM) -> 调用 PyTorch -> PyTorch 调用 CUDA 库 -> CUDA 驱动 指挥 GPU 硬件 进行计算。

介绍

• vLLM: 基于 Python 开发，主打高吞吐量和易用性。它通过 PagedAttention 技术优化显存管理，非常适合服务端部署、多用户并发场景。它的 API 兼容 OpenAI 格式，也提供直接的 Python 类调用。
• TensorRT-LLM (TRT-LLM): NVIDIA 出品，主打极致低延迟和硬件级优化。它需要将模型预先编译成特定的 Engine 文件（类似 C++ 编译过程），然后在 Python 或 C++ 中加载运行。适合对延迟极其敏感或显存极度受限的生产环境。

功能和作用

• vLLM：核心在于“显存管理”。它通过 PagedAttention 技术，像操作系统管理内存一样管理显存，消除了显存碎片，从而能在同一块显卡上塞进更多的并发请求。
• TensorRT-LLM：核心在于“计算调度”与“内核融合”。它通过 In-flight Batching 技术在 GPU 计算间隙插入新请求，并利用高度优化的 CUDA 内核（Kernel）榨干 GPU 算力。

下面详细拆解它们的原理、实现方式和作用。

一、vLLM：PagedAttention (分页注意力机制)

1. 痛点：传统 LLM 的显存浪费

在 vLLM 出现之前，传统的 LLM 推理（如 Hugging Face Transformers）在处理并发时有一个致命弱点：显存预分配。

• 场景：假设显存够存 1000 个 token。
• 问题：如果你启动了 10 个并发请求，系统通常会为每个请求平均分配显存（例如每个预留 100 个 token 的空间）。
  • 如果请求 A 只生成了 10 个 token 就结束了，它占用的 90 个 token 空间就浪费了（无法给其他请求用）。
  • 如果请求 B 需要生成 200 个 token，但因为预分配限制，它可能直接报错 OOM (Out Of Memory)，即使此时显存里还有很多空闲碎片。
• 结果：显存利用率通常只有 20%-40%，并发度上不去。

2. 原理：PagedAttention (借鉴操作系统的虚拟内存)

vLLM 将操作系统中虚拟内存分页 (Virtual Memory Paging) 的思想引入了 GPU 显存管理。

• 非连续存储：KV Cache（注意力机制需要的键值对缓存）不再需要连续的显存空间。它被切分成固定大小的块 (Blocks)（例如每个块存 16 个 token）。
• 页表映射：vLLM 维护一个页表 (Block Table)，记录每个逻辑上的 token 块对应物理显存中的哪个块。
  • 逻辑块 0 -> 物理块 5
  • 逻辑块 1 -> 物理块 2
  • 逻辑块 2 -> 物理块 9
• 动态分配：
  • 当请求生成新 token 时，vLLM 只需从空闲池子里拿一个物理块挂上去，按需分配。
  • 当请求结束时，物理块立即释放回池子，供其他请求使用。

3. 如何实现并发？

由于消除了碎片，vLLM 可以将显存利用率提升到 95% 以上。

• 效果：原本只能同时跑 10 个请求的显存，现在可以跑 50 个甚至 100 个请求。
• 调度器：vLLM 有一个调度器，每一轮迭代前，它会检查显存池里还有多少空闲块。如果有，就从等待队列中把新的请求“塞”进去（加入当前 Batch），直到显存填满。

4. 作用

• 极高吞吐量：单位时间内处理的请求数大幅提升。
• 抗长尾能力：不怕某个请求特别长，也不怕某个请求特别短，显存永远被充分利用。
• 零浪费：彻底解决了内部碎片化问题。

二、TensorRT-LLM：In-flight Batching (飞行中批处理)

1. 痛点：传统 Batch 的“木桶效应”

传统的批处理（Static Batching）是将 N 个请求打包一起送入 GPU 计算。

• 场景：Batch 里有 4 个请求。
  • 请求 A, B, C 很短，第 10 步就生成完了（遇到 EOS）。
  • 请求 D 很长，需要生成 100 步。
• 问题：在传统模式下，GPU 必须等待所有请求完成当前步。当 A, B, C 完成后，它们占用的计算资源必须空转等待 D 完成，或者整个 Batch 必须解散，重新组织新的 Batch。这导致 GPU 在很多时间片里是在做无用功。

2. 原理：In-flight Batching (Continuous Batching)

TRT-LLM (以及 vLLM 也采用了此技术) 允许在同一个迭代步骤中动态地移除已完成的请求，并插入新的请求。

• 细粒度调度：GPU 不是以“整个 Batch”为单位运行，而是以“迭代步 (Iteration)”为单位。
• 即时插拔：
  • 在第 10 步计算结束时，发现请求 A, B, C 结束了。
  • 立刻将它们从显存和计算队列中移除。
  • 立刻从等待队列中抓取新的请求 D, E, F 填补空缺。
  • 请求 D (原来的长请求) 继续在第 11 步计算，完全不受影响。
• 全程无停顿：GPU 的计算单元永远处于满载状态，没有因为等待慢请求而产生的气泡 (Bubble)。

3. TRT-LLM 的独门绝技：内核融合与量化

除了调度，TRT-LLM 还通过底层优化提升并发时的单请求速度，从而间接提升并发能力：

• Kernel Fusion (内核融合)：将多个小的 CUDA 操作（如 Add, Norm, Activation）合并成一个大操作，减少 GPU 启动内核的开销和显存读写次数。
• FP8/INT8 量化：在编译阶段将模型权重转换为低精度，显存占用减半，计算速度翻倍，意味着同样的显存能容纳更大的 Batch Size。
• Multi-Block Attention：针对长上下文优化的注意力核，减少显存访问延迟。

4. 作用

• 极低延迟：特别是对于首字延迟 (TTFT) 和 单 token 延迟，优化到了极致。
• 稳定高吞吐：即使请求长度差异巨大，GPU 利用率也能保持在高位。
• 硬件压榨：充分利用 NVIDIA GPU 的 Tensor Core 特性。

三、总结对比：它们是如何协作的？

其实，现代高性能推理引擎往往是两者结合的。vLLM 也实现了 In-flight Batching，TRT-LLM 也有自己的显存管理机制。但侧重点不同：

表格

四、实际工作流程示例

假设现在有 100 个用户同时发起请求：

结论：

• 如果你追求部署简单且并发量大，vLLM 的 PagedAttention 是首选，因为它能让你用更少的卡抗住更多的人。
• 如果你追求极致性能且有能力做复杂的工程化构建，TensorRT-LLM 通过底层的计算优化和调度，能提供比 vLLM 更高的上限（尤其是在延迟方面）。
• 趋势：两者正在融合。vLLM 正在集成更多 TRT-LLM 的优化内核，而 TRT-LLM 也在改进其显存管理策略。

使用

VLLM

核心区别一览

表格

🛠️ 第一步：环境安装 (通用)

无论选哪种方案，先执行以下命令安装依赖：

# 安装 vLLM 核心库 (会自动安装 torch 等依赖)
pip install vllm

# 安装 HTTP 客户端库 (仅方案二需要，但建议都装上)
pip install aiohttp tqdm

硬件要求：需要 NVIDIA 显卡 (显存建议 8GB 以上运行 7B 模型，若显存不足请将代码中的模型改为 Qwen/Qwen2.5-1.5B-Instruct)。

方案一：离线模式 (Offline Mode)

特点：单文件运行，无需启动服务器，无网络开销，速度最快。
原理：直接在 Python 进程中加载模型，通过 llm.generate() 函数并发处理列表。

📄 文件：vllm_offline_qwen.py

python

"""
方案一：vLLM 离线模式 (Offline Mode)
功能：直接在 Python 代码中加载千问模型，批量并发处理多个问题。
无需启动服务器，无需 HTTP 请求。
"""

import time
from vllm import LLM, SamplingParams
from tqdm import tqdm # 用于显示进度条

# ================= 配置区域 =================
# 模型名称：通义千问 7B 指令版
# 如果显存不足 (<8GB)，请改为 "Qwen/Qwen2.5-1.5B-Instruct"
MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"

# 模拟 50 个并发问题 (实际使用中可以是你的业务数据列表)
PROMPTS = [
f"问题 {i}: 请用一句话解释什么是量子纠缠？" for i in range(25)
] + [
f"问题 {i}: 写一个 Python 函数计算斐波那契数列。" for i in range(25)
]

def main():
print(f"🚀 正在初始化 vLLM 引擎并加载模型：{MODEL_NAME} …")
print("⏳ 首次运行会自动下载模型，请耐心等待…")

start_time = time.time()

# 1. 初始化 LLM 引擎
# tensor_parallel_size=1: 使用 1 张显卡。多卡可改为 2, 4, 8
# gpu_memory_utilization=0.9: 占用 90% 显存，留 10% 给系统防止 OOM
# enforce_eager=False: 启用 CUDA Graph 加速，提升推理速度
llm = LLM(
model=MODEL_NAME,
tensor_parallel_size=1,
gpu_memory_utilization=0.9,
enforce_eager=False,
disable_log_stats=True # 关闭 vLLM 内部日志，保持输出整洁
)

load_time = time.time() – start_time
print(f"✅ 模型加载完成！耗时：{load_time:.2f} 秒")

# 2. 设置采样参数
sampling_params = SamplingParams(
temperature=0.7, # 创造性：0-1，越高越随机
top_p=0.8, # 核采样概率
max_tokens=256 # 限制最大生成长度，防止测试时间过长
)

print(f"📝 开始处理 {len(PROMPTS)} 个并发请求…")
process_start = time.time()

# 3. 核心并发调用
# vLLM 的 generate 方法天然支持列表输入
# 它会自动将列表中的所有 Prompt 打包成一个 Batch，利用 PagedAttention 并行处理
# 这里不需要写任何多线程代码，vLLM 内部已经高度优化
outputs = llm.generate(PROMPTS, sampling_params, use_tqdm=True)

process_time = time.time() – process_start

# 4. 结果统计与展示
total_tokens = sum(len(out.outputs[0].token_ids) for out in outputs)
speed = total_tokens / process_time

print("\\n" + "="*50)
print("📊 测试结果统计")
print("="*50)
print(f"总请求数：{len(outputs)}")
print(f"处理耗时：{process_time:.2f} 秒")
print(f"生成总 Token 数：{total_tokens}")
print(f"平均生成速度：{speed:.2f} Tokens/秒")
print("="*50)

# 展示前 3 个结果示例
print("\\n💡 部分结果示例：")
for i, output in enumerate(outputs[:3]):
prompt = output.prompt
result = output.outputs[0].text
print(f"[{i+1}] 问：{prompt}")
print(f" 答：{result.strip()[:100]}…") # 只打印前100字
print("-" * 30)

if __name__ == "__main__":
main()

👉 如何运行：
直接运行即可：

python vllm_offline_qwen.py

方案二：在线模式 (Server Mode)

特点：分离架构。服务端一直运行监听端口，客户端通过 HTTP 发送请求。
适用：你需要部署一个服务，让网页、App 或其他程序调用。

此方案需要两个步骤（或两个终端窗口）。

第 1 步：启动服务端

📄 文件：start_vllm_server.py (或者直接在命令行运行)

"""
方案二 (步骤 1)：启动 vLLM API 服务端
功能：加载模型并启动 HTTP 服务器，监听 8000 端口。
"""

import os
import sys
import subprocess
import time

def start_server():
# 模型名称 (需与客户端一致)
MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
PORT = 8000

print(f"🚀 正在启动 vLLM API 服务器…")
print(f"📦 模型：{MODEL_NAME}")
print(f"🔌 端口：http://localhost:{PORT}")

# 构建启动命令
# –disable-log-requests: 减少日志刷屏
cmd = [
sys.executable, "-m", "vllm.entrypoints.openai.api_server",
"–model", MODEL_NAME,
"–port", str(PORT),
"–host", "0.0.0.0",
"–tensor-parallel-size", "1",
"–gpu-memory-utilization", "0.9",
"–disable-log-requests"
]

# 启动子进程
# 注意：在实际生产中，通常直接用命令行运行，而不是用 python 脚本启动
# 这里为了演示完整性，用 subprocess 启动
process = subprocess.Popen(cmd)

print("⏳ 等待服务就绪 (约 10-20 秒)…")
time.sleep(15) # 简单等待，实际应检查健康接口

print("✅ 服务已启动！请保持此窗口运行。")
print("👉 现在请打开另一个终端，运行 'start_vllm_client.py' 进行测试。")

# 阻塞主进程，保持服务运行
try:
process.wait()
except KeyboardInterrupt:
print("\\n🛑 正在停止服务…")
process.terminate()

if __name__ == "__main__":
start_server()

👉 如何运行服务端：
在 终端 1 运行：

python start_vllm_server.py

(看到 Uvicorn running on http://0.0.0.0:8000 表示成功)

第 2 步：运行客户端

📄 文件：start_vllm_client.py

python

"""
方案二 (步骤 2)：vLLM 客户端并发测试
功能：通过 HTTP 请求向 vLLM 服务端发送大量并发问题。
"""

import asyncio
import aiohttp
import time
from tqdm import tqdm

# ================= 配置区域 =================
# 服务端地址 (必须与启动服务时的端口一致)
API_URL = "http://localhost:8000/v1/chat/completions"
MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"

# 构造 50 个测试问题
PROMPTS = [f"问题 {i}: 今天天气不错，适合做什么？" for i in range(50)]

async def send_request(session, prompt, req_id):
"""发送单个 HTTP 请求"""
payload = {
"model": MODEL_NAME,
"messages": [{"role": "user", "content": prompt}],
"temperature": 0.7,
"max_tokens": 100
}

start = time.time()
try:
async with session.post(API_URL, json=payload, timeout=aiohttp.ClientTimeout(total=60)) as resp:
if resp.status == 200:
data = await resp.json()
content = data["choices"][0]["message"]["content"]
latency = time.time() – start
return {"id": req_id, "success": True, "latency": latency, "content": content[:50]}
else:
return {"id": req_id, "success": False, "error": f"Status {resp.status}"}
except Exception as e:
return {"id": req_id, "success": False, "error": str(e)}

async def main():
print(f"🚀 开始向 {API_URL} 发送 {len(PROMPTS)} 个并发请求…")

# 创建 TCP 连接器，限制最大连接数以防本地端口耗尽
connector = aiohttp.TCPConnector(limit=len(PROMPTS))

start_time = time.time()

async with aiohttp.ClientSession(connector=connector) as session:
# 创建所有任务
tasks = [send_request(session, p, i) for i, p in enumerate(PROMPTS)]

# 并发执行所有任务，并显示进度条
results = await tqdm.gather(*tasks, desc="Processing")

total_time = time.time() – start_time

# 统计结果
success_count = sum(1 for r in results if r["success"])
avg_latency = sum(r["latency"] for r in results if r["success"]) / success_count if success_count > 0 else 0

print("\\n" + "="*50)
print("📊 客户端测试结果")
print("="*50)
print(f"总请求数：{len(results)}")
print(f"成功数：{success_count}")
print(f"总耗时：{total_time:.2f} 秒")
print(f"平均延迟：{avg_latency:.3f} 秒/请求")
print(f"吞吐量：{success_count/total_time:.2f} 请求/秒")
print("="*50)

# 展示一个成功结果
for r in results:
if r["success"]:
print(f"\\n💡 示例回答 (ID:{r['id']}): {r['content']}…")
break

if __name__ == "__main__":
# 运行异步主函数
asyncio.run(main())

👉 如何运行客户端：
在 终端 2 (保持终端 1 的服务端运行) 运行：

python start_vllm_client.py

💡 总结与建议

TensorRT

⚠️ 重要前提：

• 推荐启动命令 (如果你用 Docker)：
  bash docker run –gpus all -it –rm -v $(pwd):/workspace nvcr.io/nvidia/tensorrt_llm:0.14.0-py3
• 进入容器后，执行下方的 install_deps.sh 安装额外依赖。

🚀 全自动一键运行脚本 (run_trt_qwen.py)

将以下代码保存为 run_trt_qwen.py。
逻辑：脚本会自动检测是否有编译好的引擎。如果没有，它会自动执行编译过程；如果有，直接进行并发测试。

python

import os
import sys
import subprocess
import time
import argparse
from pathlib import Path

# ================= 配置区域 =================
# 模型选择：根据显存大小选择
# 显存 < 8GB: 使用 "Qwen/Qwen2.5-1.5B-Instruct"
# 显存 >= 12GB: 使用 "Qwen/Qwen2.5-7B-Instruct"
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

# 输出目录 (编译后的引擎和中间文件存放处)
OUTPUT_DIR = "./trt_engine_output"
# 工作目录 (用于构建脚本)
WORKSPACE_DIR = "./trt_build_workspace"

# 并发测试配置
NUM_REQUESTS = 50 # 并发请求数量
MAX_OUTPUT_LEN = 100 # 每个回答最大长度

def print_header(text):
print("\\n" + "="*60)
print(f" {text}")
print("="*60 + "\\n")

def run_command(cmd, cwd=None):
"""运行 shell 命令并实时打印输出"""
print(f"🔧 执行命令: {' '.join(cmd)}")
process = subprocess.Popen(cmd, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
for line in process.stdout:
print(line, end='')
process.wait()
if process.returncode != 0:
raise Exception(f"命令执行失败: {' '.join(cmd)}")
return True

def build_engine():
"""步骤 1: 下载模型并编译 TensorRT 引擎"""
print_header("步骤 1: 构建 TensorRT 引擎 (首次运行需耗时 5-10 分钟)")

if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
if not os.path.exists(WORKSPACE_DIR):
os.makedirs(WORKSPACE_DIR)

# 1. 准备构建脚本 (使用 TRT-LLM 自带的 examples 脚本逻辑)
# 注意：这里调用的是 trtllm-build 命令，这是 TRT-LLM 的核心编译工具
# 为了简化，我们直接使用 Python API 进行构建，这样更灵活且无需写复杂的 shell 脚本

build_script = f"""
import tensorrt_llm
from tensorrt_llm.models import QWenForCausalLM
from tensorrt_llm.builders import build_engine
from transformers import AutoTokenizer
import torch

model_name = "{MODEL_NAME}"
output_dir = "{OUTPUT_DIR}"

print(f"正在加载模型：{{model_name}} …")
# 加载 HF 模型
model = QWenForCausalLM.from_hugging_face(model_name, dtype="float16")

print("正在编译引擎 (这可能比较慢，请耐心等待)…")
# 构建配置
config = tensorrt_llm.BuildConfig(
max_batch_size=64, # 支持的最大并发数
max_input_len=1024,
max_output_len={MAX_OUTPUT_LEN},
max_beam_width=1,
use_gpt_attention_plugin="float16",
use_gemm_plugin="float16",
enable_context_fmha=True # 启用 Flash Attention 加速
)

# 编译并保存
build_engine(model, config).save_to_directory(output_dir)
print("✅ 引擎构建完成！")
"""

# 将构建逻辑写入临时文件并运行
temp_build_file = "temp_build_engine.py"
with open(temp_build_file, "w") as f:
f.write(build_script)

try:
run_command([sys.executable, temp_build_file])
print("✅ 引擎构建成功！文件位于:", OUTPUT_DIR)
except Exception as e:
print(f"❌ 引擎构建失败：{e}")
print("💡 提示：请确保已正确安装 tensorrt_llm 且在 NVIDIA Docker 环境中运行。")
sys.exit(1)
finally:
if os.path.exists(temp_build_file):
os.remove(temp_build_file)

def run_concurrent_inference():
"""步骤 2: 加载引擎并进行高并发推理"""
print_header("步骤 2: 执行高并发推理测试")

try:
from tensorrt_llm.runtime import ModelRunner
from transformers import AutoTokenizer
import torch
except ImportError:
print("❌ 无法导入 tensorrt_llm。请确认安装是否正确。")
sys.exit(1)

# 1. 加载 Tokenizer
print(f"正在加载 Tokenizer: {MODEL_NAME} …")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, pad_token='<|extra_0|>')

# 2. 初始化 Runner (加载编译好的引擎)
print(f"正在加载 TensorRT 引擎从：{OUTPUT_DIR} …")
runner = ModelRunner.from_dir(engine_dir=OUTPUT_DIR, rank=0)

# 3. 准备并发数据
prompts = [f"问题 {i}: 请用简短的语言解释人工智能。" for i in range(NUM_REQUESTS)]

# 预处理输入 (Tokenize)
input_ids_list = []
for prompt in prompts:
# 应用 Chat Template (千问模型需要)
messages = [{"role": "user", "content": prompt}]
text_input = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
input_ids = tokenizer.encode(text_input, return_tensors="pt").squeeze(0).cuda()
input_ids_list.append(input_ids)

print(f"🚀 开始处理 {NUM_REQUESTS} 个并发请求…")
start_time = time.time()

# 4. 批量推理 (Batch Inference)
# TRT-LLM 的 generate 方法天然支持 Batch 输入
# input_ids 需要填充到相同长度 (pad)
max_len = max(len(ids) for ids in input_ids_list)
padded_input_ids = torch.full((len(input_ids_list), max_len), tokenizer.pad_token_id, dtype=torch.int32, device='cuda')
input_lengths = torch.tensor([len(ids) for ids in input_ids_list], dtype=torch.int32, device='cuda')

for i, ids in enumerate(input_ids_list):
padded_input_ids[i, -len(ids):] = ids

with torch.no_grad():
# 执行生成
outputs = runner.generate(
inputs=padded_input_ids,
input_lengths=input_lengths,
max_new_tokens=MAX_OUTPUT_LEN,
end_id=tokenizer.eos_token_id,
pad_id=tokenizer.pad_token_id,
temperature=0.7,
top_k=1,
streaming=False
)

end_time = time.time()
total_time = end_time – start_time

# 5. 解析结果并统计
# outputs shape: [batch, beam, seq_len]
generated_ids = outputs[:, 0, :]

success_count = 0
total_output_tokens = 0

print("\\n💡 部分结果示例:")
for i in range(min(3, NUM_REQUESTS)):
output_ids = generated_ids[i]
# 解码 (跳过输入部分)
input_len = input_lengths[i].item()
text = tokenizer.decode(output_ids[input_len:], skip_special_tokens=True)
print(f"[{i}] {text[:80]}…")
total_output_tokens += len(output_ids) – input_len
success_count += 1

# 性能统计
tokens_per_sec = total_output_tokens / total_time
req_per_sec = success_count / total_time

print("\\n" + "="*60)
print("📊 TensorRT-LLM 并发测试结果")
print("="*60)
print(f"模型：{MODEL_NAME}")
print(f"并发请求数：{NUM_REQUESTS}")
print(f"总耗时：{total_time:.2f} 秒")
print(f"总生成 Token 数：{total_output_tokens}")
print(f"吞吐量 (Tokens/s): {tokens_per_sec:.2f}")
print(f"请求处理速度 (Req/s): {req_per_sec:.2f}")
print("="*60)

runner.close()

def main():
parser = argparse.ArgumentParser(description="TensorRT-LLM Qwen 并发测试")
parser.add_argument("–rebuild", action="store_true", help="强制重新编译引擎")
args = parser.parse_args()

engine_exists = os.path.exists(os.path.join(OUTPUT_DIR, "rank0.engine")) or os.path.exists(OUTPUT_DIR)

if args.rebuild or not engine_exists:
build_engine()

run_concurrent_inference()

if __name__ == "__main__":
main()

🛠️ 如何真正“直接拿到就能用”？

由于 TensorRT-LLM 的依赖极其复杂（需要特定版本的 CUDA, TensorRT, PyTorch, Transformers），在普通 Linux/Windows 上直接 pip install 成功率极低。

最稳妥的“一键运行”方案是使用 Docker。

1. 准备环境 (只需做一次)

确保你安装了 NVIDIA Docker Toolkit。然后运行以下命令拉取官方镜像并启动容器：

# 拉取最新的 TRT-LLM 镜像 (约 10GB)
docker pull nvcr.io/nvidia/tensorrt_llm:0.14.0-py3

# 启动容器，挂载当前目录以便代码互通
docker run –gpus all -it –rm \\
-v $(pwd):/workspace \\
-w /workspace \\
nvcr.io/nvidia/tensorrt_llm:0.14.0-py3 \\
bash

2. 在容器内安装额外依赖

进入容器后（你会看到命令行提示符变了），执行：

# 安装 huggingface hub 和 tqdm (镜像里可能缺这些)
pip install huggingface_hub tqdm aiohttp

# 登录 HuggingFace (如果需要下载千问模型，有时需要 token，公开模型可跳过)
# huggingface-cli login

3. 运行脚本

将上面的 run_trt_qwen.py 保存到当前目录，然后在容器内运行：

# 第一次运行会自动下载模型并编译 (耗时 5-10 分钟)
python run_trt_qwen.py

# 第二次运行 (引擎已存在，直接秒开并发测试)
python run_trt_qwen.py

💡 代码核心原理解析

⚠️ 常见问题排查

• OOM (显存不足)：
  • 修改代码中的 MODEL_NAME 为 "Qwen/Qwen2.5-1.5B-Instruct"。
  • 或者在 BuildConfig 中减小 max_batch_size。
• 下载模型慢：
  • 设置镜像源：export HF_ENDPOINT=https://hf-mirror.com (在中国大陆地区)。
• 找不到 tensorrt_llm 模块：
  • 务必使用官方 Docker 镜像。本地手动安装极易出错。

混合使用

⚠️ 重要前置条件

由于要在一个脚本里同时加载两个重型引擎，对显存要求极高：

• 你需要一张大显存显卡（如 A100 80G, H100, 或 双卡 4090）。
• 或者，你需要加载不同大小的模型（例如：TRT 跑 1.5B 小模型处理简单问题，vLLM 跑 7B 大模型处理复杂问题），以节省显存。

下面的代码为了让你能直接运行测试逻辑，做了以下处理：

🚀 真实混合调用代码 (hybrid_direct_call.py)

python

import time
import random
import asyncio
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from enum import Enum

# ================= 配置区域 =================

# 模型配置
# 建议：TRT 跑小模型 (快/省显存)，vLLM 跑大模型 (强/长上下文)
VLLM_MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
TRT_MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" # 假设 TRT 编译的是这个小模型

# 路由策略阈值
MAX_LENGTH_FOR_TRT = 50 # 少于 50 字走 TRT
KEYWORDS_FOR_VLLM = ["代码", "写一个", "分析", "总结", "推理", "复杂", "长文"]

# ================= 核心引擎封装类 =================

class VLLMEngine:
"""
真实 vLLM 引擎封装
"""
def __init__(self, model_name: str):
print(f"🚀 [vLLM] 正在加载模型：{model_name} …")
try:
from vllm import LLM, SamplingParams
self.llm = LLM(
model=model_name,
tensor_parallel_size=1,
gpu_memory_utilization=0.4, # 只占 40% 显存，留给 TRT
enforce_eager=False,
disable_log_stats=True
)
self.sampling_params = SamplingParams(temperature=0.7, max_tokens=200)
print("✅ [vLLM] 模型加载完成")
self.available = True
except Exception as e:
print(f"❌ [vLLM] 加载失败 (可能是显存不足): {e}")
self.available = False

def generate(self, prompt: str) -> str:
if not self.available:
raise RuntimeError("vLLM engine not available")

# 真实调用 vLLM
outputs = self.llm.generate([prompt], self.sampling_params)
return outputs[0].outputs[0].text

class TRTEngine:
"""
TensorRT-LLM 引擎封装
注意：真实使用需要预先编译好 .engine 文件。
此处为了代码可直接运行演示路由逻辑，实现了真实接口结构，但内部计算做了模拟。
"""
def __init__(self, model_name: str, engine_dir: Optional[str] = None):
print(f"🚀 [TRT] 正在初始化引擎：{model_name} …")
self.model_name = model_name
self.engine_dir = engine_dir
self.runner = None
self.tokenizer = None
self.available = False

# 【真实环境开启这里】
# if engine_dir and os.path.exists(engine_dir):
# try:
# from tensorrt_llm.runtime import ModelRunner
# from transformers import AutoTokenizer
# import torch
# self.tokenizer = AutoTokenizer.from_pretrained(model_name)
# self.runner = ModelRunner.from_dir(engine_dir=engine_dir, rank=0)
# self.available = True
# print("✅ [TRT] 引擎加载完成 (真实模式)")
# except Exception as e:
# print(f"❌ [TRT] 加载失败：{e}")

# 【演示模式：如果没有 engine 文件，自动进入模拟模式以便你测试路由逻辑】
if not self.available:
print("⚠️ [TRT] 未检测到编译好的引擎文件，已切换到【模拟模式】用于演示路由逻辑。")
print(" (真实生产中，请在此处加载真实的 ModelRunner)")
self.available = True # 标记为可用，但内部是模拟的
self.is_mock = True
else:
self.is_mock = False

def generate(self, prompt: str) -> str:
if not self.available:
raise RuntimeError("TRT engine not available")

if self.is_mock:
# 模拟 TRT 的极速特性 (0.1s – 0.3s)
time.sleep(random.uniform(0.1, 0.3))
return f"[TRT 模拟响应] 针对 '{prompt[:20]}…' 的快速回答。"

# 【真实环境开启这里】
# import torch
# messages = [{"role": "user", "content": prompt}]
# text_input = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
# input_ids = self.tokenizer.encode(text_input, return_tensors="pt").squeeze(0).cuda()
# with torch.no_grad():
# outputs = self.runner.generate(…)
# return self.tokenizer.decode(outputs…)

# ================= 智能路由器 (核心逻辑) =================

class HybridRouter:
def __init__(self, vllm_engine: VLLMEngine, trt_engine: TRTEngine):
self.vllm = vllm_engine
self.trt = trt_engine

def decide_engine(self, prompt: str) -> str:
"""
纯内存决策，无网络开销
返回：'vllm' 或 'trt'
"""
length = len(prompt)

# 1. 关键词优先原则 (复杂任务强制 vLLM)
for kw in KEYWORDS_FOR_VLLM:
if kw in prompt:
return 'vllm'

# 2. 长度原则 (短任务优先 TRT)
if length <= MAX_LENGTH_FOR_TRT:
if self.trt.available:
return 'trt'
else:
return 'vllm' # TRT 不可用时降级

# 3. 默认长任务走 vLLM
return 'vllm'

def process(self, prompt: str) -> Dict[str, Any]:
start_time = time.time()

# 1. 决策
target = self.decide_engine(prompt)
reason = f"Len={len(prompt)}"

# 2. 执行 (直接函数调用，非 HTTP)
try:
if target == 'trt':
response = self.trt.generate(prompt)
engine_used = "TensorRT"
else:
response = self.vllm.generate(prompt)
engine_used = "vLLM"

latency = time.time() – start_time
return {
"prompt": prompt[:30] + "…",
"engine": engine_used,
"reason": reason,
"latency": latency,
"response": response[:50] + "…",
"success": True
}
except Exception as e:
# 故障降级逻辑 (如果首选挂了，尝试另一个)
print(f"⚠️ {target} 执行失败：{e}, 尝试降级…")
fallback = 'vllm' if target == 'trt' else 'trt'
try:
if fallback == 'trt' and self.trt.available:
response = self.trt.generate(prompt)
engine_used = "TensorRT (Fallback)"
elif fallback == 'vllm' and self.vllm.available:
response = self.vllm.generate(prompt)
engine_used = "vLLM (Fallback)"
else:
raise RuntimeError("No fallback available")

latency = time.time() – start_time
return {
"prompt": prompt[:30] + "…",
"engine": engine_used,
"reason": f"Fallback from {target}",
"latency": latency,
"response": response[:50] + "…",
"success": True
}
except Exception as e2:
return {
"prompt": prompt[:30] + "…",
"engine": "None",
"reason": "All Failed",
"latency": time.time() – start_time,
"response": "",
"success": False,
"error": str(e2)
}

# ================= 主程序 =================

def main():
print("="*60)
print("🧠 启动混合推理系统 (进程内直连模式)")
print("="*60)

# 1. 初始化引擎
# 注意：如果你的显存不够同时加载 7B+1.5B，vLLM 可能会失败，代码会自动处理
vllm_eng = VLLMEngine(VLLM_MODEL_NAME)

# TRT 需要编译好的 engine 目录，这里传 None 会进入模拟模式
trt_eng = TRTEngine(TRT_MODEL_NAME, engine_dir=None)

if not vllm_eng.available and not trt_eng.available:
print("❌ 错误：没有可用的引擎！")
return

router = HybridRouter(vllm_eng, trt_eng)

# 2. 构造测试数据
test_cases = [
"你好", # 短 -> TRT
"1+1 等于几？", # 短 -> TRT
"写一个 Python 快速排序", # 关键词 -> vLLM
"请分析量子力学的哲学意义", # 关键词 -> vLLM
"今天天气怎么样", # 短 -> TRT
"总结这篇文章的中心思想…", # 关键词 -> vLLM
"Hello", # 短 -> TRT
"帮我推理一下这个逻辑题…", # 关键词 -> vLLM
] * 3 # 重复几次

print(f"\\n⚡ 开始处理 {len(test_cases)} 个请求 (无 HTTP 开销)…\\n")

results = []
for prompt in test_cases:
res = router.process(prompt)
results.append(res)

# 实时打印
status = "✅" if res['success'] else "❌"
print(f"{status} 路由至:{res['engine']:<15} | 耗时:{res['latency']:.3f}s | 原因:{res['reason']}")
# print(f" 问：{res['prompt']}")
# print(f" 答：{res['response']}")

# 3. 统计
trt_count = sum(1 for r in results if "TensorRT" in r['engine'])
vllm_count = sum(1 for r in results if "vLLM" in r['engine'])
avg_latency = sum(r['latency'] for r in results if r['success']) / len(results)

print("\\n" + "="*60)
print("📊 最终统计")
print("="*60)
print(f"TensorRT 处理数：{trt_count}")
print(f"vLLM 处理数： {vllm_count}")
print(f"平均延迟： {avg_latency:.3f} 秒")
print(f"架构模式： 进程内直接调用 (In-Process Direct Call)")
print("="*60)

if __name__ == "__main__":
main()

💡 代码核心改动点

🛠️ 如何真正使用真实 TRT？

如果你想把模拟换成真实的：

VLLM+Ray（CPU调度）

核心思路：Ray + vLLM 分布式推理

架构变化：

• 之前 (单进程)：主进程加载模型 -> 循环处理列表。受限于单卡显存，无法利用多卡加速批量处理。
• 现在 (Ray 分布式)：
• Head Node：负责任务分发和调度。
• Worker Nodes (Actors)：每个 Worker 独占一张 GPU，加载一个 vLLM 实例。
• 并行策略：将大问题列表切分 (Chunking)，分发给多个 Worker 同时处理，最后汇总结果。

优势：

• 线性加速：如果有 4 张卡，理论速度提升接近 4 倍。
• 显存隔离：每个 Worker 独立占用显存，避免单进程管理多卡的复杂性。
• 弹性扩展：代码无需大改，即可从单机扩展到多机集群。

🛠️ 第一步：环境安装

你需要额外安装 ray 和 ray[default]：

# 基础依赖
pip install vllm tqdm

# 安装 Ray (支持分布式调度)
pip install "ray[default]"

📄 完整代码：vllm_ray_qwen.py

这段代码会自动检测你机器上的 GPU 数量，启动对应数量的 Worker，并并行处理任务。

python

import ray
import time
from typing import List, Dict, Any
from tqdm import tqdm
import os

# ================= 配置区域 =================
MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
# 如果显存较小，请改为 "Qwen/Qwen2.5-1.5B-Instruct"

# 每个 vLLM 实例占用的显存比例 (0.9 表示 90%)
# 注意：如果是多卡，每张卡都会独立加载一个模型副本
GPU_MEMORY_UTILIZATION = 0.9

# 采样参数
MAX_TOKENS = 256
TEMPERATURE = 0.7

# ================= Ray Actor 定义 =================

@ray.remote(num_gpus=1)
class VLLMWorker:
"""
Ray Actor: 每个实例独占一张 GPU，运行一个独立的 vLLM 引擎
"""
def __init__(self, model_name: str):
# 关键：在 Actor 内部导入 vllm，确保它在 GPU 上下文中初始化
from vllm import LLM, SamplingParams

# 获取当前 Actor 绑定的 GPU ID (用于日志调试)
gpu_id = os.environ.get("CUDA_VISIBLE_DEVICES", "Unknown")
print(f"🚀 [Worker] 正在 GPU {gpu_id} 上加载模型：{model_name} …")

self.llm = LLM(
model=model_name,
tensor_parallel_size=1, # 每个 Worker 只用 1 张卡
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
enforce_eager=False,
disable_log_stats=True
)

self.sampling_params = SamplingParams(
temperature=TEMPERATURE,
max_tokens=MAX_TOKENS,
top_p=0.8
)
print(f"✅ [Worker] GPU {gpu_id} 准备就绪")

def generate_batch(self, prompts: List[str]) -> List[str]:
"""
处理一批提示词，返回生成的文本列表
"""
if not prompts:
return []

outputs = self.llm.generate(prompts, self.sampling_params)
return [out.outputs[0].text for out in outputs]

# ================= 主调度逻辑 =================

def chunkify(data: List[Any], n_chunks: int) -> List[List[Any]]:
"""将列表均匀切分成 n 份"""
avg = len(data) / float(n_chunks)
out = []
last = 0.0
while last < len(data):
out.append(data[int(last):int(last + avg)])
last += avg
return out

def main():
print("="*60)
print("🧠 启动 Ray + vLLM 分布式推理系统")
print("="*60)

# 1. 初始化 Ray
# ignore_reinit_error=True 防止重复初始化报错
# 如果是在多机集群，这里需要指定 address="ip:port"
if not ray.is_initialized():
ray.init(ignore_reinit_error=True, log_to_driver=True)

print(f"✅ Ray 已初始化。集群资源：{ray.cluster_resources()}")

# 自动检测可用 GPU 数量
num_gpus = int(ray.cluster_resources().get("GPU", 0))
if num_gpus == 0:
print("❌ 未检测到 GPU 资源，请检查 NVIDIA 驱动和 Ray 配置。")
return

print(f"📡 检测到 {num_gpus} 张可用 GPU，将启动 {num_gpus} 个 Worker。")

# 2. 创建 Worker 池
# 每个 Worker 会独占一张 GPU (因为 @ray.remote(num_gpus=1))
workers = [VLLMWorker.remote(MODEL_NAME) for _ in range(num_gpus)]

# 等待所有 Worker 加载完成 (阻塞直到所有模型加载完毕)
print("⏳ 等待所有 Worker 加载模型完成…")
ray.get([w.__ready__() for w in workers]) if hasattr(workers[0], '__ready__') else time.sleep(2)
# 注：上面这行是伪代码逻辑，实际我们通过在 __init__ 打印日志来确认，
# 或者简单 sleep 等待，因为 ray.get(worker.init.remote()) 模式需要显式定义 init 方法。
# 为了简化，我们直接在第一次调用时隐式等待，或者加个简单的 barrier。
# 这里为了稳健，我们调用一个空方法或直接进行任务分发，Ray 会自动排队等待初始化完成。

# 显式Barrier: 确保所有 worker 初始化完成
ray.get([w.generate_batch.remote(["test"]) for w in workers]) # 预热/同步
print("✅ 所有 Worker 已就绪。")

# 3. 准备数据 (模拟 200 个问题)
all_prompts = [f"问题 {i}: 请简要解释量子纠缠。" for i in range(200)]
print(f"📝 待处理任务总数：{len(all_prompts)}")

# 4. 任务切分与分发
# 将 200 个问题切成 num_gpus 份，每份给一个 Worker
chunks = chunkify(all_prompts, num_gpus)

print(f"⚡ 开始并发处理 (分为 {len(chunks)} 个批次)…")
start_time = time.time()

# 异步提交所有任务
futures = []
for i, chunk in enumerate(chunks):
if chunk:
# remote() 返回的是 ObjectRef (未来对象)，不会阻塞
future = workers[i].generate_batch.remote(chunk)
futures.append(future)

# 5. 收集结果 (阻塞直到所有任务完成)
# ray.get 会并行等待所有结果返回
all_results = ray.get(futures)

end_time = time.time()
total_time = end_time – start_time

# 6. 结果汇总与统计
# all_results 是一个列表的列表，需要展平
flat_results = [text for chunk_res in all_results for text in chunk_res]

total_tokens_estimated = sum(len(text.split()) * 1.3 for text in flat_results) # 粗略估算 token 数

print("\\n" + "="*60)
print("📊 分布式测试报告")
print("="*60)
print(f"GPU 数量：{num_gpus}")
print(f"总任务数：{len(flat_results)}")
print(f"总耗时：{total_time:.2f} 秒")
print(f"平均每秒处理请求数 (QPS): {len(flat_results)/total_time:.2f}")
print(f"预估吞吐量 (Tokens/s): {total_tokens_estimated/total_time:.2f}")
print("="*60)

# 展示部分结果
print("\\n💡 部分结果示例:")
for i, text in enumerate(flat_results[:3]):
print(f"[{i}] {text[:100]}…")

# 7. 释放资源
ray.shutdown()
print("\\n👋 Ray 会话已关闭。")

if __name__ == "__main__":
main()

💡 代码核心解析

1. @ray.remote(num_gpus=1)

这是最关键的一行。

• 它告诉 Ray：“这个类 (VLLMWorker) 运行时必须独占 1 张 GPU”。
• 当你创建 workers = [VLLMWorker.remote() for _ in range(4)] 时，Ray 会自动检测你有几张卡。
  • 如果你有 4 张卡，它会启动 4 个进程，每个进程绑定一张卡。
  • 如果你只有 1 张卡，它会排队执行（或者报错，取决于配置），通常建议 Worker 数量 <= GPU 数量。

2. 进程隔离与显存管理

• 每个 VLLMWorker 都是一个独立的 Python 进程。
• 每个进程内部都会执行 from vllm import LLM，这意味着每张显卡上都加载了一个完整的模型副本。
• 优点：互不干扰，一个 Worker 崩了不影响其他。
• 缺点：显存消耗是 N * 模型大小。如果你的显存不够加载 N 个副本，需要减少 Worker 数量或改用“张量并行”（Tensor Parallelism，即多个卡合力跑一个大模型，而不是每个卡跑一个小模型）。
  • 本代码采用的是 数据并行 (Data Parallelism) 策略：多卡各跑各的，处理不同数据。适合高吞吐场景。

3. 任务切分 (chunkify)

• 我们将 200 个问题平均分给 4 个 Worker，每个处理 50 个。
• worker.generate_batch.remote(chunk) 是异步的。主程序发出指令后立刻继续，不会等待该 Worker 算完。
• ray.get(futures) 在最后统一等待所有结果回来。

4. 性能提升原理

• 单卡模式：处理 200 个问题，串行排队。
• Ray 多卡模式：4 张卡同时处理，理论上速度提升 4 倍。
• vLLM 内部的 PagedAttention 保证了单个 Worker 在处理那 50 个问题时的内部并发效率，而 Ray 保证了多卡之间的宏观并发效率。

🚀 如何运行？

场景 A：单机多卡 (最常见)

直接运行脚本即可。Ray 会自动发现本地所有 GPU。

python vllm_ray_qwen.py

观察输出：你会看到 [Worker] 正在 GPU 0 上加载…, [Worker] 正在 GPU 1 上加载… 等日志。

场景 B：多机集群 (进阶)

如果你想用 2 台机器（每台 4 卡，共 8 卡）：

⚠️ 注意事项