资讯动态

推理框架与AI编译栈:模型部署的底层逻辑

发布时间:2026/10/1 13:38:35 来源:尧图企业网站定制
我最近翻了一些推理框架的源码比如 ONNX Runtime、TensorRT、TVM、llama.cpp 这类项目最大的感受是这些不是一个个单独的库而是整套串起来的“栈”。很多人在问“模型怎么才能跑起来”的时候卡住的原因往往不是模型本身而是没搞懂这最后一层——推理框架和 AI 编译栈到底做了什么。标题里说的“第三层”就是指这条链路里最靠近硬件的那一层模型从 PyTorch 权重文件变成 GPU 或 CPU 上真实计算指令的过程。这篇文章适合正在做大模型部署、端侧推理、低显存运行优化或者想把模型接到自己业务系统里的人读。我会把推理框架和编译栈分开拆先讲它们各自解决什么问题再讲实际部署中怎么选、怎么配、怎么排查。看到最后你会发现所谓“高效映射到设备”核心就三件事图怎么组织、算子怎么调度、内存怎么复用。1. 先搞清楚模型离“能跑”还差多少步1.1 训练好的模型只是一堆数字训练完成后我们拿到的权重文件本质上是一堆浮点数加上网络结构的描述。以 Transformer 为例它的计算过程是输入 token 先经过 Embedding 查到向量然后进入若干层 Transformer Block每层里有 QKV 投影、注意力矩阵乘法、FFN 两层全连接、LayerNorm、残差连接。这些操作写出来是一张计算图节点是算子MatMul、Softmax、Add、LayerNorm边是张量流动。这张图本身不会执行。PyTorch 的model.forward()只是按 Python 逻辑把算子逐个调出来每一步都走一遍 Python 解释器效率很低。真正生产环境的推理需要把这张图交给一个专门的执行引擎让它理解算子之间的依赖关系、预先规划好内存、选择最高效的算子实现然后在 GPU 上批量启动 kernel。我习惯用一个类比训练好的模型相当于一位顶级厨师手里攥着菜谱计算图和食材权重但厨房里没有传菜路径、没有灶台分工每次做菜都要临时找锅找铲。推理框架和编译栈要做的事情就是把这个厨房重新装修一遍谁先炒、谁后蒸、哪口锅共用、哪些食材提前切好全部排定。1.2 框架层与编译层的分工边界“推理框架”和“AI 编译栈”这两个词经常混着用但职责不同。推理框架处理的是计算图的装载、算子分发、内存池管理和设备调度它知道“先做 A 再做 B”以及“A 的输出放在哪里给 B 用”。AI 编译栈则更进一步它要把计算图中的每个算子翻译成硬件能直接执行的底层指令并在这个过程中做算子融合、循环分块、向量化等优化。打个比方推理框架是项目经理负责排期和资源管理AI 编译栈是翻译官把“帮我算一个矩阵乘”翻译成“用 tensor core 加载 16x16 的 tile 乘 16x16 的 tile 累加”。两者边界有时候模糊——比如 TensorRT 既是推理引擎又自带编译器TVM 更是从图优化一路做到 kernel 生成但理解这条分工线排查问题时思路会清楚很多。分层设计这件事跟 OSI 参考模型把网络功能拆成七层是同一个道理。每一层只向上层提供稳定接口内部实现可以随便换。你今天用 ONNX Runtime明天觉得算子实现不够快可以换成 TensorRT 或 TVM但上层的业务代码不用动因为接口都是加载图、喂输入、拿输出。1.3 静态图与动态图的差异这直接关系到“映射”的难度。PyTorch 默认是动态图每次 forward 都重新建图好处是灵活坏处是很多优化在运行时漏看。TensorFlow 的 Graph Mode 和 PyTorch 的torch.compile走的是静态图路线先 trace 一遍模型把算子之间的连接关系固定下来然后编译器可以看到整张图的全貌。静态图最直观的红利是算子融合。比如 Conv 后面接 BatchNorm再接 ReLU动态图模式下这三个算子是三个 kernel每算完一步就要把中间结果写回显存、再读出来。静态图模式下编译器可以把三步合成一个 kernel中间张量直接留在寄存器或共享内存里。实际部署时同一个模型用静态图推理比动态图快 20% 到 50% 是很常见的。所以当你拿到一个模型准备部署第一件事不是急着选框架而是先问这个模型能不能转成静态图能转的话后面编译优化的空间就大。这不只针对深度学习模型像 LightGBM 这类传统模型的推理同样存在特征计算到原生决策树执行的映射过程只不过优化点不在算子融合而在特征分箱和内存布局。2. 推理框架真正“跑起来”的执行引擎2.1 一张计算图是怎么被驱动的推理框架的核心组件是执行器。ONNX Runtime 的InferenceSession、TensorRT 的ExecutionContext、llama.cpp 的llama_context名字不同做的事情一样维护输入输出 buffer按拓扑序逐层执行计算图中的节点。驱动计算图有一个很关键的动作拓扑排序。框架要先分析每个节点依赖哪些前驱节点排出严格的执行顺序。一个节点如果依赖两个输入那两个输入必须都算完它才能启动。这个依赖关系在框架里表现为 task 状态机GPU 场景下还会进一步转成 CUDA Stream 上的 kernel 启动序列。用一个最简单 ONNX Runtime 调用流程说明import onnxruntime as ort sess ort.InferenceSession( model.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider] ) outputs sess.run( None, {input_ids: input_ids, attention_mask: mask} )providers参数很值得玩味。它执行的是“按优先级挑设备”的逻辑第一个 provider 能跑的算子全部归它剩下不支持的算子自动掉到下一个 provider。这个机制叫算子回退实际部署里经常用到——GPU 上不支持的算子自动落到 CPU 上保证模型能完整跑完代价是这些算子之间有设备间数据拷贝速度慢。2.2 图优化为什么先融合再执行推理框架加载模型之后做的第一件事不是执行而是“图改写”。ONNX Runtime 里有一套 GraphOptimizerTensorRT 里也有类似模块做的工作包括常量折叠、维度简化、死分支消除还有最常见的算子融合。以 Transformer 推理为例Attention 块里的QK^T * scale、Softmax、Dropout、V 矩阵乘这串操作如果不融合中间矩阵要完整地写进显存再读出来显存占用和带宽开销都很疼。融合之后QK^T的计算结果不会落到显存而是直接流进下一个 kernel 的共享内存Softmax 算完也直接做attn V。我在部署 DeBERTa、Longformer 这类结构改动较大的模型时对融合的体会更深。这些模型里长序列的注意力计算对内存布局极其敏感torch.compile的融合策略能明显减少中间张量我实测过长了近一倍的序列也没有爆显存。反过来像 CLIP 这类视觉模型微调后导出 ONNX如果没开图优化同一个卷积后面的 BN 和激活会分成三个 kernel 跑速度可能慢两倍。2.3 内存池与显存复用低显存运行的核心很多人的第一反应是“我的显卡显存不够跑不了大模型”但其实推理框架的内存管理比想象中要精打细算得多。推理时显存分配不会每次都向驱动申请而是用内存池预分配一块大空间张量用完不是释放而是回收到空闲列表。推理过程里显存占用的大头是三类模型权重、KV Cache、激活值。KV Cache 是自回归模型推理时保存的历史键值对每一轮生成都要追加不是一次性分配。它的体积可以用一个公式估算2K和V × 层数 × 序列长度 × hidden_size × 字节数。以 7B 模型为例如果层数是 32hidden_size 是 4096序列长度 2048那么 KV Cache 大约是 1GB。权重部分更好算7B 模型用 FP16 存储权重约 14GB改成 INT8 约 7GBINT4 约 3.5GB。这就是为什么量化是低显存跑大模型最直接的手段。权重做完 INT4 量化加上 KV Cache 和激活值8GB 显存就能跑起 7B 模型。这也是 GGUF 格式能流行起来的原因——llama.cpp 把权重按量化 group 打包加载时不需要完整还原到 FP16而是直接在量化域里做计算。内存池设计里有一个容易被忽略的点碎片化。推理过程中张量大小不断变化如果频繁分配释放空闲块会被切得七零八落。框架的解法一般是分桶管理——把常用尺寸的 buffer 单独圈起来避免碎片。我自己调低显存部署时最常做的操作就是从框架的日志接口里打开内存统计确认有没有出现 allocation failure如果没有再调整 KV Cache 上限把显存余量压到最小。2.4 执行策略异步、流与批处理推理框架不只是把算子串起来执行它还研究怎么并行。GPU 场景下 CUDA Stream 允许不同 kernel 在不同 stream 上并发执行。比如 Attention 里多头计算是独立的可以拆到多个 stream 上同时跑CPU 场景下则是线程池加任务队列把算子在多核上分摊。更贴近生产的是动态批处理。在线服务场景里请求是零散进来的如果每个请求单独推理一次GPU 利用率很低。推理框架会把多个请求攒一波按相同 shape 打包成一个 batch再一次性算完。TensorRT 的 dynamic shape、ONNX Runtime 的shared memory arena、vLLM 的 continuous batching 都是在解决这个问题。这里要提醒一点动态 batching 不是实时意识攒 batch 有延迟攒得太久用户体验差攒得太短 GPU 吞吐上不去。实际部署时我一般把最大 batch 设为 4 到 8遇到并发尖峰宁可多等几十毫秒也要保证 GPU 做满。调这个参数时用 GPU 利用率做指标尽量让利用率稳定在 70% 以上。3. AI 编译栈把算子翻译成设备听得懂的指令3.1 为什么需要中间表示IR编译器领域有个共识直接从一个高级语言生成机器码容易耦合、难优化。所以主流编译器都会先转成中间表示IR在 IR 上做多轮优化最后再生成目标代码。用写文章来类比先写提纲、再写初稿、最后改成投稿格式中间的提纲就是 IR。AI 编译栈里的 IR 通常分两层。第一层是图级 IR描述算子之间的连接关系第二层是算子级 IR描述单个算子的循环与内存访问。TVM 的图级 IR 是 Relay算子级 IR 是 TIRMLIR 则有 linalg、triton 等多种方言对应不同抽象层次。分层好处是图优化不关心算子内部细节算子优化不关心图拓扑各改各的互不干扰。3.2 Lowering把抽象算子落到硬件原语“Lowering”是编译栈里出现频率最高的词意思是把高级算子拆成硬件能执行的底层指令。最典型的是把卷积下降成矩阵乘。卷积本质是滑窗乘加但直接按滑窗实现访存效率极差用 im2col 把输入转成矩阵再调 GEMM 库如 cuBLAS就能利用高度优化的矩阵乘 kernel。这也是为什么很多框架的卷积性能取决于 GEMM 优化水平。下降过程还包括 tiling、padding、向量化。tiling 是把一个大矩阵切成多个小 tile让每个 tile 的计算能塞进 GPU 共享内存或 CPU 缓存。一个典型配置是把 4096x4096 的矩阵切成 128x128 的块逐块做乘法累加。padding 是为了满足硬件对齐要求比如 AVX 指令一次处理 8 个 float如果最后剩 3 个元素就补到 8。向量化则是把标量循环改写成向量指令在 CPU 上效果尤其明显。以 FlashAttention 为例它的核心优化就是把一整段注意力计算切成 block每个 block 只占一小块 SRAM避免创建完整的N x N注意力矩阵。这一步本质上就是 tiling 思想在注意力计算上的应用。你看很多所谓“创新优化”落到编译栈视角都是老几样切块、复用、融合。Triton 这类语言进一步简化了 kernel 编写。你可以用类似 Python 的语法描述一个计算 kernel剩下的 tiling、向量化由编译器自动完成。下面是一个很简化的 Softmax kernel 示意重点看它的分块语法import triton import triton.language as tl triton.jit def softmax_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr): offsets tl.program_id(0) * BLOCK_SIZE tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) x x - tl.max(x, axis0) x tl.exp(x) y x / tl.sum(x, axis0) tl.store(y_ptr offsets, y, maskmask)3.3 自动调优让机器自己找最快路径算子实现选型有一个棘手的问题不同硬件上最快的 tiling 方案不一样。A100 上128x128的 tile 最快某些消费级显卡上可能是64x64最快。让工程师手工调每个算子不现实于是编译器引入了自动调优机制。TVM 里的 Ansor、AutoTVM会生成一批候选 schedule循环顺序、分块大小、展开方式在实际硬件上跑一遍用成本模型预估或实测结果选出最优。这个过程很吃时间一个模型全部算子自动调优可能跑几小时甚至几天但调完的 kernel 是手写级别甚至更优。实际工作中要注意自动调优结果是硬件相关的换一张卡就必须重新调。很多部署项目在 A100 上调好的 engine拿到 T4 上跑反而变慢就是因为缓存了不该缓存的配置。我建议在正式部署前对目标设备做一次专门的调优或采集tactics配置不要跨设备复用。3.4 现实中的编译栈选型对比现在主流的编译栈方案我没有必争按场景分开看方案抽象层级优势典型场景TensorRT图级 算子级对 NVIDIA GPU 深度优化、融合激进云端 GPU 高吞吐推理ONNX Runtime图级 内置优化生态广、多后端支持中场景通用部署TVM / Apache TVM图级 算子级硬件覆盖面广、算子优化灵活边缘设备、新芯片适配MLIR多层 IR适合构建自研编译流程编译器基础设施、芯片公司Triton算子级上手快、GPU kernel 开发效率高研究、算子库开发torch.compile图级 Triton 后端一行开启编译, 对 PyTorch 模型友好PyTorch 模型快速提速我的经验是如果是 NVIDIA GPU 上部署服务首选 TensorRT它做的算子融合和 kernel 自动选择比我手动配置稳定得多如果模型结构经常变、快速验证为主用 ONNX Runtime 或torch.compile性价比更高如果在做端侧或专用芯片TVM 是绕不开的因为它对算子实现的控制粒度最细。4. 设备映射实操低显存、CPU 与端侧部署4.1 低显存跑大模型量化与分块推理低显存运行模型本质是取舍显存不够要么缩小权重体积要么降低中间结果的显存占用。前者靠量化后者靠算子融合和内存池优化。两种手段可以叠加实际部署中也是叠加着用。量化的选择不是简单地“从 FP16 换成 INT8”而是要考虑量化的粒度。GGUF 格式里常见的量化级别有 Q4_0、Q4_K_M、Q5_K_M、Q8_0。Q8_0 有 8bit 权重质量几乎无损但体积只减一半左右Q4_K_M 在 4bit 基础上分组量化结合了 k-quant 策略是目前 7B 模型在 8GB 显存上跑的首选体积约 4.7GB。部署时显存占用可以大致估算。比如一个 13B 模型用 FP16 跑需要 26GB 权重显存显然 16GB 显卡装不下换 Q4_K_M 以后权重约 7.8GB加上 KV Cache 和激活值勉强能塞进 16GB。如果你还要同时处理长上下文KV Cache 的增长会很快这时候就要考虑 KV Cache 量化——把缓存里的值从 FP16 压到 INT8显存直接减半代价是生成质量有轻微波动。低显存部署时还有个技巧叫分块推理block-wise inference。对于超大模型模型权重可以按层分块加载一部分算一部分算完的层释放掉再加载下一层。llama.cpp 的--n-gpu-layers参数就是控制模型层在 GPU 和 CPU 间的分配实测下来一半层放 GPU一半层放 CPU7B 模型在 6GB 显卡上也能流畅生成只是速度比全 GPU 慢 30% 上下。4.2 CPU 场景线程、向量指令与缓存别以为 CPU 推理没有优化空间。llama.cpp 在纯 CPU 环境下通过 OpenMP 多线程和 AVX2/AVX512 向量指令也能跑得比原生 PyTorch 快一个量级。核心原因是 PyTorch 每个算子独立调度线程同步开销很大而 llama.cpp 的算子全部手动向量化且多个算子融合成一个 kernel线程只启动一次。CPU 部署调参时线程数是最关键的。设得太少多核浪费设得太多线程切换开销反而拖慢速度。我的经验是对于 7B 模型物理核数 8 到 12 的情况下线程设为物理核数即可超线程带来的额外线程收益很小。另外要注意 NUMA 架构如果机器有多个 CPU 插槽llama.cpp 默认会把线程散到两个插槽上跨插槽访问内存非常慢最好用taskset固定到单个插槽。内存带宽也是 CPU 推理的快慢分水岭。7B 模型每个 token 都要把全部权重从内存过一遍权重 4.7GB按 20 token/s 的速度算每秒要读接近 100GB 的数据这个量级只有 DDR4 双通道才能兜住。所以 CPU 跑大模型内存频率的影响甚至比 CPU 主频还大我推荐优先选大内存带宽的机器而不是堆核心数。4.3 把本地模型接到应用从 Ollama 到 LM Studio现在很流行的做法是用 Ollama 把 GGUF 模型包装成 HTTP 服务。命令很简单ollama serve ollama run llama3.2启动后本地会监听 11434 端口可以用标准的 HTTP 请求调模型curl http://127.0.0.1:11434/api/generate -d { model: llama3.2, prompt: 你好 }这类服务化的价值在于上层应用不再直接依赖某个框架的 Python API而是通过统一接口调用模型。你在 Langflow 这类可视化工具里配置自定义模型服务地址时填的通常就是http://127.0.0.1:11434这类地址逻辑上是把 LLM 推理做成了一个可插拔的组件。LM Studio 走的也是类似路线只是更偏桌面端。很多 AI 编程工具支持自定义模型端点把地址指到本地模型服务就行。我试过的体验是本地 7B 模型响应速度在 30 token/s 左右对于补全类任务足够用但要求高时还是不如更专业的云端模型。关键在于本地模型不会上传你的代码和对话内容数据隐私场景下这个优势是硬性的。4.4 一个被我踩过的坑切模型后上下文“跳闪”有段时间我在本地跑了多个模型用一套界面反复切换使用。切到另一个模型后原对话界面会反复跳闪、模型反复输出。排查到最后问题是切换模型时上层应用仍然复用了旧的上下文 buffer而底层模型已经换了两边的 token 映射表完全不一致。这不是模型本身的问题而是推理框架和应用的接口契约没对齐。解决方法是切换模型前必须重置 session、清空 KV Cache必要时新建一个上下文对象。现在主流框架都提供了 session 级别隔离一个 session 只对应一个模型不要跨模型复用会话。给模型服务建立一层隔离非常重要。在同一台机器上跑多个模型时我习惯开多个进程进程间用端口区分进程内只加载一个权重文件。这样既避免了共享内存互相污染也方便单独重启某个模型而不用动其他服务。5. 踩坑实录推理栈问题排查速查5.1 高频问题与排查思路现象可能原因排查思路模型加载极慢未启用内存池、磁盘读取权重大、模型未预热检查日志中加载耗时启用框架缓存推理爆显存KV Cache 过大、激活值未复用、算子未融合打开显存统计降低 max_seq_len 或启用量化量化后生成质量下降明显量化粒度太粗、敏感层未保留精度换 Q5_K_M/Q8_0或对指定层跳过量化CPU 推理速度远低于预期线程数设置不当、内存带宽不足检查物理核数固定 NUMA 节点切换模型后对话跳闪上下文 buffer 未重置、模型映射不一致切换模型前清空上下文、隔离 sessionComfyUI 下载模型失败下载目录权限不足、文件名与预期不匹配、校验值失败手工下载后放到指定目录核对 sha256自定义模型服务地址配置不生效地址不通、返回格式不符合预期先用 curl 验证接口再看日志这里重点说两个我常遇见的。第一个是“量化后效果差”。很多人以为只要模型能加载就是优化成功但量化模型在边缘场景数学推理、长文本上掉点可能很明显。我的做法是先用 Q8_0 跑一遍基准确认精度可用后再尝试 Q4_K_M同时把注意力层保留为高精度。如果 Q4_K_M 掉点太多放弃低比特加长上下文长度比降比特实惠。第二个是“ComfyUI 下载模型失败”。这类问题多数时候不是 ComfyUI 本身的问题而是模型文件放置路径不对。手动下载再放进models/checkpoints或models/diffusers目录重启一遍界面比反复点界面上的下载按钮可靠得多。模型文件名要严格匹配代码里的预期否则权重加载不上。5.2 部署选型时我习惯的排查顺序如果不确定用什么框架先做 profiling。拿模型在 PyTorch 里跑一遍用 PyTorch Profiler 抓出每个算子的耗时看瓶颈在哪。如果瓶颈是算子访存开销大优先考虑图优化和算子融合如果瓶颈是 kernel 启动次数太多优先考虑换更激进的编译方案如果是显存严重不足先做量化和内存池配置别急着换框架。所有优化做完以后保留一个可复现的 benchmark 脚本。我习惯在每次改动后跑三遍取中位数记录延迟和显存峰值。没有这个基准很容易陷入“调参一时爽上线火葬场”的循环。参数改动要一次只改一个避免多变量同时变化说不清是谁导致的收益或劣化。框架切换时不要迷信“支持某格式就是万能”。ONNX Runtime 支持很多算子但不同版本对同一算子的实现差异很大。部署前用onnxruntime.transformers里的优化脚本先跑一遍重点确认 attention 类算子有没有被正确融合成高效版本。很多模型导出 ONNX 后变慢就是算子停留在低效实现上。5.3 低显存用户的一点额外建议如果你的显卡只有 6GB 或 8GB别硬上 13B 模型。7B 模型用 Q4_K_M 已经能策略性地覆盖大多数场景。另一个方向是使用更小的 embedding 和 intermediate size 的模型比如同参数规模下架构更年轻的模型在同等量化下表现更好。还有一个细节很多人忽略推理时的 batch size 直接影响显存。自回归生成里 batch size 对应同时生成几条序列从 1 提到 4KV Cache 直接翻四倍。如果只是个人使用batch size 保持 1 就好换来的显存余量可以拉长上下文长度。关于上下文长度KV Cache 的显存不是线性增长的而是随序列长度增长。如果模型支持 RoPE 位置编码扩展有些框架可以通过配置动态增加长度上限但需要额外显存。每次调长上下文请检查显存峰值别等到 OOM 再回滚。我个人的体会是推理框架和编译栈这些底层技术看起来复杂真正上手以后最关键的思路反而是克制。不要一上来就追求端到端编译、全算子手写 kernel先用现成框架把链路跑通再根据 profiling 数据逐步优化。多数场景下图优化加量化已经能解决 80% 的问题剩下的 20% 才需要深入到编译栈层面去打磨。最后分享一个小技巧部署前把模型的输入输出 shape 搞清楚特别是动态维度。很多框架在 shape 不确定时不会做激进的图优化因为 kernel 选择需要知道大小。能固定 batch size、固定序列长度上限的尽量固定。这几个数字提前定死框架能做的优化空间会大一个量级。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑