自定义插件开发基本流程 TensorRT 本身只是一个推理优化框架，它需要从其他训练框架的模型中解析出网络结构，因此常常会遇到一些 TensorRT 不支持的算子，这时我们就需要用插件扩展功能来自己实现这些算子。 简单来说，开发一个 TensorRT 自定义算子插件的流程是： 首先明确要开发的算子，最好是先用 CUDA 实现； 开发插件类，实现 IPluginV2DynamicExt 接口，在这里调用前面实现的算子； 开发插件 Creator 类，实现 IPluginCreator 接口，用于创建插件实例，然后注册该 Creator 类； 编译插件项目，生成动态链接库； 在构造 engine 之前，首先加载上一步编译出来的插件动态链接库，在构造 engine...

类 GPT 模型在自回归生成过程中，每一次推理只生成一个 token，然后将生成的 token 拼接到输入序列，再进行下一次推理，直到生成 eos 或者达到最大长度。在这个过程中，每次推理的计算量都会增加，假设每个 token 的计算量为 1 个单位，则在 context 的长度为 $L$，生成序列长度为 $M$ 的情况下，总的计算量为 $(M + 1) L + \frac{M (M + 1)} 2$ 个单位。 KV Cache 的使用者认为上述过程的大部分运算是可以避免的，比如每个 token 输入模型之后，内部的 key 和 value 可以缓存起来，后续就不必再计算了。也就是说除了第一次推理是将 整个prompt 输入模型，后续每次推理都是将新生成的 token 作为输入，这样...

fp16 量化 如果显卡支持 fp16 运算加速，那么使用 fp16 量化能显著提升模型的推理速度，且由于 fp16 的表示范围为 (-65504, +65504)，一般都能包含模型的权重数值范围，所以直接截断对精度的影响很小。因此在 TensorRT 中实现 fp16 量化的方法相当简单，只需要在构建 engine 时添加一行配置即可 config = builder.create_builder_config() config.set_flag(trt.BuilderFlag.FP16) int8 量化 int8 量化能进一步压缩模型体积以及提升推理速度，但是 8bit 整数能表示的范围只有 -127~128，模型权重值一般都超过了这个范围。因此需要对权重进行 scale...

安装 TensorRT 的注意事项 TensorRT 的安装方式在官方安装导引 里已经有详细的说明了，这里提下需要注意的地方，首先是尽量保证 TensorRT 的安装方式 和 CUDA 的安装方式相同，否则可能会出现找不到 CUDA 的情况，比如在 Ubuntu 中，如果 CUDA 不是通过 deb 包安装的，后面用 deb 安装 TensorRT 就会报错找不到 CUDA 依赖。 从 PyTorch 导出 ONNX 模型 ONNX 是一种通用的神经网络模型交换格式，TensorRT 有专门的 ONNX 解析器，可以解析 ONNX 模型并创建 TensorRT 自有的网络结构并做后续的优化工作。因此，为了让 TensorRT 优化我们使用 PyTorch 训练的模型，可以先将 P...

算术强度的计算问题 前面我们在介绍算术强度的时候，非常理想化的把内存读取总量等价成了参与运算的数据总量。但实际上，由于寄存器数量限制，没法完全容纳全部的数据量，每个数据往往需要读很多次，因此实际的算术强度只可能比理论值更低。 图 1. 数据流动示意图 上图右边展示了实际的数据流动过程，每次加载只能读取一部分数据到寄存器，这部分数据算完之后，再读下一部分，由于算法的原因，两次计算之间有可能需要用到一些相同的数据，因此存在重复加载的现象。 共享内存的作用 通过上一节的分析我们可以看到，寄存器容量的限制会直接导致算法的实际算术强度达不到理论值，会对程序的性能造成不利影响。在没法无限制增加寄存器数量的情况下，一种折中的方案就是用一块速度比普通内存块，容量比寄存器大的特殊区域来做中间...