4. 理解引擎构建参数
任务目标
本案例通过修改 TensorRT-LLM 引擎构建参数（max_batch_size、max_seq_len、max_num_tokens），观察不同配置下引擎构建耗时、显存占用及推理延迟的变化，建立参数配置与性能表现之间的直观关联。
环境与依赖
- TensorRT-LLM（版本 ≥ 0.10.0）
- CUDA 12.x + 对应 cuDNN
- GPU：NVIDIA Ampere 或更新架构（如 A10/A30/L40S）
- 模型：使用内置的
核心实现
以下脚本通过循环构建同一 模型的三组不同引擎配置，记录每次构建的耗时与显存峰值，并对每组引擎执行相同的推理请求以测量端到端延迟。构建参数中 max_num_tokens 按公式 max_batch_size × max_seq_len 自动推导，避免参数冲突。
```python
# build_and_bench.py
import time
import torch
import tensorrt_llm
from tensorrt_llm.builder import Builder
from tensorrt_llm.models import
tensorrt_llm.logger.set_level("error")
CONFIGS = [
{"name": "bs1_seq128", "max_batch_size": 1, "max_seq_len": 128},
{"name": "bs8_seq128", "max_batch_size": 8, "max_seq_len": 128},
{"name": "bs8_seq512", "max_batch_size": 8, "max_seq_len": 512},
]
def build_engine(cfg, workdir):
"""构建引擎并返回 (构建耗时秒, 显存峰值GB)"""
torch.cuda.reset_peak_memory_stats()
start = time.time()
builder = Builder()
builder.max_batch_size = cfg["max_batch_size"]
builder.max_seq_len = cfg["max_seq_len"]
builder.max_num_tokens = cfg["max_batch_size"] * cfg["max_seq_len"]
builder.engine_dir = f"{workdir}/{cfg['name']}"
model = .from_pretrained("", dtype="float16")
builder.build(model)
build_time = time.time() - start
peak_mem = torch.cuda.max_memory_allocated() / 1024**3
return build_time, peak_mem
def run_inference(engine_dir, batch_size, seq_len):
"""加载引擎并执行一次生成，返回端到端延迟(ms)"""
from tensorrt_llm.runtime import ModelRunner
runner = ModelRunner.from_dir(engine_dir)
input_ids = torch.randint(0, 50256, (batch_size, seq_len)).cuda()
start = time.time()
runner.generate(input_ids, max_new_tokens=32)
torch.cuda.synchronize()
return (time.time() - start) * 1000
if __name__ == "__main__":
workdir = "/tmp/trtllm_bench"
os.makedirs(workdir, exist_ok=True)
print(f"{'配置':<14} {'构建耗时(s)':<12} {'显存峰值(GB)':<14} {'推理延迟(ms)'}")
for cfg in CONFIGS:
bt, mem = build_engine(cfg, workdir)
# 推理时使用与构建相同的 batch_size 和 seq_len
lat = run_inference(f"{workdir}/{cfg['name']}",
cfg["max_batch_size"], cfg["max_seq_len"])
print(f"{cfg['name']:<14} {bt:<12.2f} {mem:<14.2f} {lat:.2f}")
```
运行方式
# 在 TensorRT-LLM 容器内执行
python build_and_bench.py
验证方法与判定标准
脚本正常结束并输出三行结果即视为构建成功。判定参数影响是否如预期，需观察以下趋势：
- bs1_seq128 vs bs8_seq128：增大
max_batch_size 后，构建耗时与显存峰值应上升（因优化空间更大），单请求推理延迟通常略有增加（批处理调度开销）。
- bs8_seq128 vs bs8_seq512：增大
max_seq_len 后，构建显存与耗时显著上升（KV cache 空间增大），推理延迟明显增加（序列更长）。
- 若三行结果均成功输出且上述趋势成立，则本案例验证通过。若某配置构建失败，请检查 GPU 显存是否满足
bs8_seq512 的需求（建议 ≥ 16GB）。
常见问题
- 显存不足（OOM）：将
bs8_seq512 改为 bs4_seq256 或使用更小模型（如 -small）。
- 构建时间过长：首次构建需编译 CUDA kernel，属正常现象；后续重复运行会使用缓存。
- 参数冲突报错：TensorRT-LLM 要求
max_num_tokens 不超过 max_batch_size × max_seq_len，本脚本已按乘积设置，无需手动调整。