4. TensorRT LLM 核心概念

说实话，刚接触 TensorRT LLM 的时候，我也被这一堆概念搞得有点晕。Engine、Builder、Network Definition……每个词都认识，但串起来就不知道谁是谁了。后来我花了一周时间，把每个概念拆开揉碎了去理解，才真正搞明白它们之间的关系。

这一章，我就带你把这些核心概念一个个说清楚。你想想看，如果连这些基础都不牢，后面做部署肯定要踩坑。

4.1 Builder：模型的建筑师

Builder 是 TensorRT 的入口。它的任务很简单——把你的模型从框架格式（比如 PyTorch、ONNX）转换成 TensorRT 能优化的中间表示。

我习惯把 Builder 想象成一个建筑队队长。他手里拿着设计图（Network Definition），指挥工人（CUDA 核心）干活，最后盖出房子（Engine）。

Builder 的核心配置项不多，但每个都很关键：

配置项	作用	我的建议
max_workspace_size	控制 GPU 显存使用上限	别设太大，留点余量给推理
precision	指定计算精度（FP32/FP16/INT8）	FP16 是性价比最高的选择
strict_type_constraints	是否强制遵守精度约束	一般不开，开了可能影响性能

4.2 Network Definition：模型的设计蓝图

Network Definition 就是 Builder 手里的设计图。它描述了模型的网络结构——有哪些层、每层的参数是什么、数据怎么流动。

在 TensorRT LLM 里，Network Definition 有两种构建方式：

• 从 ONNX 导入：最省事，但可能丢失一些细节
• 手动构建：灵活度高，适合定制化需求

我个人更推荐手动构建，尤其是做 LLM 部署的时候。为什么呢？因为 ONNX 对 Transformer 结构的支持并不完美，经常会出现一些奇怪的算子转换问题。

举个例子，我去年部署一个 模型时，ONNX 把 Attention 里的 Softmax 拆成了好几个小算子，导致 TensorRT 优化效果很差。后来我用手动构建 Network Definition 的方式，把整个 Attention 写成一个 Plugin，性能直接提升了 30%。

4.3 Plugin：给模型装上特殊零件

Plugin 是 TensorRT 最强大的功能之一。说白了，它允许你自定义算子，把一些标准算子组合不出来的操作塞进去。

什么时候需要 Plugin？

• 模型里有 TensorRT 不支持的算子
• 标准算子组合性能太差
• 你想用一些特殊的 CUDA 优化技巧

我记得有一次，一个客户要求把 FlashAttention 集成到他们的 LLM 里。TensorRT 原生不支持 FlashAttention，怎么办？写 Plugin 呗。

Plugin 的开发流程大概是这样的：

// 1. 继承 IPluginV2DynamicExt 类
class MyFlashAttentionPlugin : public IPluginV2DynamicExt {
    // 2. 实现 getOutputDimensions、enqueue 等虚函数
    // 3. 注册 Plugin
};
REGISTER_TENSORRT_PLUGIN(MyFlashAttentionPluginCreator);

4.4 Optimization Profile：性能调优的开关

Optimization Profile 是 TensorRT LLM 里一个容易被忽视的概念。它决定了 Builder 在优化模型时，会考虑哪些输入形状。

为什么需要它？因为 LLM 的输入长度是动态的。你今天可能处理 512 个 ，明天可能处理 2048 个。如果没有 Optimization Profile，Builder 只能按固定形状优化，性能会大打折扣。

Optimization Profile 包含三个关键参数：

参数	含义	典型值
min_shape	最小输入形状	[1, 1, 512]
opt_shape	最优输入形状	[1, 1, 1024]
max_shape	最大输入形状	[1, 1, 2048]

我一般会这样设置：min_shape 取实际业务的最小值，max_shape 取硬件能支持的最大值，opt_shape 取最常见的值。这样 Builder 会优先优化常见场景，同时保证极端情况也能跑。

4.5 Engine：最终的可执行产物

Engine 是 Builder 的最终输出。它包含了优化后的网络结构、权重、CUDA kernel 代码，以及运行时需要的所有信息。

你可以把 Engine 理解成一个编译好的 .so 文件。它可以直接加载到 GPU 上执行，不需要再依赖任何框架。

Engine 有几个特点：

• 平台相关：不同 GPU 型号生成的 Engine 不通用
• 精度固定：一旦生成，计算精度就定了
• 可序列化：可以保存成文件，下次直接加载

4.6 它们之间的关系

说了这么多，你可能还是有点乱。没关系，我画了一张图，把这些概念串起来：

从这张图你能看出来，Builder 是核心调度者。它拿到 Network Definition 和 Plugin 后，再根据 Optimization Profile 的约束，进行层融合、精度校准、内存优化等一系列操作，最后输出 Engine。

嗯，这里要注意一点：Plugin 不是必须的。如果你的模型所有算子 TensorRT 都支持，那完全不需要 Plugin。但现实情况是，LLM 里总有一些特殊操作，比如 FlashAttention、Grouped Query Attention，这些都需要 Plugin 来实现。

4.7 实战：从零构建一个 Engine

光说不练假把式。我们来看一段完整的代码，把刚才的概念串起来：

# 1. 创建 Builder
builder = trt.Builder(logger)

# 2. 创建 Network Definition
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 3. 添加网络层（以 Transformer 的 Attention 为例）
input_tensor = network.add_input(
    name="input_ids",
    dtype=trt.float16,
    shape=(-1, -1, 768)  # 动态形状
)

# 4. 注册 Plugin（如果有）
plugin = get_flash_attention_plugin()
plugin_node = network.add_plugin_v2([input_tensor], plugin)

# 5. 设置 Optimization Profile
profile = builder.create_optimization_profile()
profile.set_shape("input_ids", 
    min=(1, 1, 768),   # 最小 batch=1, seq_len=1
    opt=(1, 512, 768), # 最优 batch=1, seq_len=512
    max=(4, 2048, 768) # 最大 batch=4, seq_len=2048
)

# 6. 配置 Builder
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB
config.set_flag(trt.BuilderFlag.FP16)
config.add_optimization_profile(profile)

# 7. 构建 Engine
engine = builder.build_engine(network, config)

# 8. 序列化保存
with open("model.engine", "wb") as f:
    f.write(engine.serialize())

好了，这一章的核心概念就讲到这里。你可能会觉得内容有点多，但相信我，把这些概念吃透了，后面做部署会顺畅很多。Builder 怎么配、Plugin 怎么写、Profile 怎么设，这些都会成为你的肌肉记忆。