训练好的模型离真正跑起来其实还差着“最后一公里”。模型权重只是“乐谱”推理框架和 AI 编译栈就是负责把乐谱翻译成不同乐团能直接演奏的总谱。今天想聊的正是这一层模型如何通过推理框架、AI 编译栈被高效地映射到 CPU、GPU、NPU 这类设备上解决“能跑”和“跑得快”这两个问题。这篇文章适合刚接触模型部署的人也适合已经用某个框架、但不知道为什么还有十倍性能提升空间的开发者。我不会照着某一家框架的文档念而是把链路拆开讲清楚再给一份可以直接抄作业的实操流程。1. 先理清楚推理框架、AI编译栈和设备映射分别干的是什么1.1 一个模型要跑起来中间其实隔着三层先说一个很多新手会有的误解一个 PyTorch 训练出来的.pth文件里面装的是权重和其他参数的序列化数据并不是“可执行程序”。设备真正执行的不是权重本身而是一系列算子组成的计算图。所以模型从静态文件变成运行中的程序至少要经过三层模型解析层、图优化/代码生成层、设备运行时层。推理框架负责第一道适配。它读取模型文件、解析成内部图结构并把图里的每个算子分发到具体设备后端。模型格式五花八门ONNX、TorchScript、GGUF、TensorRT engine 都由不同框架消化框架还要帮你管理 session、内存、并发执行、输入输出队列这些脏活。AI 编译栈是第二个大杀器。它做的是“翻译 优化”两件事把高层计算图翻译成设备能识别的底层指令同时在翻译过程中做算子融合、常量折叠、数据布局优化、内存规划。像 TVM、MLIR、XLA、TorchInductor 都属于这一类边界现在和推理框架越来越模糊。设备映射则是最后一跳。同一个算子GPU 上要用哪个 CUDA kernelCPU 上走什么 SIMD 向量化路径NPU 上怎么排布内存、怎么配置 dma都必须在编译阶段定下来。你看到的很多“部署后速度不对”本质就是这一层出了问题某个算子落回了通用实现性能直接掉一个数量级。1.2 为什么不能直接“双击运行”模型因为不同硬件的指令集、执行模型差太多了。GPU 靠的是成千上万个线程并行CPU 靠的是 SIMD 向量和高速缓存NPU 往往只能执行固定的算子集合。你把一个 PyTorch 模型直接丢给 CPU 跑它虽然能跑但每个算子都走通用逻辑性能可能比优化后慢几十倍。更麻烦的是同一个算子在同一个设备上也有好几种写法。以矩阵乘为例GPU 上要区分 float32、fp16、bf16、tf32还要考虑 tile 大小、线程块数量、向量化宽度、共享内存占用。这些参数组合起来空间巨大手写根本不可能在每个 case 上都写对。AI 编译栈的价值就在这里它不会写死一种实现而是根据 shape、精度、设备特性去搜索和生成最合适的 kernel。你可以这样理解模型文件是一份作曲家的总谱推理框架是排练调度AI 编译栈是编曲设备映射是最终每个乐器演奏家在台上实际发出的声音。没有编曲和排练直接让乐团照着草稿演奏效果和效率都很难保证。2. 推理框架模型部署的“第一道适配层”2.1 主流推理框架怎么选选推理框架本质上是选“场景匹配度”。没有绝对的好坏只有有没有用在正确的地方。我把常用的几个框架按场景列一下你会更容易判断。框架主要场景优势需要接受的取舍TensorRTNVIDIA GPU 生产部署性能极致、算子融合激进、支持 fp8/int8绑定 NVIDIA 硬件engine 构建慢算子黑盒化ONNX Runtime跨平台、跨硬件生态广、部署简单、EP 机制灵活默认优化相对保守复杂模型需要手动调 EPllama.cpp本地 LLM 推理GGUF 量化方便、CPU/GPU 混合跑、内存占用低主要围绕 transformer 类模型算子覆盖有限vLLMLLM 高并发服务PagedAttention、continuous batching 吞吐高对非 LLM 模型支持弱依赖 CUDA 生态TNN/MNN移动端、端侧部署针对手机硬件优化包体小算子库相对窄超大规模模型支持一般以 ONNX Runtime 为例它最核心的设计是 Execution ProviderEP。EP 相当于一个插槽插上 CUDA EP 就能把算子分发到 GPU插上 DNNL EP 就能在 CPU 上用 Intel 的底层库跑。你写业务代码只需要创建一次 session底层换 EP 不用动业务逻辑这对工程化非常友好。2.2 框架只是入口真正的效率在 backend 与 execution provider很多人以为“用了 ONNX Runtime 就自动快了”。实测下来如果你不做任何配置默认情况下算子可能一半跑在 GPU、一半因为不支持而落回 CPU性能反而比 PyTorch 还慢。这里有一个关键动作检查 session 的 provider 列表和实际生效算子。ONNX Runtime 有个 profiling 工具开起来之后可以看到每个算子落在哪个 EP 上。我排查慢推理问题时第一步永远是生成 profile看有没有大量算子走了 CPU EP而不是直接怀疑模型本身。PyTorch 生态也在朝同一个方向走。PyTorch 的 eager 模式为什么慢因为每个算子独立分派每个 kernel 启动都有固定开销卷积完结果写回显存、ReLU 又读一遍访存浪费严重。PyTorch 2.0 开始推 TorchDynamo TorchInductor本质就是加了一个 AI 编译栈先把 Python 层面的模型抓成 FX Graph再做图优化和代码生成最后输出一个融合后的 CUDA kernel。你会发现它已经不是一个单纯“框架”而是一个“带编译器的框架”。这也是我建议你理解推理框架时不要只看 API 的原因。框架只是入口底层是谁在干活、有没有走对 EP、有没有触发图优化才是决定速度的胜负手。3. AI 编译栈把计算图变成真正高效的机器指令3.1 图优化是第一个大杀器算子融合、常量折叠、数据布局AI 编译栈做优化第一步大多是对计算图做改写。最典型的是算子融合operator fusion。把“卷积 ReLU”合成一个 kernel避免中间结果写回显存把“LayerNorm 残差”合起来也能省掉一次张量读写。这些看起来微小的省访存对性能的影响非常大。我举一个实际数字假设一个卷积输出是256×128×64也就是约 200 万个 float。如果单独写回显存再读出来做 ReLU光是这个中间 TENSOR 来回搬运就是大约 64MB 的额外访存。模型推理时同样的算子被执行几万次累加起来就是 GB 级以上的访存浪费。融合成单 kernel 后这段数据可以留在寄存器或共享内存里省掉的访存对 latency 是实打实的提升。除了算子融合编译栈还会做常量折叠。模型里很多 shape 相关的计算比如seq_len // 2、dim * 4如果输入是确定的编译阶段直接算出结果不要在运行时重复计算。数据布局优化也很常见。NCHW 和 NHWC 的转换会让 CPU 和 GPU 有完全不同的访存表现编译器可以在图里插入最少的转置甚至把前后 transposed 算子抵消掉。这些优化都不是免费的。编译时间会增加而且动态 shape 会让很多优化失效。所以实际工程里我会尽量固定输入 shape哪怕只固定 batch 和 seq_len都能让图优化效果大幅提升。3.2 从图优化到代码生成TVM/MLIR 的一条龙流程TVM 是我觉得最适合拿来理解“AI 编译栈”的开源项目。它的流程很清晰先用前端把 ONNX 或 PyTorch 模型转换成 Relay IR然后在 Relay 上做一系列图优化再往下变成 TETensor Expression最后通过 AutoTVM 或 Ansor 自动调优搜索最优 schedule生成 CUDA/C/OpenCL 代码编译成可加载的库。这里有个非常反直觉的点编译器不是“直接生成最优代码”而是“搜索最优代码”。因为 GPU kernel 的参数空间太大编译器干脆把常见调度方式都列出来在一个子空间里做自动搜索和实测。比如卷积的循环展开、线程绑定、共享内存切分方式不同组合在不同显卡上表现差异极大干脆让编译器自己跑实验找最优解。这就是 AutoTVM/Ansor 的底层思路。MLIR 则是另一条更通用的路线。它用 dialect 表达不同抽象层次TOSA 描述算子语义Linalg 描述循环结构LLVM dialect 负责映射到机器指令。你可以把 MLIR 看作一套积木不同硬件厂商、不同前端框架都能插入自己的模块。实践中新的 AI 芯片厂商往往优先做一套 MLIR 后端而不是从零写全套算子库。编译栈因此成了 AI 硬件竞争里绕不开的基础设施。为什么编译栈不直接用 CUDA 手写算子因为手写 kernel 没法跟上模型变化的速度。Transformer 里一个 FlashAttention手写版本要考虑硬件特性换个新 GPU 又要调编译器把“注意力计算”和“访存策略”分开描述一次优化可以在多个算子结构上复用。3.3 内存规划低显存跑大模型的隐藏功夫很多人部署大模型第一反应是“卡里放不下”于是先量化。但编译栈的另一个隐藏贡献是内存规划。模型运行时不止有权重还有激活值、中间张量、KV cache。如果每算一个算子就分配一块新内存显存再多也不够造。编译器会做 liveness analysis分析每个张量的存活区间把互不冲突的张量复用到同一块显存上。这种 arena 内存分配策略在编译后的推理引擎里几乎是标配。我见过一个 7B 模型用 naive 逐算子分配显存时峰值需求比编译后高 30% 以上这 30% 可能就决定了你的 8G 卡能不能跑起来。“低显存运行模型”这个话题近年特别火核心就是三件事权重压缩、KV cache 管理、分层 offload。我算过一笔账一个 7B 模型fp16 权重大约 14GB用 4bit 量化和分组打包可以压到 4GB 左右。KV cache 也不便宜假设 32 层、hidden size 4096、序列长度 2048、fp16光一个 batch 的 KV cache 就约 1GB如果 batch 提到 32就是 32GB。vLLM 的 PagedAttention 本质就是不让 KV cache 一次性全分配而是像操作系统分页一样按需分配、按页管理这样显存利用率立刻上来了。还有 offload 策略把权重切片放在 CPU 内存GPU 只保留当前计算层需要的部分用的时候拷上去算完再换下一层。代价是 PCIe 带宽成为瓶颈所以适合“能跑但慢”的场景。模型能不能在低显存设备上跑起来从来不是单看权重大小而是整个推理链路的内存管理是不是够聪明。4. 实操把一个小模型完整映射到设备全流程4.1 实操起点从 PyTorch 到 ONNX理论讲再多不如亲手走一遍链路。我选一个最常见的 ResNet18 分类模型演示从 PyTorch 导出 ONNX再到 ORT 和 TensorRT 推理的完整过程。第一步是用torch.onnx.export导出计算图。import torch import torchvision.models as models model models.resnet18(pretrainedTrue).eval() x torch.randn(1, 3, 224, 224) with torch.no_grad(): torch.onnx.export( model, x, resnet18.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version17, do_constant_foldingTrue, )这里有两个关键参数opset_version决定导出 ONNX 算子集的版本版本太低很多新算子不支持版本太高有些推理框架跟不上一般 17 左右比较稳dynamic_axes把 batch 维度声明成动态这样同一个模型可以接受1、8、32不同 batch 的输入。但注意动态轴会降低图优化效果如果业务场景 batch 固定建议先把它写死。导出时踩坑最多的三类问题模型里有 Python 控制流、切片 shape 依赖输入、自定义算子没有 ONNX 映射。规避办法是导出前先用固定输入跑一次 tracing再检查输出的 ONNX 计算图是否符合预期。不要嫌这一步麻烦很多转换失败问题在图上两三分钟就能看出端倪。4.2 在 ONNX Runtime 和 TensorRT 上跑起来模型转成 ONNX 后用 ONNX Runtime 推理非常简单。import numpy as np import onnxruntime as ort opts ort.SessionOptions() opts.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess ort.InferenceSession( resnet18.onnx, sess_optionsopts, providers[CUDAExecutionProvider, CPUExecutionProvider], ) input_data np.random.randn(1, 3, 224, 224).astype(np.float32) outputs sess.run([output], {input: input_data}) print(outputs[0].shape)注意 providers 列表的顺序就是优先级CUDA EP 在前遇到不支持的算子会 fallback 到 CPU EP。这看起来是容错实际上是性能陷阱。我建议在开发环境开启 profiling明确确认算子是否真的都跑在 GPU 上。ORT_ENABLE_ALL开满图优化后很多融合逻辑会生效但若你的模型里有非标准算子也可能导致 build 变慢或行为异常可以先从ORT_ENABLE_BASIC开始验证正确性。如果要上 NVIDIA GPU 的极致性能下一步就是转成 TensorRT engine。命令行最简单的方式是用trtexectrtexec \ --onnxresnet18.onnx \ --fp16 \ --saveEngineresnet18_fp16.engine \ --shapesinput:1x3x224x224TensorRT 10 之后的版本里workspace 参数改成了--memPoolSizeworkspace:2G老版本还是--maxWorkspaceSize用之前先看版本。一个容易忽略的点是生成的.engine文件和显卡型号、CUDA 版本、shape 上限强绑定换机或者改输入上限通常要重新构建。所以生产环境一般会在部署机上执行转换不在开发机转好再拷过去。fp16 通常能带来近乎翻倍的推理速度提升但不是所有算子在 fp16 下都安全。如果发现输出精度明显变差可以先用--fp16 --strictTypeConstraints让所有层保持 fp16再逐层排查是哪一个算子对精度敏感把该层手动回调到 fp32。这个排查过程看起来繁琐却是部署落地时绕不开的一步。4.3 用 TVM 把算子自动调优如果你要部署的目标设备不是 NVIDIA GPU或者想试试编译栈带来的额外优化空间可以用 TVM 走一遍自动调优流程。以下代码是把同一个 ONNX 模型交给 TVM 做自动调度和编译的最小示例。import tvm from tvm.contrib import graph_executor from tvm import auto_scheduler # 这里可用 onnx 或 relay 前端 mod, params tvm.relay.from_onnx(resnet18.onnx) target tvm.target.Target(cuda) tasks, task_weights auto_scheduler.extract_tasks( mod[main], params, target ) tuner auto_scheduler.TaskScheduler(tasks, task_weights) tuning_options auto_scheduler.TuningOptions( num_measure_trials200, measure_callbacks[auto_scheduler.RecordToFile(resnet18_tuning.log)], ) tuner.tune(tuning_options)这一段里最关键的是num_measure_trials它决定搜索次数。次数太少搜不到好 schedule次数太多调优时间会非常长。我一般先在100~200次左右跑通确认输出正确性没问题再加大到1000做正式调优。实测中TVM 自动搜出来的卷积 kernel可以比常规 cuDNN 实现快 10% 到 30%但也可能更慢所以最终要以设备上的真实 benchmark 为准。TVM 这条链路最大的价值不是替代 TensorRT而是让你理解“算子性能”是可以被编译器搜索出来的。以后碰到某些奇奇怪怪的硬件你至少有思路图优化 自动调度比硬碰硬写 kernel 要靠谱得多。5. 常见问题与排查技巧实录5.1 一张“部署问题速查表”部署踩坑不是一次性的我把从业几年最常遇到的五类问题整理成速查表基本覆盖了 80% 的场景。现象可能原因排查手段常用解法模型文件转换失败算子不兼容、opset 版本低、动态 shape 依赖用onnx.checker校验、打印 graph 节点更换算子、升级 opset、固定输入 shape推理速度比预期慢很多算子 fallback 到 CPU、未开图优化开启 profiling查看算子落在哪个 provider调 EP 顺序、开启 ORT_ENABLE_ALL、转 TensorRT显存 OOM内存未复用、KV cache 突增、权重未量化记录峰值显存分析激活张量 liveness换 arena 分配、开量化、用 PagedAttention、offloadfp16 下精度异常敏感算子溢出、量化校准数据不够对照 fp32 输出逐层对比误差敏感层保持 fp32、重新校准、按 tensor 调试动态 shape 下时快时慢动态维度让优化失效、每帧重新编译固定 batch/seq_len 对比限制动态轴范围、按常见 shape 预 build这五类问题并不孤立。我遇到过最典型的场景是一个 transformer 模型在静态 shape 下跑得飞快一开动态 batch 就卡得离谱。原因是动态维度让编译栈无法预分配 kernel 和内存某些算子开始走通用路径。你只需要把 batch 限制在几个固定档位比如 1、4、8、16四个静态 engine 轮换性能立刻回来。5.2 一个值得养成的调试习惯与其等出问题再修不如在一开始就把“链路验证”做扎实。拿到一个模型后我通常按这个顺序检查固定输入跑一遍参考输出转成 ONNX 后先 CPU 推理对输出再用目标设备 EP 推理对输出最后才做性能优化。每一步只引入一个变量出问题就能立刻锁定是哪一步引入的。还有一个非常容易被忽略的小细节不管用什么推理框架都要检查输入张量的内存是否连续。很多模型部署后第一帧很慢是因为输入做了非连续切片框架被迫做一次额外拷贝。先把输入numpy.ascontiguousarray处理一下能省掉很多看不见的开销。我在实际工作中感受最深的一点是推理框架和 AI 编译栈的界限越来越模糊但底层的原理始终是那些——图优化、内存复用、算子调度、精度取舍。你只要能拆开这一层任何新的框架、新的硬件平台出现时都能快速迁移过去而不是被某一个工具的 API 绑死。最后再分享一个小技巧当你怀疑“为什么我这个模型在别人机器上很快、在我这儿很慢”的时候先别急着调框架用nvidia-smi -l 1盯一眼 GPU 利用率和显存占用。如果利用率波动大、显存长期不释放问题大概率在内存复用和 batch 策略上如果利用率跑满但速度上不去才轮到算子融合和 kernel 选型的问题。这个顺序反了很容易白忙活两三天。模型部署这条路看起来是被各种框架推着走但真正拉开差距的永远是你会不会从“能跑”往下多挖一层去理解设备是怎么把模型吃进去、再吐出来的。希望这篇梳理能让你下次面对几个 G 的模型文件时心里更有底。