第4章 NVIDIA TensorRT（下）

4、NVIDIA TensorRT（下）：插件开发、多流执行、与Triton Server集成实践

好，咱们接着聊 TensorRT。上一章我们把 TensorRT 的基础搭建、模型转换和 INT8 量化讲透了。这一章，我带你深入几个真正让 TensorRT 发挥威力的“硬核”话题：插件开发、多流执行，以及和 Triton Server 的集成。

说实话，这三个东西，才是你在生产环境中真正拉开差距的地方。光会跑个 trtexec 可不够。

4.1 插件开发：当标准算子不够用时

TensorRT 内置了上百个算子，但现实世界里的模型，总有些“奇奇怪怪”的操作。比如自定义的激活函数、特殊的池化逻辑、或者某些论文里的新结构。这时候，你就得自己写插件了。

为什么需要插件？

模型里有 TensorRT 不支持的算子
你想把多个小算子融合成一个，减少 kernel launch 开销
你需要极致的性能优化，手写 CUDA kernel

我个人习惯把插件开发分成三步：定义、实现、注册。

4.1.1 插件接口与生命周期

TensorRT 的插件体系，说白了就是让你继承几个基类。核心接口是 IPluginV2DynamicExt（动态 shape 场景）或 IPluginV2IOExt（静态 shape）。

嗯，这里要注意：TensorRT 8.x 之后，官方推荐用 IPluginV2DynamicExt。为什么？因为动态 shape 是主流，你写一次插件，输入从 1x3x224x224 变成 1x3x416x416，不用改代码。

一个插件生命周期大概是这样：

构造：解析参数，比如 kernel size、padding
getOutputDimensions：告诉 TensorRT 输出张量的 shape
supportsFormatCombination：声明你支持的数据类型和布局（比如 FP16 + NHWC）
configurePlugin：拿到输入输出的具体信息，做初始化
enqueue：真正的执行函数，里面写你的 CUDA kernel
clone/destroy：拷贝和销毁

我在项目中遇到过最坑的一件事：enqueue 函数里忘了同步流，结果推理结果全是错的。排查了整整两天。记住，enqueue 是异步的，你提交完 kernel 就得返回，但 必须保证 kernel 在流上正确排队。

4.1.2 一个简单的插件示例：LeakyReLU

咱们写个最简单的 LeakyReLU 插件。别看它简单，但麻雀虽小五脏俱全。

class LeakyReluPlugin : public IPluginV2DynamicExt {
public:
    LeakyReluPlugin(float alpha) : mAlpha(alpha) {}

    // 获取输出维度
    DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, 
                                   int nbInputs, IExprBuilder& exprBuilder) override {
        return inputs[0]; // 输入输出 shape 一样
    }

    // 支持 FP32 和 FP16
    bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, 
                                    int nbInputs, int nbOutputs) override {
        return inOut[pos].type == DataType::kFLOAT || 
               inOut[pos].type == DataType::kHALF;
    }

    // 执行函数
    int enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc,
                const void* const* inputs, void* const* outputs, 
                void* workspace, cudaStream_t stream) override {
        int n = 1;
        for (int i = 0; i < inputDesc[0].dims.nbDims; i++) {
            n *= inputDesc[0].dims.d[i];
        }
        // 调用 CUDA kernel
        leakyReluKernel<float><<<(n+255)/256, 256, 0, stream>>>(
            (float*)inputs[0], (float*)outputs[0], n, mAlpha);
        return 0;
    }

    // 注册插件
    REGISTER_TENSORRT_PLUGIN(LeakyReluPluginCreator);
private:
    float mAlpha;
};

避坑指南：我曾经在 supportsFormatCombination 里只写了 FP32，结果模型用 FP16 跑的时候，插件直接报错。后来我养成了习惯：插件默认支持 FP32 和 FP16，除非你有特殊理由。

4.2 多流执行：把 GPU 用满

你想想看，一个 GPU 上有几十个 SM（流多处理器），如果你只用一个 CUDA 流，那大部分 SM 都在闲着。多流执行，说白了就是让 GPU 同时干好几件事。

什么时候用多流？

同时处理多个请求（比如在线服务）
一个模型里有多条独立的计算路径
你想把预处理、推理、后处理流水线化

4.2.1 多流的基本用法

TensorRT 的 enqueueV2 函数接受一个 cudaStream_t 参数。你只需要创建多个流，然后分别提交任务。

cudaStream_t streams[4];
for (int i = 0; i < 4; i++) {
    cudaStreamCreate(&streams[i]);
}

// 为每个流创建独立的 context
ICudaEngine* engine = ...;
IExecutionContext* contexts[4];
for (int i = 0; i < 4; i++) {
    contexts[i] = engine->createExecutionContext();
}

// 多流并行推理
for (int i = 0; i < 4; i++) {
    contexts[i]->enqueueV2(buffers[i], streams[i], nullptr);
}

// 等待所有流完成
for (int i = 0; i < 4; i++) {
    cudaStreamSynchronize(streams[i]);
}

注意：每个流必须有自己的 IExecutionContext，不能共享。因为 context 里保存了推理过程中的中间状态。我见过有人用一个 context 在多个流上调用 enqueue，结果数据全乱了。

4.2.2 性能调优：流数不是越多越好

我做过一个测试：在 A100 上跑 ResNet-50，单流延迟 2ms，吞吐 500 QPS。改成 4 个流，吞吐到了 1800 QPS，延迟只涨到 2.3ms。但改成 16 个流，吞吐反而降了。

为什么会这样？因为 GPU 的并发能力是有限的。每个流都会占用一些资源（寄存器、共享内存），流太多反而导致资源争抢。

我的建议：

小模型（如 ResNet-50）：2-4 个流
大模型（如 BERT-Large）：1-2 个流
用 nvidia-smi 看 GPU 利用率，低于 80% 就加流

4.3 与 Triton Server 集成：从模型到服务

模型部署的最后一公里，就是把它变成一个服务。NVIDIA Triton Inference Server 是目前最流行的方案之一。它支持 TensorRT、ONNX、PyTorch 等多种后端。

说白了，Triton 就是个“模型路由器”。你给它一个请求，它帮你选模型、分配 GPU、管理并发、返回结果。

4.3.1 模型仓库结构

Triton 要求模型按特定目录结构存放：

model_repository/
├── resnet50/
│   ├── 1/
│   │   └── model.plan          # TensorRT 引擎文件
│   └── config.pbtxt            # 模型配置
├── bert/
│   ├── 1/
│   │   └── model.plan
│   └── config.pbtxt

config.pbtxt 是核心配置文件，告诉 Triton 模型的输入输出、最大 batch size、动态 shape 等信息。

name: "resnet50"
platform: "tensorrt_plan"
max_batch_size: 32
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [3, 224, 224]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [1000]
  }
]

4.3.2 动态 batching：Triton 的杀手锏

Triton 最让我喜欢的功能就是动态 batching。你想想看，如果每个请求单独推理，GPU 利用率很低。Triton 会把多个请求攒起来，凑成一个 batch 再送进 TensorRT。

配置起来很简单：

dynamic_batching {
  preferred_batch_size: [4, 8, 16]
  max_queue_delay_microseconds: 100
}

max_queue_delay_microseconds 是关键参数。设得太小，batch 凑不齐；设得太大，延迟会变高。我一般从 100 微秒开始调，看实际效果。

  经验之谈：我曾经在一个 OCR 服务里，把动态 batching 的延迟从 100 微秒调到 500 微秒，吞吐提升了 3 倍，但 P99 延迟只涨了 20%。这个 trade-off 非常划算。

4.3.3 集成自定义插件

如果你的 TensorRT 引擎里用了自定义插件，Triton 怎么知道？答案是：把插件编译成共享库，让 Triton 加载。

步骤很简单：

编译插件为 .so 文件（比如 libmyplugins.so）
启动 Triton 时加上 --backend-config=tensorrt,plugins=libmyplugins.so
Triton 加载引擎时，会自动注册你的插件

嗯，这里有个坑：插件库的 CUDA 版本必须和 Triton 内置的 CUDA 版本一致。否则加载时会报符号找不到的错误。我因为这个折腾过一整个下午。

4.4 本章知识体系

下面这张图，帮你理清本章的核心逻辑：

从这张图你能看到，插件开发解决的是“能不能跑”的问题，多流执行解决的是“跑得快不快”的问题，而 Triton 集成解决的是“怎么给别人用”的问题。三者层层递进，缺一不可。

好了，这一章的内容就到这里。插件开发、多流执行、Triton 集成，这三板斧你掌握了，基本就能应对大部分生产环境的部署需求了。下一章我们聊聊更高级的话题——模型压缩与知识蒸馏，敬请期待。

📚 课程目录

4、NVIDIA TensorRT（下）：插件开发、多流执行、与Triton Server集成实践

4.1 插件开发：当标准算子不够用时

4.1.1 插件接口与生命周期

4.1.2 一个简单的插件示例：LeakyReLU

4.2 多流执行：把 GPU 用满

4.2.1 多流的基本用法

4.2.2 性能调优：流数不是越多越好

4.3 与 Triton Server 集成：从模型到服务

4.3.1 模型仓库结构

4.3.2 动态 batching：Triton 的杀手锏

4.3.3 集成自定义插件

4.4 本章知识体系