推理框架与AI编译栈：从PyTorch到RK3568的部署优化指南

上个月帮客户调一块RK3568的板子模型是内部训练好的一个LSTM回归模型用来做设备状态预测。本地GPU上跑得好好的一到盒子里面就卡成PPTCPU占用拉到90%推理一次要300多毫秒。客户甩过来一句“你们这模型是不是太大了”其实模型只有几十MB。问题恰恰出在“用训练框架直接推理”这个错误姿势上。PyTorch的动态图、Python运行时、自动求导机制每一层都为了灵活性付出了额外开销而这些开销在设备端根本不需要。后来我沿着“推理框架 AI编译栈”这条路把部署链路捋了一遍模型最终跑进了20毫秒以内CPU占用降到了15%。这篇文章就把这条链路拆开讲清楚。1. 先搞清楚“第三层”卡在哪个位置训练灵活性与设备执行的语义鸿沟1.1 训练框架的“自由”是部署的“负担”PyTorch和TensorFlow这类训练框架设计目标只有一个让算法工程师方便地组装网络、算梯度、做训练实验。动态图意味着每层计算都要被解释器跟踪一遍自动求导意味着内存里保留了完整的反向计算图Python的GIL又决定了多线程推理很难吃到多核红利。这些机制在训练阶段是优点在部署阶段就是纯粹的浪费。举个例子你在PyTorch里跑一个Transformer模型forward过程中框架要维护一个自动求导的计算图每个算子的输入输出都要有引用计数张量分配走的是缓存分配器。到了部署场景我们不反向传播、不更新权重只需要前向推理那维护计算图、保留梯度信息这些动作全部可以去掉。可如果你直接拿torch.load的模型在设备上用Python做推理这些无用开销一样都跑不掉。热词里那个“transformer模型详解”“deberta模型结构图”我能理解模型结构网上讲得很多但真正落地时大家普遍缺的是“怎么把结构变成设备上高效执行的指令序列”。这就要靠中间那一层来解决。1.2 模型翻译成设备语言要经过的三道工序把训练好的模型变成设备上运行的二进制需要经历三个翻译阶段。第一阶段是计算图提取。从PyTorch的动态图里冻结出静态计算图导出成ONNX这样与框架无关的中间表示。这个阶段解决的问题是“把模型的逻辑结构固定下来”。第二阶段是图级别的优化。ONNX图拿到手之后推理框架会做算子融合、常量折叠、冗余去除、精度改写。这些优化不依赖具体硬件属于通用优化。比如把Conv BN ReLU融合成一个算子把两个相邻的Reshape合并掉这些动作在写训练代码时是不可能做的因为训练需要保留中间的梯度流。第三阶段是算子到指令的映射。图优化完之后每个算子要落在具体硬件上执行。同样的一个卷积在英伟达GPU上可能要走cuDNN或TensorRT在瑞芯微的NPU上要走RKNN的算子库在ARM CPU上要走CMSIS-NN或手写NEON汇编。这个映射过程就是推理框架的算子调度加AI编译器的代码生成在做的事。1.3 我理解中的四层部署栈这里还得把“第三层”的上下文交代清楚。按我习惯的分层方式一套AI部署链路是这样切的第一层是算法训练层就是你用PyTorch、TensorFlow训练模型、验证精度的那一层第二层是模型压缩层负责把模型变小变快包括量化、剪枝、知识蒸馏第三层就是本文要聊的推理框架与AI编译栈它把压缩后的模型文件接收过来做图优化、算子调度、内存规划、甚至直接生成设备代码第四层是硬件执行层也就是GPU、NPU、DSP、CPU这些计算单元。这四层之间不是严格串行的实际工程中经常交叉。比如量化这件“第二层的事”往往需要推理框架的算子支持才能落地而芯片厂商提供的编译器比如瑞芯微的RKNN-Toolkit本质上也是第三层的组成部分。所以很多人觉得部署很乱是因为没意识到这个栈本来就是层层嵌套的。2. 推理框架的三板斧图优化、算子调度和内存规划2.1 Execution Provider同一个图不同后端推理框架里最有代表性的设计是ONNX Runtime的Execution ProviderEP机制。所谓EP就是“执行提供者”一个EP对应一类后端设备。你同一个ONNX模型加载时给CPU EP跑它就调度到CPU给CUDA EP跑算子就下沉到GPU给TensorRT EP跑整个图会被交给TensorRT去优化执行。这个机制的意义在于框架层完成了与设备的解耦。算法工程师只需要负责导出正确的ONNX部署工程师根据设备情况选择EP不需要为每个设备重新写一套推理逻辑。但EP也有非常现实的坑。热词里有个“当前设备已离线请确认coze-bridge已连接后重试”我第一反应就是这个桥接思想跟EP一致——中间有一层服务负责把请求转给正确的后端。实际部署时如果你的图里某个算子在你选的EP上没有实现框架不会直接报错而是悄悄回退到CPU执行。你从日志里看到的信息通常是“op XX not supported”然后模型还是能跑只是某个算子慢得离谱。这种“半CPU半加速器”的状态非常难排查因为整体时延看起来没有爆炸但性能就是上不去。2.2 算子融合为什么能同时提速又降内存算子融合是推理框架最核心的优化手段原理一句话就能说清把多个算子的中间结果留在寄存器或缓存里不写回内存。以Conv BN ReLU这个组合为例。如果你三个算子分开跑Conv的输出要写进内存BN要读一次再写一次ReLU再读一次写一次。如果融合成一个算子Conv计算完的结果可以直接在寄存器里做BN和ReLU最后再写回。省下的不只是三次内存读写还有三次kernel dispatch的固定开销。有个很直观的类比你整理房间每整理一个角落就把工具放回工具箱然后再拿出来用下一个角落效率肯定低真正的做法是把要用的工具摆在一起一次干完一片区域的活。算子融合就是干这个事。TensorRT在这一点上做得尤其激进它会把整个ResNet的卷积、激活、残差连接融合成几个大的“kernels”这就是为什么TensorRT跑ResNet能比直接调cuDNN快好几倍。你在深度学习框架里写一百个层TensorRT执行时可能只有二十个真正的计算内核。2.3 静态内存规划让显存复用率翻倍训练框架为了灵活性张量内存是动态分配的PyTorch的缓存分配器会维护一块内存池避免频繁申请释放。但部署场景可以用更狠的方式静态内存规划。因为计算图在部署前是已知的框架可以静态分析每个张量的生命周期算出哪些张量不会同时存活然后把它们放到同一块内存上。一个典型的Transformer推理如果不做静态规划所有激活值都要单独分配空间做了规划之后激活内存能从几百MB降到几十MB。这里要特别提一下LLM推理的场景。你现在用Claude Code或ChatGPT这类Agent时如果本地模型服务是用LM Studio这类工具跑的背后实际上就有一层推理框架在做KV Cache的空闲块管理、beam search的批次调度。热词里“Claude Code调用LMStudio的本地模型”就是这个套路Agent工具通过OpenAI兼容的HTTP接口连到本地推理服务这个服务内部承担了所有内存规划和调度活。很多人以为只是“启动一个模型”其实服务端做了大量你看不到的优化。3. AI编译栈再往下钻从高层算符到设备指令的翻译工程3.1 手写Kernel的瓶颈催生了编译方案推理框架能帮你选好算子实现但解决不了“新算子怎么办”的问题。芯片厂商支持的算子库就那么几百个模型里出现一个不支持的算子你说怎么办很多团队的答案是手写Kernel但这是个不可持续的路。原因很简单算子种类在涨设备指令集也在涨。一个卷积算子在x86上要做AVX512向量化在ARM上要用NEON在GPU上要考虑共享内存和线程编排在NPU上要走硬件硬指令。就算你是一个系统专家为同样一个算子适配四种设备也要好几天。这不是工程师不够努力是这活根本干不完。AI编译器对这个问题的核心思想是算法与调度分离。算法是“这个算子的数学逻辑是什么”调度是“这个算子应该怎么拆分成循环、分块、向量化、线程绑定”。算法写一次调度可以针对不同设备做不同模板。3.2 Relay IR到Tir一个卷积的拆解过程以TVM为例它有两个关键的中间表示层Relay IR负责图级别优化对应前面说的算子融合、常量折叠Tir负责算子内部的循环变换和代码生成。看一个卷积被拆解的过程你就明白了。Relay IR里一个二维卷积算子到了Tir层会被展开成多层嵌套的for循环外层是输出通道中间是输出空间位置内层是输入通道的累加。然后编译器会针对目标设备的缓存大小决定怎么分块tiling、怎么向量化、怎么排布线程。这一层的信息量非常大。同一个卷积在CPU上最优的调度可能是“把输入通道分成两半循环顺序是NCWH”在GPU上最优的调度可能是“每个线程算一个输出点共享内存缓存输入块”。这些知识靠写代码的人是很难全部掌握的所以编译器引入了自动搜索机制。3.3 自动调优用搜索代替经验AutoTVM和Ansor的思路很简单粗暴与其让工程师猜一个最优调度不如把循环顺序、分块大小、向量化宽度、访存策略这些参数组合成一个巨大的搜索空间然后让编译器在真实硬件上跑实验用性能结果反过来挑最优参数。这个搜索空间有多大呢一个普通的卷积算子可能的调度组合轻松超过10的12次方。但Ansor这类系统会先用模板剪枝、再用贝叶斯优化或进化算法先粗搜再精调通常几十分钟就能收敛到接近最优的参数。这套机制的意义在于它把“硬件玄学”变成了“可复现的数值优化”。你不需要预先知道NPU的缓存是几路组相联不需要翻几百页的优化指南只要给足时间编译器自己会试出最优解。热词里“瑞芯微RK3568设备树”之所以被搜得多侧面说明很多工程师在嵌入式Linux上配置硬件资源时已经习惯了“查文档、看手册、手调参数”的模式而AI编译器正在把这套经验主义流程变成自动化。3.4 MLIR把编译器变成超级工厂的模块化流水线MLIRMulti-Level Intermediate Representation是编译器领域的另一个大方向。它的核心贡献是把“中间表示”这件事抽象成了可以自由组合的Dialect也就是不同抽象级别的中间语言。你可以把MLIR理解成一个超级工厂的流水线每个Dialect是车间里的一台机器有的机器负责把高层数学表达变成循环有的机器负责把循环映射到硬件有的机器负责做内存布局优化。不同硬件厂商可以只实现自己需要的那几台机器剩下的复用公共部分。谷歌的XLA、英伟达的TensorRT、AMD的ROCm底层都在往这个方向收敛。但我要提醒一句如果只是做部署你现阶段不一定需要直接写MLIR。它的价值更多体现为“编译器背后的架构”而不是“工程师日常使用的接口”。搞懂MLIR能让你看懂框架日志里那些lowering pass的名字能让你理解为什么一个模型在设备上会生成出奇怪的代码但别一上来就钻进去。4. 从PyTorch到RK3568的一条完整部署链路4.1 导出ONNX的第一个坎动态Shape和自定义算子实际部署项目中第一个拦路虎往往是ONNX导出。torch.onnx.export这件事看着简单但踩过坑的人都知道难点全在细节。第一个坑是动态shape。训练阶段你可以任意改变输入长度但导出ONNX时如果没有指定dynamic_axes导出的模型就把输入尺寸固定死了。我见过一个LSTM模型训练时支持变长序列导出时没设动态维度部署后稍微换个输入长度就报错。解决办法是导出时明确标记动态维度的范围比如torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch, 1: seq_len}, output: {0: batch}} )第二个坑是自定义算子。很多科研模型里写了自定义的autograd.Function这类算子ONNX里没有标准定义。如果你是自用模型有两条路可以走一是用ONNX的Custom Op扩展机制注册实现在推理框架侧单独提供这个算子的kernel二是干脆把这个自定义算子改写成标准算子的组合比如把某个复杂的注意力实现改成多个标准的矩阵乘和Softmax。我一般优先选第二条因为后续在NPU上适配标准算子容易得多。4.2 量化入门FP32到INT8的精度博弈设备端部署基本绕不开量化尤其是RK3568这类边缘盒子NPU对INT8的支持远好于FP16。量化的基本思路是把连续的浮点数值映射到离散的整数公式很简单q round(r / scale) zero_point难点在于scale和zero_point怎么选。最简单的后训练量化PTQ会拿一批校准数据跑一遍模型统计每个张量的数值分布然后确定缩放系数。更稳的方式是量化感知训练QAT在训练阶段就模拟量化误差让模型参数提前适应低精度。实际项目中精度损失跟任务类型强相关。分类任务量化后掉0.5个百分点可能完全无感但回归任务比如工业设备的状态预测对数值敏感量化后的误差会直接放大到输出结果上。热词里提到的“lightgbm回归模型”经常被拿来跟神经网络回归模型对比LightGBM部署不需要量化、对小数据量和高噪声场景很友好但问题是它不能直接用到NPU上。如果你的硬件只有NPU没有CPU算力那还得回到神经网络路上来此时就要把量化误差当作第一优先级来测。4.3 设备树、NPU驱动和算子的“最后一公里”这里必须把“设备树”这个概念讲清楚因为嵌入式部署绕不开它。设备树是嵌入式Linux用来描述硬件资源的数据结构告诉内核哪个设备挂在哪个地址上、用的是哪条中断线、时钟频率是多少。你在部署模型到RK3568上时如果设备树里NPU节点的状态不对比如power-domain没使能NPU驱动根本加载不起来模型自然跑不了。这跟AI编译栈有什么关系关系很大。设备树负责把硬件“点亮”NPU的运行时库负责接收推理框架下发的计算图然后把算子的指令推到硬件上执行。这条链条里任何一环断了模型都跑不起来。我遇到过一次板子的设备树里GPU和NPU共用了一个power-domainNPU在高负载时会把GPU的算力也拖住导致整个系统画面卡顿。这种问题你在纯软件层面排查是查不出来的必须把设备树打开看一眼。算子不支持的问题也会在这一步集中爆发。RKNN-Toolkit在转换模型的时候会报告每个算子的转换情况一旦出现“该算子不支持fallback到CPU”整条pipeline的速度就被拉低。实操中的策略是优先改模型结构把不支持的算子换成NPU友好的标准算子如果实在改不了再接受这个算子走CPU但要给CPU预留足够的算力。4.4 部署验收指标不能只看单次时延部署完第一件事千万别急着宣布成功。要建立一套完整的验收指标包括精度、时延、吞吐、稳定性四个方面。精度好理解就是把量化前后模型的输出对比一遍看误差有没有超出业务容忍范围。时延要看p99而不是平均因为设备端最怕的是偶发卡顿——一次推理从20毫秒变成80毫秒用户体验就会明显变差。我在验收时固定跑1000次推理取p50、p95、p99三个分位三个都要达标才算过。吞吐则要看单位时间内能处理多少请求这在多人并发的场景下尤其重要。稳定性测试往往被忽略但它才是设备端部署的重头戏。很多问题不是模型跑不起来而是跑了几小时后开始出问题。5. 算力量级决定映射策略GPU、边缘NPU和MCU的差异化部署5.1 GPU上的TensorRT层合并与FP16服务器端部署GPU时TensorRT基本是绕不开的。它跟ONNX Runtime的GPU模式最大的差别在于ONNX Runtime只是把算子逐个交给CUDA执行TensorRT会把整个计算图做深度的层合并生成专门优化过的kernel流向。TensorRT的优化还体现在精度选择上。它可以在不损失太多精度的情况下把算子改成FP16这在半精度计算能力强的GPU比如安培架构之后上效果非常明显。实测同一个ResNet-50模型在A10上FP32批处理时延约10毫秒改FP16加TensorRT优化后能降到3毫秒以内。这里也顺带说个经验TensorRT对动态shape的支持没有ONNX Runtime那么宽松。如果你的模型输入尺寸经常变用TensorRT的时候要么限定固定尺寸要么用它的优化配置文件设置几个shape档位。很多团队在GPU上跑得好好的一上生产环境就报shape mismatch基本都是动态输入尺寸没处理好。5.2 边缘盒子RK3568与Jetson的NPU派活策略边缘设备的资源比服务器紧张得多部署策略也要随之调整。RK3568上的NPU算力大约1 TOPSJetson Orin NX的GPU算力可以到100 TOPS这两类设备的部署方式完全不同。RK3568这类中低端NPU对模型结构极其挑剔。它最擅长的是CNN类的规则算子像卷积、池化、全连接对Transformer的注意力机制支持有限。如果你非要把一个大模型塞进去部分算子掉CPU不说内存带宽还会成为瓶颈。这就是为什么热词里“瑞芯微RK3568设备树”被大量搜索——本质上大家都卡在“模型与NPU算子的适配”这道坎上。Jetson的情况好一些Flexible可以跑TensorRTCUDA生态也比较完整但代价是功耗高。边缘部署的本质是在算力、功耗、模型精度之间做三向权衡。我常用的思路是先定功耗墙再定时延目标最后看剩余算力能不能塞下目标模型。塞不下就压缩模型再塞不下就换更小的骨架而不是一上来就优化算子。5.3 MCUSTM32上的TFLite Micro和CMSIS-NN再往下走到MCU级别比如STM32就是完全不同的游戏了。这类设备常常是电池供电、无操作系统、Flash和RAM以KB为单位。STM32的型号不同算力差异很大但即便最强的型号也不适合跑大模型通常就是跑TFLite Micro转换出来的Int8模型或者直接调用CMSIS-NN提供的ARM优化算子。对MCU部署我不建议引入太重的东西TFLite Micro足够了。整个流程是PyTorch训练出模型 → 量化到int8 → 转成TFLite格式 → 用TFLite Micro在MCU上执行。算子只能用标准数学函数做不了复杂算子融合所以模型设计阶段就要保持“朴素”用标准卷积、标准激活、标准全连接。MCU还有个特点是外设交互。热词里“stm32 如何做usb设备”跟模型部署结合得更紧因为MCU上跑推理往往只是系统的一部分它还得通过USB或串口跟主机通信负责把推理结果上报。这意味着MCU的内存规划要同时考虑推理和通信两块不能把缓冲区都分给模型。我在STM32项目里的做法是双缓冲一块DMA内存专门做USB收发另一块做推理输入输出两者隔离避免CPU在推理时还要处理中断。5.4 异构调度让CPU、GPU、NPU各干各的活真实的设备端系统很少只有一个计算单元RK3568就同时有CPU、GPU、NPU。异构调度解决的核心问题是每个算子应该派到哪个单元去执行。这个“派活”的逻辑要综合考虑算子的特性、设备当前的负载、数据搬运的成本。比如一个模型里有Embedding层、卷积层、全连接层你可能想把Embedding放到CPU它访存密集、计算稀疏把卷积放到NPU它并行度高把全连接放到CPU或GPU看芯片的矩阵乘法能力。但异构系统最大的坑是数据搬运成本。每次跨计算单元交换数据都要通过总线拷贝这个开销可能比算子本身还贵。所以异构调度不是简单地把算子分到最快的地方而是要考虑到整体流水线的数据局部性。有些推理框架允许你配“跨设备融合”的选项本质上就是通过算子调度减少搬运次数。工业场景里热词“modbus、opc ua协议读取plc、传感器、数控机床等设备的运行状态数据”让我想到PLC设备的数据采集也是异构系统的一部分。模型推理的结果往往要写回PLC或者通过Modbus返回到SCADA系统那这个数据链路跟模型推理链路就是串在一起的。做部署计划时一定要把采集频率、通信带宽、推理时延这三者放到一起估算瓶颈。6. 部署只是开始时延监测、滤波预处理与设备老化回归测试6.1 设备端“软故障”内存碎片、温度降频与精度漂移模型部署上线后真正的挑战才开始。硬件设备经过数周持续运行后会出现一类“软故障”进程没崩但性能越来越差。第一种是内存碎片化。推理框架虽然做了静态内存规划但如果你的设备上还跑着别的进程频繁申请释放小块内存会让堆碎片化最终导致某次大块内存申请失败。我排查过一起现场问题设备跑了两周后突然推理失败重启又好了最终定位就是内存碎片。第二种是温度降频。NPU和GPU在高负载下温度升高触发降频保护推理时延缓慢爬升。这种问题看日志很难发现因为它不报错只有连续监控时延才能看到趋势。第三种是精度漂移。工业场景里传感器老化、环境变化都可能让输入数据分布偏移模型在训练时见过的分布跟现场完全不匹配输出开始偏离。这不是模型参数变了是数据变了。这也是为什么部署之后必须保留一个持续监控机制。6.2 用LightGBM回归模型给推理时延做体检既然模型推理过程有这么多变量影响时延我们干脆用机器学习来“预测时延”。我自己做过一套这样的监控采集每个推理请求的输入长度、batch size、设备温度、NPU占用率、内存余量这些特征目标是推理时延。样本量足够后用LightGBM回归模型训练一个时延预测器。训练好之后这个模型的作用有两个一是在线异常检测实际时延比预测值高出很多时说明设备状态有问题立即触发告警二是容量规划预测在给定并发下系统会不会超时提前扩容或限流。LightGBM在这里比神经网络更适合原因很简单特征维度低数据量不大训练快可解释性强。你可以直接看特征重要性搞清楚到底哪个因素对时延影响最大——这在调优设备配置时非常有用。我在几个项目里用这个方案准确率能到95%以上。6.3 滑动窗口滤波设备端信号预处理的标配部署到IoT设备上的模型输入信号往往很脏。传感器数据里掺着噪声、毛刺、偶发尖峰直接喂给模型会影响判断。滑动窗口滤波是解决这类问题最实用的手段。具体做法是维护一个固定长度的窗口每次新数据进来窗口滑动一格把窗口内的数据做平均或加权平均。比如一个窗口大小为5的滑动平均滤波当前输出是最近5个采样点的均值。这个方法的优点是完全不需要额外算力几个移位的开销适合在MCU上实时跑。热词里“滑动窗口滤波模型”其实是个更大的概念它不仅指信号预处理还指时间序列预测里的“用过去N个时间步预测未来”。这两种用法在工业设备预测项目里都会遇到。部署时要注意训练时用的窗口长度和实时推理时用的窗口长度必须一致否则模型输入分布会变。6.4 自动化老化测试脚本7x24小时跑出来的稳定性数据设备的稳定性不能靠“感觉稳定”必须用自动化脚本跑出来的数据说话。我一般会在交付前跑一轮7x24小时的老化测试脚本实现四件事按业务频率循环发推理请求模拟真实负载每30秒记录一次时延、内存、温度、NPU占用率跑一批固定输入集每1小时做一次精度比对检查精度漂移如果出现时延突刺或推理失败自动记录现场信息并保留log这个脚本的价值在于它能复现我前面说的三类软故障。我遇到过最典型的一次是设备跑到第19小时的时候内存占用顶到了上限之后每次推理都开始抖动但硬件本身没报任何错误。如果没有老化测试的数据曲线这个问题根本没法向客户解释。写这个脚本的时候有几个细节要注意。一是时间戳要对齐时延记录、温度记录、日志片段必须要能关联到同一个时间点否则事后定位就是灾难。二是负载要带随机性固定频率的请求测不出真实问题最好加入随机间隔和偶发峰值。三是断电恢复要测很多设备的稳定性问题出在异常重启后的状态恢复上设备树加载、NPU初始化、推理框架冷启动这些都要在脚本里模拟。到这里从名字到落地推理框架和AI编译栈这条链路大概就串起来了。说一点个人体会遇到部署性能问题先别急着搬TensorRT或者上TVM多数情况下的瓶颈在框架配置和模型结构设计上。你先把ONNX Runtime的图优化、内存规划这些机制吃透再把具体的设备算子约束摸清基本能解决90%的问题。自动调优和MLIR这些高级玩法留给真正上了规模、追求极致性能的项目去用。工具不是越复杂越好能稳定扛住业务的那一套就是最好的方案。