第3章基础推理流程

一、基础推理流程：非流式文本生成

各位同学，今天咱们来聊聊 TensorRT-LLM 最基础的操作——非流式文本生成。

说白了，就是给模型一个输入，它一次性把完整结果吐给你。

没有流式那种一个字一个字往外蹦的效果，但胜在简单、直接、好理解。

1.1 环境准备与模型加载

在开始之前，我假设你已经装好了 TensorRT-LLM 环境。

如果还没装，建议先跑一遍官方文档里的 Quick Start。

我个人习惯用 Docker 镜像，省去很多依赖冲突的麻烦。

模型加载这块，我踩过不少坑。

记得有一次，我直接加载了一个没转换过的 HuggingFace 模型，结果报了一堆算子不支持的错。

后来才明白，TensorRT-LLM 需要先把模型转成 TRT-LLM 格式。

注意： 模型转换这一步不能省。直接用 HuggingFace 权重跑推理，会触发大量 fallback 操作，性能直接打折扣。

正确的加载方式是这样的：

from tensorrt_llm import LLM, SamplingParams

# 加载已转换的模型
model = LLM(model="path/to/your/trt-llm-model")

# 或者指定 engine 目录
model = LLM(engine_dir="path/to/engine_dir")

这里有个小细节：engine_dir 指向的是你转换后生成的 engine 文件夹。

里面通常包含 config.json 和多个 .engine 文件。

1.2 构建采样参数

采样参数，说白了就是控制模型怎么「选词」的。

你想想看，模型每次预测下一个 token 时，其实会给出一个概率分布。

采样参数决定了它从分布里怎么挑一个出来。

我常用的几个参数：

参数名	作用	推荐值
temperature	控制随机性。值越大，输出越随机	0.7 ~ 1.0
top_p	核采样，只保留概率累加和 top_p 的 token	0.9
top_k	只保留概率最高的 k 个 token	50
max_tokens	最大生成长度	根据任务定

代码实现很简单：

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.9,
    top_k=50,
    max_tokens=512
)

小技巧： 如果你做的是代码生成或数学推理，建议把 temperature 设低一点（0.1~0.3）。我试过，输出更稳定，幻觉也少。

1.3 执行推理

好了，模型加载好了，参数也配好了，接下来就是最核心的一步——跑推理。

prompt = "请用 Python 写一个快速排序算法"

# 生成
outputs = model.generate([prompt], sampling_params)

# 提取结果
generated_text = outputs[0].outputs[0].text
print(generated_text)

这段代码看起来简单，但背后发生的事情可不简单。

我简单拆解一下：

model.generate 接收一个 prompt 列表，支持批量推理
内部会先对 prompt 做 tokenize
然后跑自回归生成，直到遇到 EOS token 或达到 max_tokens
最后把生成的 token 序列 decode 回文本

嗯，这里要注意：outputs[0].outputs[0].text 这个写法有点绕。

outputs 是列表，每个元素对应一个 prompt 的生成结果。

每个结果里又有多个 outputs（如果你设置了 num_return_sequences > 1）。

1.4 完整示例与避坑

我把上面几段拼起来，给你一个可以直接跑的完整代码：

from tensorrt_llm import LLM, SamplingParams

# 1. 加载模型
model = LLM(model="models/llama-7b-trt")

# 2. 设置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=256
)

# 3. 准备输入
prompts = [
    "什么是深度学习？",
    "解释一下 Transformer 架构"
]

# 4. 生成
outputs = model.generate(prompts, sampling_params)

# 5. 打印结果
for i, output in enumerate(outputs):
    print(f"Prompt {i+1}: {prompts[i]}")
    print(f"Generated: {output.outputs[0].text}")
    print("-" * 50)

避坑指南： 我曾经遇到过一个问题——生成的文本开头总是重复 prompt 的内容。后来发现是 add_special_tokens 参数没设对。如果你也遇到类似情况，检查一下 tokenizer 的配置。

1.5 核心流程总结

非流式推理的流程，其实就三步：

加载模型——确保是 TRT-LLM 格式
配采样参数——控制生成行为
调用 generate——拿到完整结果

为了让你更直观地理解，我画了一张流程图：

这张图把流程串起来了。

你从左边开始，一步步走到右边，最后拿到结果。

是不是很清晰？

1.6 性能与适用场景

非流式推理最大的优势是 吞吐量高。

因为它可以一次性处理多个 prompt，充分利用 GPU 的并行能力。

我做过测试，在 A100 上跑 LLaMA-7B，batch size 设为 8，吞吐量能达到单条推理的 5 倍以上。

但缺点也很明显——延迟高。

用户必须等所有 token 生成完才能看到结果。

如果 max_tokens 设得大，等待时间会很长。

所以，非流式适合这些场景：

离线批量处理（比如文档摘要、数据标注）
对实时性要求不高的应用
需要完整结果才能做后续处理的场景

而如果你做的是聊天机器人、实时翻译这类应用，那就要用流式推理了。

下一章我会详细讲。

  核心要点： 非流式推理是 TensorRT-LLM 最基础的用法。掌握它，你就掌握了推理的骨架。后续的流式、并发、优化，都是在这个基础上做加法。

好了，这一章的内容就到这里。

代码不多，但都是实打实能跑起来的。

建议你打开电脑，跟着敲一遍，感受一下 TensorRT-LLM 的推理流程。

遇到问题别慌，翻翻官方文档，或者来群里问我。

📚 课程目录

一、基础推理流程：非流式文本生成

1.1 环境准备与模型加载

1.2 构建采样参数

1.3 执行推理

1.4 完整示例与避坑

1.5 核心流程总结

1.6 性能与适用场景