📚 课程目录第4章 TensorRT-LLM核心概念

4. TensorRT-LLM核心概念：引擎、执行上下文、张量、层与优化策略

在正式编写模型编译脚本之前，必须先建立一套统一的“心智模型”。TensorRT-LLM 的推理系统可以拆解为五个核心概念：引擎（Engine）、执行上下文（ExecutionContext）、张量（Tensor）、层（Layer）与优化策略（Optimization Strategy）。它们之间的关系不是并列的，而是层层递进、相互约束的。理解这套关系，是后续调试性能问题和编写生产级推理代码的基石。

4.1 核心概念与层次关系

我们可以把整个推理生命周期想象成“建造工厂”与“使用工厂”两个阶段。引擎是图纸与生产线的结合体，而执行上下文则是某一次具体的生产任务单。

引擎（Engine）：编译期的最终产物，一个二进制文件（通常以 .engine 结尾）。它包含了网络结构的精确计算图、每层选择的 kernel 实现、权重常量以及所有算子融合后的指令序列。引擎是不可变的，一旦生成，其内部的计算路径就固定了。
执行上下文（ExecutionContext）：引擎的“运行时实例”。同一个引擎可以被实例化多个上下文，每个上下文拥有独立的中间激活值缓冲区、KV Cache 指针和 Profiler 状态。它负责维护推理过程中的可变状态。
张量（Tensor）：数据流的基本单位。在编译期，它描述维度与数据类型；在运行期，它绑定到具体的内存地址。TensorRT-LLM 中的张量通常显式携带形状（Shape）与数据类型（dtype）。
层（Layer）：计算图的基本单元。在 TensorRT-LLM 中，层不仅包括 Conv、GEMM，更常见的是 Attention、RaggedSoftmax、GptAttention 等高级融合层。
优化策略：连接“层”与“引擎”的桥梁。包括层融合、精度校准（FP16/INT8/FP8）、内核自动调优（Tactic Selection）等。

关键认知：引擎是“静态的图”，执行上下文是“动态的态”。不要试图在推理时修改引擎中的层参数，所有可变数据（如输入 token 序列）必须通过张量绑定传入执行上下文。

4.2 张量的特殊语义：形状与内存布局

TensorRT-LLM 中的张量与 PyTorch 张量有一个显著区别：它不拥有内存，而是通过 setMemoryAddress 或 setAddress 绑定到外部显存指针。这种设计是为了避免推理时的显存拷贝。

在构建引擎时，张量维度分为三种：

静态维度：编译时完全确定，如 hidden_size=4096。
动态维度（优化维度）：用 -1 表示，运行时可以变化，例如 batch size 或序列长度。
形状张量（Shape Tensor）：作为输入参与计算，用于控制循环或条件分支。

一个典型的输入张量声明如下：

# 构建期定义输入张量
input_ids = network.add_input(name="input_ids", dtype=trt.int32, shape=(-1, -1))
# 第一个 -1 是 batch size，第二个 -1 是序列长度（均为动态）

# 运行期绑定实际显存地址
context.set_tensor_address("input_ids", input_ids_ptr)
context.set_tensor_address("logits", output_ptr)

易错点：动态维度并非越多越好。每增加一个动态维度，引擎内部的优化器就需要考虑更多形状组合，可能导致 Tactic 选择时间变长，甚至产生额外的显存预留。通常只将 batch size 与 序列长度 设为动态，而隐藏层维度保持静态。

4.3 优化策略的落地形态

优化策略不是抽象的，它们最终体现在引擎的构建参数与网络结构上。下表总结了常见策略及其影响层面：

优化策略	作用对象	典型收益	风险/代价
层融合（Fusion）	多个层 → 单个 kernel	减少显存读写与 kernel 启动开销	增加编译时间
精度校准（FP8/INT8）	权重与激活张量	显存减半，吞吐提升	精度损失，需校准数据集
KV Cache 量化	KV Cache 张量	支持更长上下文	可能影响生成质量
Tactic 自动调优	单个层实现选择	选择最快 kernel	编译时间剧增（可限制）
In-Flight 批处理	执行上下文调度	提高 GPU 利用率	需要显式管理请求队列

4.4 案例：从层到引擎的编译流程

假设我们要编译一个仅包含 GptAttention 与 MLP 的简单 Decoder 层。实际构建流程如下：

创建网络定义：使用 Builder 创建网络，指定 precision=FP16。
添加层：调用 add_attention 与 add_mlp，此时层尚未绑定具体权重。
设置张量形状：标记输入张量的动态维度。
应用优化策略：设置 builder_opt_level=3（最高优化），允许层融合与算子替换。
序列化引擎：调用 build_serialized_network 生成引擎字节流。

from tensorrt_llm import Builder
from tensorrt_llm.network import Network

builder = Builder()
network = Network()
with builder.create_builder_config(opt_level=3) as config:
    config.set_precision("fp16")
    # 添加输入张量
    input_t = network.add_input("input_ids", "int32", (-1, -1))
    # 添加注意力层（内部包含 QKV GEMM + Softmax + 输出投影）
    attn = network.add_attention("self_attn", input_t, hidden_size=4096, num_heads=32)
    # 添加 MLP 层
    mlp = network.add_mlp("mlp", attn.get_output(0), hidden_size=4096, ffn_size=11008)
    # 标记输出
    network.mark_output(mlp.get_output(0))
    # 构建引擎
    engine = builder.build_engine(network, config)
    with open("decoder.engine", "wb") as f:
        f.write(engine)

这段代码展示了“层”如何通过“优化策略”最终固化为“引擎”。注意 add_attention 是一个高级融合层，它内部已经包含了多头注意力的所有计算，这正是层融合策略在 API 层面的体现。

4.5 常见误区与复盘

误区一：混淆“引擎”与“模型权重”。引擎文件不仅包含权重，还包含计算图与 kernel 选择。因此，即使权重相同，不同 TensorRT 版本或不同 GPU 架构生成的引擎文件不能互换。

误区二：认为执行上下文是线程安全的。同一个 ExecutionContext 不能同时被多个线程调用。如果使用多线程并发推理，必须为每个线程创建独立的上下文，或者使用线程池加锁串行化。

误区三：忽略张量的“形状对齐”。在运行期绑定张量地址时，如果实际显存大小小于引擎期望的优化维度大小，会导致越界读写。例如引擎期望最大序列长度 2048，但运行时只分配了 1024 长度的 KV Cache，就会产生未定义行为。

复盘建议：当遇到推理结果错误或 CUDA error 时，请按以下顺序排查：

检查引擎构建时的 max_seq_len 与 max_batch_size 是否覆盖了运行时的实际请求。
确认执行上下文绑定的张量地址是否指向正确的设备内存（而非主机内存）。
验证输入张量的 dtype 是否与引擎期望一致（例如 FP16 引擎不能喂 FP32 张量）。
如果使用了 INT8/FP8，检查校准是否在代表性数据集上完成。