最近被问到一个很直接的问题:TensorRT 不是已经在 NVIDIA GPU 上把推理压榨到接近极限了吗,为什么还要做 MLIR-TensorRT 这种听起来像“底层套底层”的东西?我反手抛了个问题回去:如果你的模型里同时有 PyTorch 导出的子图、ONNX 转换过来的子图,还有团队自己用 C++ 堆出来的推理逻辑,能不能把它们编到同一份可分析的中间表示里,再做统一调度?对方愣了一下。在传统 TensorRT 的使用语境下,这个问题确实不好答,因为 TensorRT 更像一个高吞吐的加速库,而不是一座能让不同前端自由汇入的编译桥梁。NVIDIA 做 MLIR-TensorRT,本质上就是在把推理编译器这件“工厂流水线”下沉到 MLIR 生态里,让 TensorRT 不再是黑盒终端,而是一段可被编排、可被改写、可被观察的编译后端。
很多做部署优化的同学第一次看到 MLIR-TensorRT 这个项目名,都会产生一个疑问:它到底是 MLIR 插件,还是 TensorRT 的替代品,抑或只是 NVIDIA 给 HPC/自动驾驶团队准备的“内部玩具”?这篇文章我会从我的实际认知出发,把项目背景、核心概念、工作链路、横向对比和上手路径一次讲透。文章不会有太多需要你“闭眼照着做”的 magic command,反而会花更多篇幅解释为什么这套设计值得你关注,以及哪些地方目前仍然有坑。
1. TensorRT 这些年一直是“加速库”,为什么现在要谈编译器
1.1 加速库的“舒服”是有代价的
如果你过去两年一直在用 TensorRT 做推理部署,一定熟悉这条流水线:把 PyTorch 模型导出 ONNX,然后用 trtexec 或者 TensorRT Python API 构建 engine,再用 runtime 加载 engine 去执行。这套流程对绝大多数单模型场景都非常省心,TensorRT 会帮你做常量折叠、层融合、精度选择、kernel 自动调优,甚至量化校准。
但它的“省心”其实是建立在几个隐藏前提上的:你的模型结构相对完整、你的前端格式足够标准、你只在一个 GPU 环境上长期服役。一旦需求变成“多个模型拼接成一个推理服务”“某个子结构想用自定义 kernel 替换”“想对算子融合规则做细粒度控制”,传统 TensorRT 的 API 模型就开始暴露边界。你可以通过 plugin 塞自定义算子,也可以把多个 engine 串起来,但这些操作更像是在一座已经很满的仓库里硬塞新货架,而不是重新设计仓库的货架体系和通行路径。
从工程视角看,TensorRT 本质上是一个“链接器加运行时”:它把网络定义变成可执行的二进制 engine,再把输入输出暴露给你。问题不在于它做得不够快,而在于它没有给你提供一个足够开放的中间层来做图级别的代码变换。过去大家总说“ONNX 是模型界的通用语言”,可实际上 ONNX 在部署链路上更像一个传输格式,它适合做模型交换,却很难在此基础上推理大规模的组合优化。于是就需要一个更稳定的 IR,而 IR 恰恰是现代编译器的基础设施。
1.2 一个 IR 能带来的自由:可看、可改、可复用
如果你熟悉 LLVM,会知道 IR 为编译器带来了三个关键能力:可视化、可改写、可复用。模型推理也需要类似的东西。当我们把一张推理图从框架里抽出来之后,第一步不是急着把它变成 kernel,而是先让它进入到一种可以被人和工具共同阅读的表示里。对这张图做算子融合、内存复用和并行调度,才算真正进入“编译”环节。
MLIR 的厉害之处在于,它不是一套固定的 IR,而是一套 IR 框架。你可以定义自己的 dialect,把 TensorRT 的层、CUDA 的 kernel、显存分配模型都表达成方言里的 op。接着就可以复用 MLIR 生态里已有的通用 pass,例如常量传播、死代码消除、循环不变量外提等,也可以针对自己的 dialect 写自定义 pass,把若干 TensorRT 层融合成复合 op。
用生活化的类比来说,传统方式是你把一段中文直接交给一位资深翻译,让他直接翻成英文,翻得又快又好,但你无法让另一位编辑在中间阶段对“中文到英文的初稿”做风格调整。MLIR 的逻辑是先翻译成一种规范且可编辑的中间语言,再由另一套编译规则把它变成最终机器码。每个环节都能暴露在工具链之下,被检查、被优化、被复用,这才是编译器思维的真正价值。
1.3 “融合 MLIR 生态”不是把 TensorRT 重写一遍
NVIDIA 的 MLIR-TensorRT 项目并不是要把 TensorRT 的每一行实现都搬到 MLIR 里,然后抛弃原有框架。事实恰恰相反,它希望把 TensorRT 放在更加开放的工具链“下游”,让 MLIR 生态中的各种前端和优化 pass 都能更顺畅地触达 TensorRT。
具体来说有三层含义。第一层,输入侧:它希望接收来自不同 MLIR 前端的模型表示,比如把 StableHLO、TorchScript 或 ONNX 转换到统一的 MLIR 模块里。第二层,优化侧:图级别的优化和调度不再完全封闭在 TensorRT 内部,而是由用户可以看见的 MLIR pass 来完成,哪里融合、哪里拆开、保持多少精度,都可以插一手。第三层,输出侧:真正落到 GPU 上执行的仍然是 TensorRT engine 或等价的运行时组件,保证 NVIDIA 多年积累的 kernel auto-tuning 和算子库能力不被浪费。
一句话说清楚:这不是“用 MLIR 重写 TensorRT”,而是“把 TensorRT 纳入到一个更大的编译器体系里”。这样既保住了 NVIDIA 在底层 kernel 上的积累,又把优化过程的解释权交还给开发者。听上去确实很理想,但也要理性看待它目前的成熟度。
2. Dialect、Pass、Lowering:进入 MLIR-TensorRT 前需要先补的底子
2.1 Dialect 就像同一份代码在不同圈子里说不同方言
很多第一次接触 MLIR 的人会被 dialect 这个概念挡住。其实它并没有那么玄。如果你去过不同地方的团队开技术评审会,会发现大家虽然都在讲“性能优化”这件事,但用词完全不同:算法工程师说 batch size 和 latency,框架工程师说 graph capture,CUDA 工程师说 kernel 和 occupancy。MLIR 中的 dialect 本质上就是给 op 定义了不同的“词汇表”。
同一个卷积加偏置加 ReLU 的小结构,在 ONNX 的表示里可能是Conv -> Add -> Relu三个 op,每个 op 都很干净;在 TensorRT 的表示里它可能已经是一个融合过的 “conv bias relu” 复合 op;在更接近硬件表示里,它可能还要带上 tensor memory accelerator 的布局要求。MLIR 允许这些不同抽象层次的 op 同时存在于同一个模块中,并通过 pass 把它们逐层降到更具体的方言上。这样写通用优化的人不必关心硬件细节,写硬件细节的人也不必操心前端语义。
MLIR 中常见的通用 dialect 其实很多:arith负责加减乘除等基础算术,tensor描述单个张量值,linalg用来表达可复用的算子结构,tosa提供一组面向神经网络算子的稳定接口,stablehlo则是许多框架前端共同瞄准的表示。MLIR-TensorRT 的价值就是在这个“方言谱系”里加入了自己的方言,让模型可以一路顺畅地落到 TensorRT 的地盘上。
2.2 Pass 和 Lowering:谁在什么时候做改写
Pass 是编译器里的“代码改写动作”。最直白的一个例子是算子融合:识别并匹配某几个相邻 op,把它们合并成一个 op,减少中间结果写回显存的开销。还有显存规划:把那些生命周期不重叠的 tensor 安排到同一块显存区域,从而显著降低总显存占用。MLIR 的标准框架会给你一套 pass 管理机制,你可以规定某个 pass 在什么时候跑、跑多少遍、依赖前面哪个 pass 的结果。
Lowering 则是从一个 dialect 切到另一个 dialect 的过程。这个过程不是一个 op 对应一个 op 的简单翻译,而是要把语义、数据类型、内存布局都对齐。比如把一个tensor表示下很优雅的 element-wise add,lower 到带显存 buffer 的形式时,就必须考虑这个 buffer 有多大、对齐方式是什么、是不是可以复用前一个 op 的 buffer。不掌握这层思维,很容易在看 MLIR 输出时觉得它啰嗦且不直观。
完整理解这些概念,再去看 MLIR-TensorRT 的工程结构,会轻松很多:你在命令行工具里看到的那些 pass pipeline、IR dump、内存规划信息,其实都是这套编译框架的标准产物,并不是 TensorRT 单独给你发明的黑话。
2.3 为什么可以用 MLIR 做前端,却把 TensorRT 留在内核层
稍微有编译器经验的读者可能会问:既然我们已经走到 MLIR 了,为什么不干脆把算子也生成到 CUDA 代码,做一条完全开放的 GPU kernel 编译器?答案是成本和风险。
模型推理编译器最难的地方并不在“能不能生成一个正确 kernel”,而在于要在一个巨大的算子空间里找出每个 shape、每个 GPU、每个精度组合下性能最优的实现。TensorRT 的价值恰恰在于它手里有大量已经自动调优过的 kernel,有各种经过验证的融合模板,还有对半精度、BF16、INT8、FP8 等细节的长期打磨。从头写一套能跟它掰手腕的 kernel 编译器,工程量会非常夸张,而且很难在短时间内达到同样高的性能下限。
所以 MLIR-TensorRT 的策略很务实:图优化、IR 层面的分析和跨模型调度用 MLIR 体系来解决;到真正的计算密集算子时,仍然调用 TensorRT 或者其背后的 CUDA 内核库。这样既保证了“能看、能改”的灵活性,又不牺牲 NVIDIA 在算子上积累多年的性能优势。理解这一点,你才不会对这个项目产生不切实际的期待。
3. 从 ONNX 到 Plan:一条推理模型在 MLIR-TensorRT 里的完整路径
3.1 compile 与 runtime:一条离线路、一条执行路
MLIR-TensorRT 的架构可以粗略分成两部分:编译期工具链和运行时组件。编译期那部分负责把前端模型变成一种可执行的“Plan”,运行时则负责加载并执行 Plan。
在离线路线上,你首先需要一个入口把模型读进 MLIR。如果你手里是一个 ONNX 文件,可以把它导入成基于 tensor 运算的 MLIR 模块;如果你从 PyTorch 导出 TorchScript,也可以把 TorchScript 先转成 StableHLO 之类的表示,再进入同一条优化管道。拿到 MLIR 模块之后,一系列 pass 会开始工作:形状推断、转成 TensorRT dialect 或更通用的计算表示、精度规划、算子融合、显存 buffer 规划,最后生成一个带完整执行顺序的 Plan。
在执行路线上,runtime 接收这个 Plan。它不像传统 TensorRT API 那样要求你手握 engine 句柄和 context 手动绑定输入输出,而是可以依据 Plan 里的依赖关系完成显存分配、kernel launch、流同步等操作。对多模型联合推理而言,这个设计比自己在外面手工拼 engine 要安全得多,因为依赖和生命周期都在同一份中间表示里被记录。
3.2 TRT Dialect 与 Plan Dialect:一个描述“怎么算”,一个描述“怎么排”
许多刚接触项目的人会混淆两个核心 dialect,一个是用来描述 TensorRT 层网络的 dialect,另一个是用来描述整个 Plan 的 dialect。你可以把前者理解成“这张图里每个节点计算什么”,把后者理解成“这些节点之间的依赖关系、资源分配和执行对象到底是怎样的”。
TRT Dialect 会拿着神经网络的核心计算逻辑,例如卷积、归一化、全连接、注意力结构等表达成 TensorRT 语义下的 op。这里面的信息足够多,可以支撑编译期将图转化成一个或多个 TensorRT engine 表示。Plan Dialect 则更上层一点,它描述的是最终运行时如何把这些算子组织起来:要用哪几个 CUDA stream、哪些输入输出 tensor 需要做 host 端拷贝、哪个 engine 的 workspace 单独分配、哪个又必须复用上一个 engine 的显存区域。
这两层分开之后,你可以做很多过去做不了的事:单独分析某一段子网络的计算类型、诊断显存碎片、把某个子图替换成自研 plugin,并且不会牵扯到其它无关部分。对需要在同一张卡上跑多个模型、共享显存池的服务端场景来说,这种“描述与控制分离”的设计比传统 TensorRT 手动管理多个 engine 更干净。
3.3 动态 shape、显存规划和多流执行能走到哪一步
动态 shape 历来是 TensorRT 使用里最头大的问题之一。过去你用 TRT 时,通常要显式指定 optimization profile,为不同 batch size 或者输入分辨率单独设范围,非常繁琐。MLIR 生态本身对符号 shape 和可变维度的表达有很成熟的支持,所以在 MLIR-TensorRT 里,一些跟 shape 相关的推断和推理可以被放在更早的阶段完成,并且把“到底为哪些 shape 生成 kernel”变成一个显式的策略决策。
显存规划在传统部署中往往是个“跑起来才知道”的黑洞。而在 MLIR 工作流里,你能看到类似 bufferization 和分配规划留下的痕迹,知道当前 Plan 大概需要多少显存、中间的临时 tensor 放在哪里。多流执行也是一样,Plan 里可以表达哪些 engine 之间没有数据依赖、可以并行跑在不同 CUDA stream 上,调度逻辑不再散落在业务代码里,而是下沉到可检查的 IR 里。
这些能力听起来非常美好,但要注意,它们仍然处在迭代过程中,并不是每个模块都足够稳定。如果你想在生产环境里直接依赖这套机制处理非常极端的动态场景,还是得先做好实验验证和版本锁定。
4. 与 XLA、TVM、Torch-TensorRT 横向对照后,真正多出来的能力边界
4.1 与 XLA 的对照:同样都是 MLIR 思想,切入角度不同
提到 MLIR 和 NVIDIA GPU 推理编译器,自然会想到 XLA。XLA 是一个历史悠久、用编译器思维做模型加速的框架,尤其跟 TensorFlow 和 JAX 的关系极其紧密。它的前端主要是 HLO/StableHLO,优化链路里也大量使用 MLIR,最后通过 GPU backend 生成并行计算逻辑。
XLA 和 MLIR-TensorRT 的差异,与其说是技术栈不同,不如说是“目标边界”不同。XLA 更强调从框架前端到 GPU kernel 生成的全链路打通,它追求让 JAX/TensorFlow 用户能在不切换工具链的前提下获得较高性能。MLIR-TensorRT 则更像一个面向“终极 N 卡推理性能”的专门通道:它不试图用一套纯编译器覆盖所有 GPU 后端,而是把大量计算密集算子交给 TensorRT 手调内核,因此对单一厂商硬件的优化空间保留得更多。
这也就造成了两种完全不同的使用体验:用 XLA,你通常是在某个框架生态内部完成从 model 到 executable 的构建;用 MLIR-TensorRT,你更可能是在一个多语言、多框架、多模型的中间地带,把它当作统一推理底座的编译器。实际项目里,你也完全可以让 JAX 模型的子图先落成 StableHLO,再由 MLIR-TensorRT 接手做 NVIDIA GPU 上的极致编译,两者并不互斥。
4.2 与 TVM / IREE / Torch-TensorRT 的差异
TVM 和 IREE 都走“前端多框架、后端多硬件”的路线,理念上跟 MLIR-TensorRT 有重叠。TVM 最大的特点是提供了相当完整的 kernel 自动调优和代码生成体系,你可以自己写调度模板,也可以让 AutoTVM 去搜索配置。IREE 则是 MLIR 社区里很有代表性的推理编译器,它非常强调模块化,支持多种前端和后端,目标是把模型编译成可移植、可嵌入的二进制。
MLIR-TensorRT 跟它们最本质的区别在于,它不在“通用编译器 + 任何硬件”这条路上跟所有人竞争。它的目标后端非常明确:NVIDIA GPU。它也不追求把所有算子的代码都开放成可以任意修改的源代码,因为底层那些融合策略和 kernel 实现许多仍属于 TensorRT 的引擎能力。你可以把 MLIR-TensorRT 看作是“MLIR 的开放外壳 + TensorRT 的底层内核”组合体,更适合那些硬件路线已经收敛到 NVIDIA、同时又在编译层需要更高可编程性的团队。
Torch-TensorRT 是另一个容易混淆的项目。它帮你把 TorchScript 模块里的可优化子图转换到 TensorRT,但主要面向 PyTorch 用户,做了很多和 TorchScript 编译器对接的工作。MLIR-TensorRT 的前端和 IR 设计则更通用,它关心的是模型进入 MLIR 之后怎么被多种工具链共同处理,而不是某一个框架的原生格式。
4.3 对你作为使用者的直接价值清单
横向比较之后,我会把 MLIR-TensorRT 的实际收益归结成几条。
- 图级可编程性:你能用 pass 和 dialect 来改动图,而不是只能依赖 TensorRT 内置的融合策略。
- 多前端统一入口:ONNX、TorchScript、StableHLO 等途径都有机会进入同一套 pipeline。
- 多模型与显存的可控性:Plan 层面的内存规划和依赖描述,对多模型服务更友好。
- 可观察性:IR 是可以随时 dump 出来慢慢看的,出了问题更有可能定位到具体环节。
同时也要记住它的边界:它不是一个万能的端到端部署框架;算子覆盖范围依然受 TensorRT 支持能力限制;如果只是为了把一个成熟单模型做固化,拆掉现成方案转投 MLIR-TensorRT 未必划算。
5. 上手跑通一个最小推理流程:环境、命令与我能预见的坑
5.1 环境准备里的第一道坎:版本不是“选新的”而是“选配的”
想体验 MLIR-TensorRT,建议先从官方镜像或与项目 CI 绑定的依赖版本开始,而不是自己手动去凑最新的 TensorRT 和 CUDA。原因在于 MLIR 项目对 LLVM 版本的敏感度非常高,某个自定义 dialect 可能只兼容特定范围内的 MLIR 接口,一旦 LLVM 版本漂移,代码可能就编译不过了。
我自己的习惯是先检查宿主机 GPU 驱动能用,然后在容器里挂载 GPU 跑nvidia-smi确认设备可见。如果你是第一次在干净机器上配环境,最大的概率不是 MLIR-TensorRT 本身报错,而是驱动与容器运行时对不上,导致 TensorRT 找不到设备。遇到这类问题,先不要怀疑编译流程,先排查驱动版本、容器里能否看到设备以及 TensorRT 库的路径是否被正确加载。
从源码构建时,我会先把仓库 README 里锁定的 LLVM/MLIR commit 找到,再按官方脚本拉取依赖。某些子项目还会带独立的构建脚本,不要一股脑全开,最好按你需要的那条链路最小化构建。以你拉到的版本为准,这一条放到任何时候都适用。
5.2 编译一个 ONNX 模型并落地的典型动作
动手之前,先准备一个小模型,比如 ResNet-18 或一个简单的卷积网络,导出成 ONNX。通过项目提供的 compile 工具或等价 Python 绑定,把 ONNX 导入进来,跑一遍 pass pipeline,最后生成 Plan。
大致的动线是:先导入 ONNX 到 MLIR 模块;然后跑类型和形状推断,确保所有 tensor 的 shape 尽量可判定;接着做面向 TensorRT 的算子转换和融合 pass;随后做内存规划与依赖排序,输出 Plan;最后用 runtime 加载 Plan,绑定输入输出的 buffer,执行推理并验证结果。
听起来步骤不算多,但真实操作里最需要留意的,是“ONNX 里某个算子在 TensorRT 里没有直接对应实现”这种情况。MLIR-TensorRT 一般会尽量把不支持的算子保留在通用 dialect 中,后续可能还有机会用自定义 runner 执行,但你不能再要求 TensorRT engine 像原生算子那样优化它。项目中不少示例都会展示这种混合执行的能力,这是它与传统 TensorRT 的重要区别。
5.3 动态 shape 和缓存策略的两个实践提醒
第一点关于动态 shape。如果你明确知道服务上线后会遇到几种输入分辨率,就不要把动态范围设得过大。范围太宽,编译器要为边界形状准备保守策略和更多 kernel 变体,容易让显存占用和首次编译时间都显著上升。最好先在符号 shape 层面把边界清晰地表达出来,再用少量有代表性的 shape 去验证性能。
第二点关于 engine 缓存。MLIR-TensorRT 里的编译过程并不便宜,想把编译好的 Plan 或 engine 保存下来复用,机制上跟 TensorRT 的 engine 序列化类似,但缓存往往会跟 TensorRT 版本、GPU 架构和 CUDA 版本绑定。团队里如果有人升级了驱动或者换了卡,先前缓存可能无法加载。建议把缓存作为可再生成的产物而不是不可替代的资产,至少在 CI 里留一条“无缓存从零编译”的可执行路径。
5.4 遇见报错时的排查材料
如果编译过程中遇到算子不支持的报错,我会先换一个小而可复现的模型,把问题隔离出来,再用命令行的 IR dump 开关打印中间结果,看是形状推导失败、某 op 没找到对应模式还是类型不被支持。能看懂 MLIR 和 TensorRT 两头的报错,是这个项目调试的主要门槛,没有什么捷径。
启动时若遇到 CUDA 相关的找不到库或初始化失败,先跑一遍 TensorRT 自带的样例或trtexec验证底层环境。这样可以把问题边界切清楚:是环境问题还是 MLIR-TensorRT 项目本身的问题,不要混在一起猜。
6. 在业务里引入它之前,我的判断标准与现阶段槽点
6.1 我的评估清单:什么团队适合切进去
我不会建议所有做推理部署的团队立刻把管线切到 MLIR-TensorRT,但在下面几类痛点上,它确实值得认真评估。
- 你有多个模型或多个子图需要在同一 GPU 上协同工作,希望更系统地管理显存和流。
- 你觉得 TensorRT 的自动融合策略够强,但你想要一些“能自己插入规则”的可编程空间。
- 你需要跨框架(PyTorch、JAX、ONNX)汇总来源不同的模型,而不是只服务于一种前端。
- 你们团队已经有懂 MLIR 或编译器的人,愿意投入时间做底层调试,而不是只想调用一个黑盒 API。
如果没有这些需求,只是想把一个稳定模型部署到线上,传统 TensorRT 或 PyTorch 自带的导出推理链路可能已经足够,没有必要为一个新项目承担额外复杂度。
6.2 现阶段几个槽点,我不会美化
首先,项目仍然处在快速演进期。MLIR 接口本身变动频繁,MLIR-TensorRT 又紧跟上游,一段时间不看就会发现 API 和 pass pipeline 变了。依赖锁定和版本升级会持续消耗维护成本。
其次,文档和示例虽然不少,但比较分散。很多东西需要你把示例代码、官方 issue 和源码实现放在一起看,才能真正理解设计意图。对初学者来说,学习曲线比传统 TensorRT 陡峭不少,不仅要知道模型怎么部署,还要明白编译器基本概念。
最后,算子覆盖边界依然是它绕不开的限制。TensorRT 能做的它大都能做,TensorRT 做不了的新算子,它只能靠通用 dialect、自定义实现或者回退逻辑去补。也就是说,它本身不会凭空扩大 NVIDIA 推理算子的覆盖范围,更多是把已有能力放进一个更灵活的编译框架里。
6.3 如果想给团队讲清楚这个项目,我会用三句话兜底
如果要在评审会上用三句话讲清楚它的定位,我会这么说:第一,MLIR-TensorRT 是 NVIDIA 把 TensorRT 能力嵌入 MLIR 生态的一次工程化尝试,目标是让图优化过程更开放、更可编程。第二,它没有否定 TensorRT 的价值,而是把 TensorRT 放在新工具链的最底层,保留其内核性能优势。第三,它最适合的场景是异构来源、多模型、多 GPU 推理系统的统一编译与调度,而不是替代你手里已经跑得很好的单模型部署脚本。
从我个人的实际体会看,花时间研究 MLIR-TensorRT,最大的收获不一定是你马上能把线上延迟降低多少,而是你会开始用“编译器”的视角重新审视推理部署。过去那些“黑盒 engine”里令人焦虑的部分——算子怎么融合、内存怎么分配、多个模型怎么互相影响——会逐步变成可以被观察、被操纵、被复用的中间产物。哪怕最终没有全量切换过去,这种视角对做高性能推理服务也会有很大帮助。如果你以后真的需要在一个复杂推理系统里打通多个框架、多个模型,这套项目理念会是一个非常好的设计出发点。