一文搞懂 AI 推理部署:算子融合、TVM、TensorRT 与 ONNX Runtime

写给刚接触 AI Infra 的你:为什么模型训练完不能直接上线?中间那一堆"编译器""运行时"到底是什么关系?这篇文章带你从头理清。

引言:一个模型的上岗之路

假设你用 PyTorch 训好了一个图像分类模型,准确率不错。现在老板说:“把它部署到线上服务,要求单次推理延迟低于 10ms。”

你兴冲冲地把模型加载进 Flask 接口,一压测——延迟 50ms,GPU 利用率才 30%。问题出在哪?

原来,训练好的模型和"能高效跑在生产环境的模型"之间,隔着一整个优化世界。这个世界的主角,就是本文要讲的几个概念:算子融合(Operator Fusion)、TVM、TensorRT、ONNX Runtime。

在展开之前,先记住一张全景图,后文会反复回到它:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
┌─────────────────────────────────────────────┐
│  训练框架层:PyTorch / TensorFlow / JAX        │
└──────────────────┬──────────────────────────┘
                   │ 导出
┌──────────────────▼──────────────────────────┐
│  模型交换格式:ONNX(通用中间表示)              │
└──────────────────┬──────────────────────────┘
                   │
┌──────────────────▼──────────────────────────┐
│  AI 编译器 / 推理运行时层:                     │
│    TVM        TensorRT      ONNX Runtime    │
└──────────────────┬──────────────────────────┘
                   │
┌──────────────────▼──────────────────────────┐
│  内核库层:cuDNN / cuBLAS / CUTLASS / 厂商 SDK │
└──────────────────┬──────────────────────────┘
                   │
┌──────────────────▼──────────────────────────┐
│  硬件层:NVIDIA GPU / CPU / NPU / FPGA        │
└─────────────────────────────────────────────┘

一、算子融合:推理优化的第一课

1.1 为什么模型跑得慢?

深度神经网络本质上是一张"算子(Operator)的计算图"。算子就是卷积、矩阵乘、ReLU、BatchNorm 这些基本运算单元。

问题在于:现代 GPU 的算力增长远快于内存带宽增长,导致"算得快,喂不饱"。当模型执行到一串小算子时,会出现两个经典瓶颈:

瓶颈一:kernel 启动开销。 GPU 上每执行一个算子,都要经历一次 kernel 启动、同步、读写显存的过程。如果一张图里有 100 个连续的逐元素小算子,光是启动开销就可能比计算本身还贵。

瓶颈二:无效的显存读写。 看这个最常见的组合:

1
Y = Conv(X)  →  Z = BN(Y)  →  A = ReLU(Z)

按最朴素的方式执行,每一步都要:把结果写回显存 → 下一个算子再从显存读出来。中间结果 Y 和 Z 明明只"活"了一瞬间,却要经历两次完整的显存往返。对于带宽敏感的场景,这是巨大的浪费。

1.2 融合:把三趟运输变成一趟

算子融合的思路非常直观——把多个算子合并成一个 kernel,数据在寄存器/共享内存里直接传递,不落回显存:

1
2
融合前:Conv → [写显存] → BN → [写显存] → ReLU    (3 次 kernel,2 次显存往返)
融合后:Conv + BN + ReLU 单个 kernel              (1 次 kernel,0 次中间往返)

经典的 Conv + BN + ReLU 三层融合,在真实场景中可以带来 30% 以上的延迟降低。而在大模型推理中赫赫有名的 FlashAttention,本质上就是一次极致的手工算子融合——把 Attention 中的矩阵乘、缩放、Softmax、再矩阵乘融合成一个访存高度优化的 kernel,避免了巨大中间矩阵的显存读写。

1.3 训练侧和推理侧,谁在做融合?

一个常见误区是认为算子融合只是"推理侧的活"。实际上两边都做,但关注点不同:

维度 训练侧 推理侧
核心目标 提升 GPU 利用率、省显存(前向 + 反向都要跑) 降低延迟、提升吞吐、支持更高并发
融合范围 前向图 + 反向图,还要考虑与梯度 AllReduce、优化器更新的融合 只需前向图,可激进删减(推理不需要 Dropout 和反向计算)
硬约束 必须保证融合后可求导 无梯度约束,还能顺手做量化(FP16/INT8)
典型工具 XLA、torch.compile、NVIDIA Apex TensorRT、ONNX Runtime、vLLM
典型案例 Norm + 残差 + 激活融合;优化器逐元素更新融合 Conv + BN + ReLU 融合;FlashAttention

一句话:训练侧的融合要"小心翼翼保正确",推理侧的融合可以"放开手脚拼性能"。随着 torch.compile 这类 JIT 编译技术普及,一套框架同时服务训练和推理越来越常见,但两侧的技术侧重点依然存在。

二、三大主角登场:TVM、TensorRT、ONNX Runtime

理解了算子融合,我们来看执行这类优化的工具。先给三者的"身份卡":

2.1 TVM:学院派的"全能编译器"

  • 出身:华盛顿大学陈天奇团队,2017 年开源,Apache 顶级项目
  • 定位:端到端深度学习编译器栈
  • 最大特点:硬件无关。CPU、GPU、NPU、DSP、FPGA,甚至国产芯片,都能通过它的编译流程生成代码

TVM 的思路借鉴了 LLVM:把模型看作计算图,经过多层中间表示(IR)的转换——图优化、算子融合、循环分块、缓存优化、自动调优——最终生成针对目标硬件的机器码。你甚至可以用它的调度语言(TE/TIR)手写算子优化策略,或用 AutoTVM/Ansor 让机器自动搜索最优实现。

适合谁:需要部署到非 NVIDIA 硬件、做算子级研究、适配自研芯片的团队。灵活度最高,上手门槛也最高。

2.2 TensorRT:NVIDIA 阵营的"性能天花板"

  • 出身:NVIDIA,2016 年发布,2024 年宣布开源
  • 定位:NVIDIA GPU 专属的推理优化器和运行时
  • 最大特点:极致性能 + 深度绑定 NVIDIA

TensorRT 拿到 ONNX 模型后,会做一整套硬件感知的优化:层融合、精度校准(FP32/FP16/INT8/FP8 甚至 FP4)、针对具体 GPU 架构的 kernel 自动调优、CUDA Graph 消除启动开销。最终产出一个高度定制化的 .engine 文件——这份"编译产物"和你的 GPU 型号是绑定的,换一张卡就要重新编译。

在 MLPerf 推理基准测试中,基于 TensorRT 的方案长期是 NVIDIA 平台的标配答案。大模型时代,它的兄弟项目 TensorRT-LLM 更是 LLM 推理优化的主力工具。

适合谁:硬件锁定 NVIDIA、追求最低延迟最高吞吐的生产环境(数据中心、自动驾驶、大模型服务)。

2.3 ONNX Runtime:微软的"通用胶水层"

  • 出身:微软,开源(MIT 协议)
  • 定位:跨平台高性能推理引擎
  • 最大特点:兼容性 + 编排能力

ORT 有两个核心设计:

其一,ONNX 标准接口。 ONNX(Open Neural Network Exchange)是一种开放的模型交换格式,好比 AI 界的"通用翻译语言"——PyTorch、TensorFlow 训练的模型都能转成 ONNX 格式,任何支持 ONNX 的运行时都能执行它。这解决了"训练框架五花八门、部署环境五花八门"的对接难题。

其二,Execution Provider(EP)架构。 这是 ORT 最精妙的设计,也是理解它与另外两者关系的关键,下一节细讲。

适合谁:需要一套代码跨 CPU/GPU/多厂商硬件部署、或混合硬件环境的团队。

三、层级关系:谁在谁之上?

这是最容易混淆的部分。先纠正一个常见误解:三者不是简单的上下游流水线,而是"同层协作、角色分工"的结构——TVM 和 TensorRT 是平级的优化执行引擎,ORT 则常常站在它们之上,作为统一的推理入口和调度者。

3.1 一张图看清协作关系

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
                PyTorch / TensorFlow(训练框架)
                        │
                        ▼ 导出
                    ONNX 模型
                        │
                        ▼ 加载
                  ┌───────────────┐
                  │ ONNX Runtime  │  ←── 通用图优化:常量折叠、
                  │  (编排者)     │       冗余节点消除
                  └───┬───┬───┬───┘
                      │   │   │
        ┌─────────────┘   │   └─────────────┐
        ▼                 ▼                 ▼
┌──────────────┐  ┌──────────────┐  ┌──────────────┐
│ TensorRT EP  │  │  CUDA EP     │  │  CPU EP      │
│ (子图交给     │  │ (cuDNN 等)   │  │ (MLAS)       │
│  TensorRT)   │  └──────────────┘  └──────────────┘
└──────────────┘
        │
        ▼
   NVIDIA GPU

ORT 的 Execution Provider 机制是这样工作的:

  1. ORT 先对整张模型图做通用优化(常量折叠、算子融合、死代码消除);
  2. 根据 EP 的算子支持情况,把计算图切分成若干子图;
  3. TensorRT 能处理的子图交给 TensorRT 编译加速;不支持的算子回退到 CUDA EP 或 CPU EP 执行;
  4. 跨 EP 的数据拷贝由 ORT 自动插入和管理。

对使用者来说,你只跟 ORT 的 Session 打交道,背后的子图切分和调度完全透明——这就是"编排者"的角色。

3.2 用传统编译栈类比

如果熟悉编译原理,这张对照表能让你秒懂:

AI 推理栈 传统编译栈类比
PyTorch / TensorFlow 模型 C++ / Rust 源代码
ONNX 格式 LLVM IR(中间表示)
TVM / TensorRT 后端编译器(生成目标平台机器码)
ONNX Runtime 运行时调度器(类似 JVM:解释 + JIT 触发)
cuDNN / cuBLAS 数学库(类似 BLAS / MKL)

核心区别在于:TVM 和 TensorRT 负责"把子图编译到极致性能",ORT 负责"决定何时、何地、由谁执行哪段代码"。前者是编译器,后者更偏运行时编排。

3.3 三者定位对比速查表

维度 TVM TensorRT ONNX Runtime
角色 全能编译器 NVIDIA 专属优化引擎 推理编排运行时
硬件支持 几乎所有硬件 仅 NVIDIA GPU 多后端(CUDA/ROCm/OpenVINO/CPU/NPU…)
上游输入 ONNX / 多框架前端 ONNX(也可从框架直接转换) ONNX / 直接接 PyTorch
输出形态 编译好的代码模块 序列化的 .engine 文件 Session 对象,直接收发 tensor
灵活度 最高(可自定义算子、后端) 低(绑定 NVIDIA 生态) 中(靠 EP 机制扩展)
上手难度 高 中 低

四、实战选型:什么场景用什么组合?

理解了层级关系,选型思路就清晰了。三条典型路线:

路线一:NVIDIA 独占,追求极致性能

1
PyTorch → ONNX → TensorRT 编译 → .engine 文件 → GPU 部署

跳过 ORT 的编排层,让 TensorRT 直接接管整张图,性能最优。适合硬件环境单一、性能敏感的场景(如自动驾驶、线上大模型服务)。

代价:.engine 文件与 GPU 型号绑定,换卡需重新编译;对不支持的算子需要自己写 plugin。

路线二:混合/跨平台环境,一套代码走天下

1
2
3
4
PyTorch → ONNX → ONNX Runtime(统一入口)
                    ├─ NVIDIA 卡 → TensorRT EP 加速
                    ├─ Intel 卡  → OpenVINO EP
                    └─ 无 GPU    → CPU EP 兜底

以 ORT 为统一推理接口,NVIDIA 卡上通过 TRT EP 享受接近原生的性能,其他硬件走对应后端。适合云厂商、SaaS 服务、桌面软件等硬件环境多样的场景。

路线三:非标准硬件或深度定制

1
任意模型 → TVM 编译(自定义算子 + 硬件后端)→ 目标芯片

绕开 ONNX 标准算子集的限制,对任意计算图做端到端编译。适合自研芯片适配、边缘设备、算子研究。

一个现实中常见的真相:大型 AI Infra 中三者常常配合出现而非二选一——ORT 管调度和兼容,TensorRT 管 NVIDIA 上的性能,TVM 管长尾硬件和定制算子。

五、总结:一张知识卡片带走全文

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
┌────────────────────────────────────────────────────┐
│ ① 算子融合                                          │
│   · 把多个小算子合并成一个大 kernel,消除显存往返       │
│   · 训练侧做(保可微),推理侧也做(拼性能)             │
│   · FlashAttention 本质就是一次极致的手工融合           │
│                                                    │
│ ② 三大工具的层级关系                                  │
│   · TVM     → 全能编译器,硬件无关,灵活度最高          │
│   · TensorRT → NVIDIA 专属引擎,性能天花板             │
│   · ORT     → 运行时编排者,通过 EP 机制调度上面两位     │
│   · TVM 和 TensorRT 平级(编译后端),ORT 站在它们之上    │
│                                                    │
│ ③ 选型口诀                                          │
│   · 锁定 NVIDIA 拼极致 → PyTorch → ONNX → TensorRT   │
│   · 跨硬件统一部署    → ONNX Runtime + 各类 EP        │
│   · 非标硬件/深度定制  → TVM                          │
│   · 三者不互斥,大型系统里常常配合使用                    │
└────────────────────────────────────────────────────┘

最后留一个思考题:如果让你部署一个 70B 的大语言模型到由 8 张 NVIDIA H100 组成的服务器上,同时还要支持一部分客户使用国产 NPU,你会怎么组合这三件工具?

欢迎在评论区聊聊你的方案。下一篇,我们聊聊大模型推理特有的优化技术——KV Cache、PagedAttention 和连续批处理。