3、TensorRT-LLM初探：架构解析、核心组件（Builder、Runtime、Plugin）
好，咱们正式开始聊 TensorRT-LLM。说实话，我第一次接触这个框架的时候，第一反应是「这玩意儿跟 TensorRT 到底啥关系？」后来用了一段时间才摸清楚门道。今天我就带你把它拆开看看，到底里面装了些什么。
3.1 整体架构：一张图看懂
先别急着看代码，咱们先搞清楚 TensorRT-LLM 的「骨架」。我个人习惯把它的架构分成三层：
- 前端层：负责模型定义、权重加载、图构建。说白了就是把你手里的 PyTorch 模型「翻译」成 TensorRT 能理解的东西。
- 中间层：核心引擎，包括 Builder（构建器）和 Runtime（运行时）。Builder 负责优化和编译，Runtime 负责跑起来。
- 后端层：Plugin（插件）系统，用来处理那些标准算子搞不定的「特殊操作」，比如 FlashAttention、PageAttention 这些。
下面这张图是我自己画的，你可以对照着看：
嗯，这里要注意：这三层不是严格串行的。实际开发中，你经常需要在前端和后端之间来回调整，比如发现某个算子性能不好，就得写个 Plugin 来替换它。
3.2 Builder：模型的「锻造师」
Builder 是 TensorRT-LLM 里我最喜欢的组件。为什么？因为它把一堆「脏活累活」全包了。你想想看，一个 70B 的模型，几十亿参数，要在 GPU 上跑得又快又省显存，这中间有多少优化要做？
Builder 主要干这几件事：
- 图优化：合并算子、消除冗余、常量折叠。我记得有一次，一个简单的 LayerNorm 被 Builder 优化后，速度提升了 30%。
- 精度校准
- 内核自动调优：Builder 会尝试不同的 CUDA kernel 实现，选最快的那个。这个过程叫「tactic selection」，说白了就是「哪个快用哪个」。
核心概念：Engine
Builder 的输出是一个 .engine 文件。这个文件包含了优化后的模型结构、权重、以及 CUDA kernel 的二进制代码。说白了，它就是 TensorRT-LLM 的「可执行文件」。
下面是一个典型的 Builder 配置示例：
from tensorrt_llm import Builder
from tensorrt_llm.models import LLaMAForCausalLM
# 加载模型
model = LLaMAForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
# 创建 Builder
builder = Builder()
builder_config = builder.create_builder_config(
precision="float16", # 精度
max_batch_size=8, # 最大 batch
max_input_len=2048, # 最大输入长度
max_output_len=512, # 最大输出长度
)
# 构建引擎
engine = builder.build_engine(model, builder_config)
engine.save("llama2_7b_fp16.engine")避坑指南
我曾经在构建引擎时把 max_batch_size 设得太大（比如 64），结果显存直接爆了。建议你先用小 batch 测试，确认没问题再往上加。
3.3 Runtime：模型的「驾驶员」
Builder 造好了引擎，Runtime 负责让它跑起来。Runtime 的核心职责就两个：加载引擎、执行推理。
Runtime 的工作流程其实很简单：
- 加载
.engine文件到显存 - 创建推理上下文（Context）
- 绑定输入输出张量
- 执行推理
- 获取结果
代码示例：
from tensorrt_llm import Runtime
# 加载引擎
runtime = Runtime()
engine = runtime.deserialize_engine("llama2_7b_fp16.engine")
# 创建上下文
context = engine.create_execution_context()
# 准备输入
input_ids = ... # shape: [1, seq_len]
input_lengths = ... # shape: [1]
# 绑定张量
context.set_input_shape("input_ids", input_ids.shape)
context.set_tensor_address("input_ids", input_ids.data_ptr())
# 执行推理
context.execute_async()
# 获取输出
output_ids = context.get_tensor("output_ids")注意
Runtime 是线程安全的，但 Context 不是。如果你要用多线程推理，记得每个线程创建自己的 Context。
3.4 Plugin：模型的「特种兵」
Plugin 是 TensorRT-LLM 最灵活的部分。标准算子搞不定的操作，比如 FlashAttention、PageAttention、GQA（Grouped Query Attention），都需要 Plugin 来实现。
我刚开始用 TensorRT-LLM 时，遇到一个自定义的激活函数，标准算子不支持。当时我硬着头皮写了一个 Plugin，结果发现其实没那么难。Plugin 的核心就三个接口：
| 接口 | 作用 | 说明 |
|---|---|---|
getPluginType() |
返回插件类型 | 用于唯一标识这个 Plugin |
enqueue() |
执行计算 | 这里是 CUDA kernel 的入口 |
getOutputDimensions() |
计算输出维度 | 告诉 Builder 输出张量长什么样 |
举个例子，FlashAttention 的 Plugin 大概长这样：
class FlashAttentionPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
const char* getPluginType() const noexcept override {
return "FlashAttention";
}
int32_t enqueue(
const nvinfer1::PluginTensorDesc* inputDesc,
const nvinfer1::PluginTensorDesc* outputDesc,
const void* const* inputs, void* const* outputs,
void* workspace, cudaStream_t stream) noexcept override {
// 调用 FlashAttention CUDA kernel
flash_attn_kernel(
inputs[0], // Q
inputs[1], // K
inputs[2], // V
outputs[0], // O
...,
stream
);
return 0;
}
};我的经验
写 Plugin 时，最容易踩的坑是维度计算。我曾经因为 getOutputDimensions() 返回的维度不对，导致 Builder 报错「shape mismatch」。建议你在写 Plugin 之前，先用 PyTorch 把计算逻辑跑通，确认输入输出维度没问题。
3.5 三者如何协作？
Builder、Runtime、Plugin 三者之间的关系，我用一个实际场景来说明：
假设你要部署一个 LLaMA 模型，用了 GQA（Grouped Query Attention）。
- Builder 阶段：Builder 读取模型，发现 GQA 这个操作标准算子不支持。于是它去 Plugin 注册表里找，找到了 GQA 的 Plugin。
- 优化阶段：Builder 把 GQA Plugin 嵌入到计算图中，然后进行全局优化。比如把连续的矩阵乘法和 GQA 融合成一个 kernel。
- Runtime 阶段：推理时，Runtime 遇到 GQA 操作，直接调用 Plugin 的
enqueue()方法执行。
说白了，Builder 是「设计师」，Runtime 是「执行者」，Plugin 是「特种工具」。三者缺一不可。
小技巧
如果你不确定某个操作有没有现成的 Plugin，可以查看 TensorRT-LLM 的 tensorrt_llm/plugins/ 目录。里面列出了所有官方支持的 Plugin。
好了，这一章的内容就到这里。TensorRT-LLM 的架构其实不复杂，关键是理解 Builder、Runtime、Plugin 各自扮演的角色。下一章我们会深入 Builder 的配置细节，到时候再聊。