TensorRT-LLM架构与核心概念

聊到TensorRT-LLM的架构，我习惯把它拆成三个层面来看：前端、后端、运行时。这三个东西各司其职，配合起来才能把一个大模型跑得又快又稳。你想想看，一个模型从训练框架出来，到最终在GPU上跑推理，中间要经历多少步骤？嗯，这就是我们今天要讲清楚的事。

整体架构：前端、后端、运行时

先说前端。前端主要负责模型解析和中间表示（IR）的生成。说白了，就是把你从PyTorch、TensorFlow或者ONNX导出的模型，转成TensorRT-LLM能理解的数据结构。我个人习惯把这一步叫做“翻译官”——它把各种框架的方言，统一成一种内部语言。

后端呢，就是做优化和编译的。它拿到前端生成的IR后，会做算子融合、内存复用、量化等一系列操作。我在项目中遇到过最头疼的事，就是后端优化时发现某个算子不支持，得自己写Plugin。嗯，这个后面会细说。

运行时负责最终的执行。它管理GPU资源、调度推理请求、处理动态形状。说白了，就是让模型真正跑起来的那一层。我记得有一次线上服务抖动，查了半天发现是运行时内存池配置不合理，调整之后吞吐直接翻倍。

核心概念：Engine、Plugin、Builder、Runtime

这四个概念，是TensorRT-LLM的基石。我一个个说。

Engine（引擎）

Engine是优化后的模型产物，说白了就是一个可执行的推理计划。它包含了所有算子的CUDA kernel、内存分配策略、执行流水线。你训练好的模型，经过Builder编译后，就变成了一个Engine文件（.engine）。这个文件可以直接加载到Runtime里跑推理。

我曾经犯过一个错：在A100上编译的Engine，拿到V100上跑，直接报错。因为不同架构的GPU，指令集和寄存器数量都不一样。所以记住：Engine是硬件绑定的。

Plugin（插件）

Plugin是TensorRT-LLM的扩展机制。当内置算子不够用时，你可以自己写Plugin。比如FlashAttention、Fused MoE这些，都是通过Plugin实现的。我建议你把Plugin理解成“乐高积木”——官方给你基础块，你自己拼出想要的形状。

避坑指南：我曾经写过一个自定义Plugin，性能反而比内置算子还差。后来发现是共享内存没用好。写Plugin时，一定要关注寄存器压力和内存带宽。

Builder（构建器）

Builder是编译过程的指挥官。它接收模型定义和配置参数，输出Engine。配置参数包括：精度（FP16/INT8/FP8）、批量大小范围、序列长度范围、优化策略等。

嗯，这里要注意：Builder的配置直接影响最终性能。我习惯的做法是，先用默认配置跑一遍，然后逐步调整批量大小和精度，找到最优组合。

Runtime（运行时）

Runtime负责加载Engine并执行推理。它提供C++和Python API，支持同步/异步执行、多流并发、动态形状。说白了，Runtime就是那个“按下启动按钮”的角色。

推理流程：模型加载 -> 图优化 -> 编译 -> 执行

整个推理流程，我习惯画成一条流水线。下面这张SVG图，就是我对这个流程的理解。

这张图把流程分成了四个阶段。我一个个拆开讲。

第一步：模型加载

这一步，Builder读取你的模型文件。支持PyTorch（通过torch.export导出）、ONNX、以及TensorRT-LLM自定义的模型定义。我个人建议优先用PyTorch导出，因为ONNX有时会丢失一些动态控制流信息。

代码示例：

from tensorrt_llm import Builder
builder = Builder()
# 加载PyTorch模型
network = builder.create_network()
parser = builder.create_parser(network)
parser.parse("model.onnx")  # 或者用PyTorch的torch.export

第二步：图优化

Builder会对计算图做一系列优化。包括：

• 算子融合：把多个小算子合并成一个，减少kernel启动开销。比如把LayerNorm和后面的Add融合。
• 常量折叠：把编译时就能算出来的常量提前算好，省掉运行时计算。
• 死代码消除：去掉那些不会被用到的计算节点。

我在项目中遇到过一个问题：图优化后，某个算子的精度从FP16变成了FP32，导致显存暴涨。后来发现是优化策略里有个“安全模式”开关，关掉就好了。

第三步：编译

这一步是最耗时的。Builder会根据你的配置（精度、批量大小、序列长度等），为每个算子生成最优的CUDA kernel。这个过程叫“tactic selection”——说白了，就是试各种实现方式，选最快的那个。

配置示例：

config = builder.create_builder_config()
config.set_memory_pool_limit(8 * 1024 * 1024 * 1024)  # 8GB显存限制
config.set_precision(Precision.FP16)
config.set_max_batch_size(32)
config.set_max_seq_len(2048)

engine = builder.build_engine(network, config)

嗯，这里要注意：编译时间可能很长，尤其是大模型。我见过一个70B参数的模型，编译花了40分钟。所以生产环境里，我习惯把编译好的Engine存成文件，部署时直接加载。

第四步：执行

Runtime加载Engine，开始推理。支持同步和异步两种模式。异步模式适合高并发场景，可以同时处理多个请求。

执行示例：

from tensorrt_llm import Runtime
runtime = Runtime()
engine = runtime.deserialize_engine(engine_data)
context = engine.create_execution_context()

# 准备输入输出
inputs = {"input_ids": input_tensor}
outputs = {"output": output_tensor}

# 执行推理
context.execute_v2(inputs, outputs)

总结一下

TensorRT-LLM的架构，说白了就是“三段式”：前端解析、后端优化、运行时执行。核心概念就四个：Engine是产物，Plugin是扩展，Builder是工具，Runtime是执行者。推理流程就是“加载->优化->编译->执行”这条流水线。

我个人觉得，理解这个架构的关键在于“分层”思维。每一层只关心自己的事，层与层之间通过标准接口通信。这样设计的好处是：你可以替换前端（比如从PyTorch换成JAX），或者替换后端（比如从FP16换成INT8），其他层不用动。

嗯，今天就讲到这里。下一章我们会深入Builder的配置细节，聊聊怎么调参才能让模型跑得又快又省显存。

---

公众号:蓝海资料掘金营，微信deep3321