1. 先把算子和框架放在同一张桌子上聊大模型开发做到一定阶段几乎都会撞上同一个问题算子框架怎么选。模型训练慢、推理延迟高、显存不够用这些问题表面上看是显卡不行、显存不够但往深了查很多根子上的问题出在算子层面。大模型跑起来的每一秒背后都是成千上万次算子调用算子实现得好不好、框架调度得到不到位直接决定你的训练速度、推理时延和成本。先说清楚本文定位这是一篇面向大模型算法工程师、AI平台开发者和做模型部署落地同学的选型参考文章。核心讨论范围是大模型场景下算子框架如何做技术选型包括算子层的基本概念、主流框架分类、决策路径、优化思路和避坑经验。不写谁吊打谁只写每个方案适合什么场景、有什么代价。很多人一上来就纠结用PyTorch还是用厂商自带的算子库自己写算子还是调现成的先别急这类问题根本没有统一答案得回到你的具体场景里看。下面我按几个层次把这件事拆开讲。1.1 算子到底是个什么东西算子是个经常被提起、但很少被真正讲透的词。把一次深度学习计算掰开看最终都是对张量做变换矩阵乘法、卷积、归一化、激活、Softmax、注意力计算、位置编码……每一个计算步骤就是一个算子。深度模型本质就是一堆算子的编排组合。PyTorch、TensorFlow这些框架写的模型最终要落到硬件上执行靠的就是算子。打个比方算子就像菜谱里的刀法。切丝、切丁、剁蓉是刀法红烧、清蒸、爆炒是烹饪流程。模型结构是菜谱算子框架就是你手上那套刀。刀的好坏不只决定速度还决定了你做得出来做不出来某些菜式。在大模型时代算子的粒度在变粗。过去一个算子可能就是一个矩阵乘法现在一个大模型里经常出现FlashAttention这种融合了访存、运算、mask处理的复合算子。算子数量看起来少了但每个算子更复杂也更需要精细优化。1.2 大模型给算子框架的加试题大模型给算子框架提出的要求和当年的CV小模型完全不同。首先是内存压力。大模型动辄几十亿、上百亿参数激活值巨大一个不合适的算子实现就可能让显存多占几个GB。很多时候训练OOM不是因为模型太大而是某个算子的临时缓冲区分配得太随意。其次是并行规模。大模型的训练已经离不开张量并行、流水并行和数据并行这就意味着算子不只要单设备跑得快还要能在多卡多机之间做高效的集合通信配合。通信算子和计算算子的重叠策略直接决定分布式训练的扩展效率。再次是精度控制。混合精度训练中哪些算子用FP16算哪些用BF16哪些必须回到FP32这些细节对收敛质量影响极大。同一个算子在不同框架里的精度行为可能不一样这也是一种隐藏的兼容性坑。再就是图优化能力。大模型前向计算图很重如果框架能对计算图做算子融合、内存复用、死代码消除收益巨大。这一点在推理阶段尤其重要。2. 主流算子框架全景每个方案都有自己的主场市面上的算子相关框架其实可以分成几个阵营。我给它们排了一下方便你在做选型时快速对号入座。2.1 PyTorch原生算子体系生态最全但不是万能PyTorch到现在依然是AI训练的主战场它的原生算子体系覆盖了你能想到的绝大多数需求。ATen是PyTorch的C张量库提供了几百个标准算子torch.nn里的模块绝大部分都封装了对应的算子。因为用户基数太大PyTorch的算子实现被各种场景反复锤炼过稳定性是相当高的。训练一个模型90%的算子可以直接用PyTorch自带实现不用动脑。但PyTorch原生的算子在推理场景里常常不是最优解。原因是它的算子调度有额外的Python和Dispatch开销算子内部实现也偏通用策略对特定硬件、特定形状未必做极致优化。你会发现同一张卡跑同一个模型用PyTorch推理比用TensorRT慢不少就是因为TensorRT做了针对性优化。另外PyTorch的卷积、矩阵乘其实底层调用了厂商库比如英伟达GPU上用的就是cuDNN、cuBLAS。所以它更多扮演调度者和图生成器真正的性能大头还是在厂商算子库里。理解这一点很重要你在PyTorch里改一个算子实现可能是绕过了cuDNN的优化反而变慢了。2.2 厂商专用算子栈性能天花板在这里再往下就是厂商自己的算子栈。英伟达这边cuDNN、cuBLAS、CUTLASS、TensorRT这些构成了CUDA生态的算子家族。CUTLASS尤其值得关注它提供了像模板一样的算子实现框架帮你基于CUDA C写出高性能的矩阵乘和卷积。很多量化推理引擎里的核心算子就是用CUTLASS改写出来的。昇腾这边对应的是CANNCompute Architecture for Neural Networks。CANN里提供了算子开发工具链包括AscendCL编程接口、TBE算子开发框架等。国内大模型私有化部署昇腾平台的占比越来越大CANN算子的适配和优化就成了很多团队的必修课。这些厂商算子栈最大的优势是性能。同样的注意力算子在PyTorch原生实现里访存效率一般而在厂商库或定制核心里经常会利用上各类硬件特性比如Tensor Core、向量指令、高速缓存或专用的矩阵计算单元。劣势同样明显开发门槛高绑平台学习成本不低。2.3 推理优化框架面向部署的精装修推理场景是算子框架另一个主战场。TensorRT是英伟达官方推理优化库它的核心能力是把训练好的PyTorch模型转成推理引擎在这个过程中做算子融合比如把ConvBiasReLU合成一个算子、层间融合、精度校准、内核自动调优。ONNXRuntime的定位不太一样它走的是中间表示的路线。先把模型导出成ONNX格式再交给runtime执行。好处是框架中立可选的执行提供方非常多可以换CPU、GPU、NPU的执行后端。在实际选型时TensorRT更适合英伟达GPU推理场景ONNXRuntime更强的问题是通用性能不一定极致。如果模型部署在昇腾设备上就优先考虑CANN配套的推理方案硬套TensorRT路径是行不通的。2.4 表格对比做一张速查表帮你快速抓重点。算子框架/方案归属阵营典型代表擅长的场景主要代价PyTorch原生算子通用框架层ATen、torch.nn训练、快速验证、学术研究推理性能一般算子调度开销大GPU厂商算子库硬件厂商层cuDNN、cuBLAS、CUTLASS高端GPU上的训练推理性能极限绑定硬件开发门槛高NPU厂商算子栈硬件厂商层CANN、AscendCL昇腾平台训练/推理部署工具链成熟度仍在追赶中通用推理优化框架部署优化层TensorRT、ONNXRuntime模型部署、服务化推理转模型有兼容性风险部分算子不支持轻量级运行框架极简部署层AirLLM等小显存设备大模型推理性能一般适合特定场景这种分层思维很重要不同层级的框架可以配合使用未必是你死我活的关系。3. 大模型算子框架选型的决策路径做选型这件事本质上是在约束条件里找最优解。我梳理了一套相对通用的决策路径按这个顺序往下走不会太偏。3.1 按场景选框架训练优先推理次之先分清楚你是做训练还是做推理。训练阶段核心诉求是灵活性、分布式扩展能力和生态兼容性。绝大多数情况下训练底座可以围绕PyTorch原生算子生态来搭。原因是大模型的训练逻辑复杂要改模型结构、试损失函数、加分布式策略PyTorch这套体系让你能在这个维度上快速折腾。但训练也不全是PyTorch的天下。当你把训练跑在昇腾上就需要考虑PyTorch和CANN的适配方案。华为官方会提供适配版的torch_npu让PyTorch算子调度落到昇腾算子上。这个路径的坑在于算子兼容性——有些PyTorch操作在torch_npu里还没有实现或者实现有偏差。推理阶段则优先考虑TensorRT、ONNXRuntime或者厂商配套的推理框架。这时候性能优先灵活性可以让位。曾经有个项目里训练部分用的PyTorch推理部分导成ONNX再转TensorRT中间模型算子不兼容调了整整两天。这个兼容性问题是部署环节里最高频的坑。3.2 按硬件选框架手里有什么牌决定怎么出牌硬件约束在选型里的权重比你想象的大得多。用英伟达GPU主流选择自然是CUDA生态。训练用PyTorch底下的cuDNN/cuBLAS推理用TensorRT需要极致性能再上CUTLASS手写算子。这条路径最成熟社区资料最丰富。平台如果切到昇腾就必须接纳CANN这套体系。好消息是CANN在持续迭代基础算子的覆盖率越来越高坏消息是大模型里的一些新算子比如各种FlashAttention变体在CANN上的支持速度可能滞后。你在英伟达GPU上能改一行代码调度的算子迁移到昇腾可能就得翻算子文档甚至用TBE自己补算子。如果遇到国产化要求、私有化部署又选了非英伟达硬件我的建议是选型时先做算子兼容性清单评审把模型里的关键算子逐个对照目标平台的支持情况列出来不要等部署的时候再发现缺算子。3.3 按开发方式选框架什么时候必须自己写算子这是另一个高频问题自己写算子还是调框架自带的盘点一下场景再判断。如果只是拼接模型、改改结构那不需要自己写算子直接用框架原生算子组合就行。自己写算子的收益在常规场景下根本体现不出来反而带来维护负担和精度一致性风险。真正需要自己写算子的场景有三类一是性能瓶颈很明确某个算子在剖析图里占了50%以上耗时且框架自带实现无法满足要求二是目标硬件上没有现成实现比如新硬件上部署Transformer动不动缺一个算子三是你想把多步计算融合成单算子减少访存和调度开销——这在推理优化里很常见。写算子这件事不是会写C就能做。往上要做性能剖析和硬件特性分析往下要对齐框架的函数签名、自动微分规则、动态形状支持。尤其是大模型场景下算子动辄涉及多维张量的布局、切分策略自研算子的工程量要按人周级别来评估不能拍脑袋说两天干完。4. 算子性能优化的核心思路与一条实操路径选定框架只是第一步。真正拉开性能差距的是算子层面的优化做不做得到位。这章节直接用一条实操路径来说清楚发现问题、定位热点、选择优化策略直到落地实现。4.1 先剖析再动手找到真正该优化的算子很多人一上来就找各种优化技巧结果优化了半天整体没快多少。原因很简单改了个不占时间的算子收益当然趋近于零。正确的第一步永远是性能剖析。PyTorch自带torch.profiler可以直接打出每个算子的耗时和显存占用。跑一小段训练或推理拿到profile记录以后把耗时排名前几的算子列出来——这就是你真正要优化的目标。注意别只看单个算子耗时还要看算子的调用次数。一个被调用一万次的单算子即使单次只要几微秒积累起来也是大头。在分布式训练场景里还有一类容易被忽略的耗时集合通信。通信算子在profile里往往显示为NCCL/AllReduce一类优化思路主要是让通信和计算重叠或者调整通信切分方式。这类问题在算子层面靠改实现是解决不了的要去改分布式策略。4.2 选择优化策略从融合到精度到调优定位到热点算子后选择优化策略的顺序很重要。第一步考虑算子融合。把两个相邻且可以合并的算子合成一个减少中间数据的写回和再读取。比如计算图中的矩阵乘偏置加ReLU激活这三个算子在底层可以合成一个融合算子。像TensorRT做的算子融合原理就是这个把计算图里的小节点合并成大节点减少调度开销和内存访问。第二步检查精度策略。确认算子内哪些计算必须用FP32哪些可以放到BF16/FP16。如果是推理场景还可以做INT8量化把权重和激活值映射到8位整数。量化算子如果实现得好加速比往往比浮点算子优化更可观。但要小心量化后的精度损失一般需要一小批校准数据来调整scale和zero point。第三步再做微架构级调优。这一步涉及向量化、循环分块、数据布局重排等细节。举个例子一个矩阵乘算子循环的嵌套顺序直接把数据在cache里的命中率拉高一个档次性能能差出好几倍。这一步很硬核做完收益也最实。4.3 在CANN上做算子优化的参考流程国内大模型部署现在很多平台走昇腾简要说一下CANN的算子优化实操路径供你在同类需求时参考。CANN平台算子开发通常走TBE算子开发框架。流程大致是先定义算子的计算逻辑然后用TBE提供的API实现计算函数对输入张量进行形状和数据类型推导再把算子注册到框架中。跑通“冒烟阶段”后用profiling工具看耗的是计算时间还是搬运时间。昇腾的计算单元和GPU的CUDA Core策略差异很大。它的矩阵计算单元适合大块的乘加计算多核调度和内存搬运策略要跟着硬件走。算子写出来以后记得做多shape测试因为静态shape和动态shape下的调度路径不同不能只测一组形状就交差。4.4 一个可参考的算子融合示例用矩阵乘偏置激活的融合来做个示例这是大模型里最常见也最典型的融合场景。未融合时三个算子各自执行MatMul(A, B)输出临时张量TAdd(T, bias)输出临时张量UReLU(U)输出最终结果。中间有两个临时张量要写显存再读回来每次读写都是一次完整的内存往返。对一个超大K维度的大矩阵乘来说这些往返浪费相当可观。融合之后变成FusedMatMulBiasReLU(A, B, bias)一次算完矩阵乘在寄存器里或片上缓存里完成加偏置和激活直接把最终结果写回。中间那两次显存访问删掉了。效果就是延迟降低、显存占用减少。在推理时这样的融合往往配合内核自动调优再做一步Tactic选择找出当前GPU/输入形状下的最优kernel配置。你不用一上来就自己实现这种融合。TensorRT在转模型时会自动干这件事你只需要理解这个融合逻辑在手动优化的时候也用得上。5. 常见问题与排查技巧实录这章节写一些我在实际项目里遇到过的典型问题帮你少踩坑。5.1 算子兼容性速查表现象可能原因排查思路模型转ONNX失败某个自定义算子没有注册导出函数查看onnx导出日志定位到具体算子用onnxregistry或手动添加自定义映射推理结果和训练不一致算子在推理框架里被替换成了数学等价的低精度版本逐层输出对比保持FP32的算子优先排查量化层昇腾跑PyTorch报算子不支持torch_npu里没有对应算子实现查CANN算子支持列表考虑TBE自研算子或换等价算子组合分布式训练扩展率上不去集合通信算子和计算算子没有重叠用NCCL/通信profile定位等待时间调整通信切分和流优先级显存占用异常高某个算子申请了大块临时buffer用显存profiler观察峰值时的分配栈尝试算子融合或内存复用建议团队在项目一开始就建立算子兼容性checklist把模型里所有算子列出来逐个标注使用框架/Target平台的支持情况。这张表后面会在部署、迁移、优化三个阶段反复用。5.2 三条独家避坑心得第一不要迷信框架越大越好。完整的PyTorch算子栈在训练时确实是好东西但它在大模型推理服务里可能引入不必要的Python层开销。如果你的业务是稳定服务某个固定结构的模型可以考虑跳过PyTorch直接走TensorRT或ONNXRuntime这类执行框架效果往往更稳。第二算子优化要设置收益阈值再动手。我的习惯是先看一个算子在整体耗时里的占比。如果占不到20%先别花精力去优化它真正要动手的是把住排名前几的算子和通信环节。优化出一个占3%耗时算子的成就感弥补不了后面联调的时间成本。第三大模型上下文的公开热词里大模型微调相关的算子优化最容易忽视。微调场景和预训练场景的算子分布差异很大微调通常batch较小、激活值比例高有些算子在小batch下的kernel并不是最优实现。如果你只照着预训练的标准做算子选型微调性能可能差上一截。5.3 一点后续建议如果项目已经跑起来了但又想做算子这块的深度优化建议考虑构建一个轻量的算子基准测试集。把模型里高频使用的算子摘出来固定输入shape、精度和硬件环境做成一个benchmark脚本。后面每次换框架版本、换硬件、或者改算子实现都可以对着这个基准测一测。这个基准测试集的维护成本不高但能在版本升级的时候帮你挡住很多隐性性能回退。根据我个人的实操体会算子框架的选型从来不是一次性的决定。项目早期选什么框架可能只是影响开发效率到了优化部署阶段算子层面的技术深度才真正决定项目能做到什么程度。提前摸清楚算子层的事情比等项目出瓶颈了再临时补课省下来的时间是以周为单位的。 
长河君 · 2026/10/27 8:24:37 · 阅读 1.2k 大模型算子框架选型：PyTorch、CANN与TensorRT实战解析
1. 先把算子和框架放在同一张桌子上聊大模型开发做到一定阶段几乎都会撞上同一个问题算子框架怎么选。模型训练慢、推理延迟高、显存不够用这些问题表面上看是显卡不行、显存不够但往深了查很多根子上的问题出在算子层面。大模型跑起来的每一秒背后都是成千上万次算子调用算子实现得好不好、框架调度得到不到位直接决定你的训练速度、推理时延和成本。先说清楚本文定位这是一篇面向大模型算法工程师、AI平台开发者和做模型部署落地同学的选型参考文章。核心讨论范围是大模型场景下算子框架如何做技术选型包括算子层的基本概念、主流框架分类、决策路径、优化思路和避坑经验。不写谁吊打谁只写每个方案适合什么场景、有什么代价。很多人一上来就纠结用PyTorch还是用厂商自带的算子库自己写算子还是调现成的先别急这类问题根本没有统一答案得回到你的具体场景里看。下面我按几个层次把这件事拆开讲。1.1 算子到底是个什么东西算子是个经常被提起、但很少被真正讲透的词。把一次深度学习计算掰开看最终都是对张量做变换矩阵乘法、卷积、归一化、激活、Softmax、注意力计算、位置编码……每一个计算步骤就是一个算子。深度模型本质就是一堆算子的编排组合。PyTorch、TensorFlow这些框架写的模型最终要落到硬件上执行靠的就是算子。打个比方算子就像菜谱里的刀法。切丝、切丁、剁蓉是刀法红烧、清蒸、爆炒是烹饪流程。模型结构是菜谱算子框架就是你手上那套刀。刀的好坏不只决定速度还决定了你做得出来做不出来某些菜式。在大模型时代算子的粒度在变粗。过去一个算子可能就是一个矩阵乘法现在一个大模型里经常出现FlashAttention这种融合了访存、运算、mask处理的复合算子。算子数量看起来少了但每个算子更复杂也更需要精细优化。1.2 大模型给算子框架的加试题大模型给算子框架提出的要求和当年的CV小模型完全不同。首先是内存压力。大模型动辄几十亿、上百亿参数激活值巨大一个不合适的算子实现就可能让显存多占几个GB。很多时候训练OOM不是因为模型太大而是某个算子的临时缓冲区分配得太随意。其次是并行规模。大模型的训练已经离不开张量并行、流水并行和数据并行这就意味着算子不只要单设备跑得快还要能在多卡多机之间做高效的集合通信配合。通信算子和计算算子的重叠策略直接决定分布式训练的扩展效率。再次是精度控制。混合精度训练中哪些算子用FP16算哪些用BF16哪些必须回到FP32这些细节对收敛质量影响极大。同一个算子在不同框架里的精度行为可能不一样这也是一种隐藏的兼容性坑。再就是图优化能力。大模型前向计算图很重如果框架能对计算图做算子融合、内存复用、死代码消除收益巨大。这一点在推理阶段尤其重要。2. 主流算子框架全景每个方案都有自己的主场市面上的算子相关框架其实可以分成几个阵营。我给它们排了一下方便你在做选型时快速对号入座。2.1 PyTorch原生算子体系生态最全但不是万能PyTorch到现在依然是AI训练的主战场它的原生算子体系覆盖了你能想到的绝大多数需求。ATen是PyTorch的C张量库提供了几百个标准算子torch.nn里的模块绝大部分都封装了对应的算子。因为用户基数太大PyTorch的算子实现被各种场景反复锤炼过稳定性是相当高的。训练一个模型90%的算子可以直接用PyTorch自带实现不用动脑。但PyTorch原生的算子在推理场景里常常不是最优解。原因是它的算子调度有额外的Python和Dispatch开销算子内部实现也偏通用策略对特定硬件、特定形状未必做极致优化。你会发现同一张卡跑同一个模型用PyTorch推理比用TensorRT慢不少就是因为TensorRT做了针对性优化。另外PyTorch的卷积、矩阵乘其实底层调用了厂商库比如英伟达GPU上用的就是cuDNN、cuBLAS。所以它更多扮演调度者和图生成器真正的性能大头还是在厂商算子库里。理解这一点很重要你在PyTorch里改一个算子实现可能是绕过了cuDNN的优化反而变慢了。2.2 厂商专用算子栈性能天花板在这里再往下就是厂商自己的算子栈。英伟达这边cuDNN、cuBLAS、CUTLASS、TensorRT这些构成了CUDA生态的算子家族。CUTLASS尤其值得关注它提供了像模板一样的算子实现框架帮你基于CUDA C写出高性能的矩阵乘和卷积。很多量化推理引擎里的核心算子就是用CUTLASS改写出来的。昇腾这边对应的是CANNCompute Architecture for Neural Networks。CANN里提供了算子开发工具链包括AscendCL编程接口、TBE算子开发框架等。国内大模型私有化部署昇腾平台的占比越来越大CANN算子的适配和优化就成了很多团队的必修课。这些厂商算子栈最大的优势是性能。同样的注意力算子在PyTorch原生实现里访存效率一般而在厂商库或定制核心里经常会利用上各类硬件特性比如Tensor Core、向量指令、高速缓存或专用的矩阵计算单元。劣势同样明显开发门槛高绑平台学习成本不低。2.3 推理优化框架面向部署的精装修推理场景是算子框架另一个主战场。TensorRT是英伟达官方推理优化库它的核心能力是把训练好的PyTorch模型转成推理引擎在这个过程中做算子融合比如把ConvBiasReLU合成一个算子、层间融合、精度校准、内核自动调优。ONNXRuntime的定位不太一样它走的是中间表示的路线。先把模型导出成ONNX格式再交给runtime执行。好处是框架中立可选的执行提供方非常多可以换CPU、GPU、NPU的执行后端。在实际选型时TensorRT更适合英伟达GPU推理场景ONNXRuntime更强的问题是通用性能不一定极致。如果模型部署在昇腾设备上就优先考虑CANN配套的推理方案硬套TensorRT路径是行不通的。2.4 表格对比做一张速查表帮你快速抓重点。算子框架/方案归属阵营典型代表擅长的场景主要代价PyTorch原生算子通用框架层ATen、torch.nn训练、快速验证、学术研究推理性能一般算子调度开销大GPU厂商算子库硬件厂商层cuDNN、cuBLAS、CUTLASS高端GPU上的训练推理性能极限绑定硬件开发门槛高NPU厂商算子栈硬件厂商层CANN、AscendCL昇腾平台训练/推理部署工具链成熟度仍在追赶中通用推理优化框架部署优化层TensorRT、ONNXRuntime模型部署、服务化推理转模型有兼容性风险部分算子不支持轻量级运行框架极简部署层AirLLM等小显存设备大模型推理性能一般适合特定场景这种分层思维很重要不同层级的框架可以配合使用未必是你死我活的关系。3. 大模型算子框架选型的决策路径做选型这件事本质上是在约束条件里找最优解。我梳理了一套相对通用的决策路径按这个顺序往下走不会太偏。3.1 按场景选框架训练优先推理次之先分清楚你是做训练还是做推理。训练阶段核心诉求是灵活性、分布式扩展能力和生态兼容性。绝大多数情况下训练底座可以围绕PyTorch原生算子生态来搭。原因是大模型的训练逻辑复杂要改模型结构、试损失函数、加分布式策略PyTorch这套体系让你能在这个维度上快速折腾。但训练也不全是PyTorch的天下。当你把训练跑在昇腾上就需要考虑PyTorch和CANN的适配方案。华为官方会提供适配版的torch_npu让PyTorch算子调度落到昇腾算子上。这个路径的坑在于算子兼容性——有些PyTorch操作在torch_npu里还没有实现或者实现有偏差。推理阶段则优先考虑TensorRT、ONNXRuntime或者厂商配套的推理框架。这时候性能优先灵活性可以让位。曾经有个项目里训练部分用的PyTorch推理部分导成ONNX再转TensorRT中间模型算子不兼容调了整整两天。这个兼容性问题是部署环节里最高频的坑。3.2 按硬件选框架手里有什么牌决定怎么出牌硬件约束在选型里的权重比你想象的大得多。用英伟达GPU主流选择自然是CUDA生态。训练用PyTorch底下的cuDNN/cuBLAS推理用TensorRT需要极致性能再上CUTLASS手写算子。这条路径最成熟社区资料最丰富。平台如果切到昇腾就必须接纳CANN这套体系。好消息是CANN在持续迭代基础算子的覆盖率越来越高坏消息是大模型里的一些新算子比如各种FlashAttention变体在CANN上的支持速度可能滞后。你在英伟达GPU上能改一行代码调度的算子迁移到昇腾可能就得翻算子文档甚至用TBE自己补算子。如果遇到国产化要求、私有化部署又选了非英伟达硬件我的建议是选型时先做算子兼容性清单评审把模型里的关键算子逐个对照目标平台的支持情况列出来不要等部署的时候再发现缺算子。3.3 按开发方式选框架什么时候必须自己写算子这是另一个高频问题自己写算子还是调框架自带的盘点一下场景再判断。如果只是拼接模型、改改结构那不需要自己写算子直接用框架原生算子组合就行。自己写算子的收益在常规场景下根本体现不出来反而带来维护负担和精度一致性风险。真正需要自己写算子的场景有三类一是性能瓶颈很明确某个算子在剖析图里占了50%以上耗时且框架自带实现无法满足要求二是目标硬件上没有现成实现比如新硬件上部署Transformer动不动缺一个算子三是你想把多步计算融合成单算子减少访存和调度开销——这在推理优化里很常见。写算子这件事不是会写C就能做。往上要做性能剖析和硬件特性分析往下要对齐框架的函数签名、自动微分规则、动态形状支持。尤其是大模型场景下算子动辄涉及多维张量的布局、切分策略自研算子的工程量要按人周级别来评估不能拍脑袋说两天干完。4. 算子性能优化的核心思路与一条实操路径选定框架只是第一步。真正拉开性能差距的是算子层面的优化做不做得到位。这章节直接用一条实操路径来说清楚发现问题、定位热点、选择优化策略直到落地实现。4.1 先剖析再动手找到真正该优化的算子很多人一上来就找各种优化技巧结果优化了半天整体没快多少。原因很简单改了个不占时间的算子收益当然趋近于零。正确的第一步永远是性能剖析。PyTorch自带torch.profiler可以直接打出每个算子的耗时和显存占用。跑一小段训练或推理拿到profile记录以后把耗时排名前几的算子列出来——这就是你真正要优化的目标。注意别只看单个算子耗时还要看算子的调用次数。一个被调用一万次的单算子即使单次只要几微秒积累起来也是大头。在分布式训练场景里还有一类容易被忽略的耗时集合通信。通信算子在profile里往往显示为NCCL/AllReduce一类优化思路主要是让通信和计算重叠或者调整通信切分方式。这类问题在算子层面靠改实现是解决不了的要去改分布式策略。4.2 选择优化策略从融合到精度到调优定位到热点算子后选择优化策略的顺序很重要。第一步考虑算子融合。把两个相邻且可以合并的算子合成一个减少中间数据的写回和再读取。比如计算图中的矩阵乘偏置加ReLU激活这三个算子在底层可以合成一个融合算子。像TensorRT做的算子融合原理就是这个把计算图里的小节点合并成大节点减少调度开销和内存访问。第二步检查精度策略。确认算子内哪些计算必须用FP32哪些可以放到BF16/FP16。如果是推理场景还可以做INT8量化把权重和激活值映射到8位整数。量化算子如果实现得好加速比往往比浮点算子优化更可观。但要小心量化后的精度损失一般需要一小批校准数据来调整scale和zero point。第三步再做微架构级调优。这一步涉及向量化、循环分块、数据布局重排等细节。举个例子一个矩阵乘算子循环的嵌套顺序直接把数据在cache里的命中率拉高一个档次性能能差出好几倍。这一步很硬核做完收益也最实。4.3 在CANN上做算子优化的参考流程国内大模型部署现在很多平台走昇腾简要说一下CANN的算子优化实操路径供你在同类需求时参考。CANN平台算子开发通常走TBE算子开发框架。流程大致是先定义算子的计算逻辑然后用TBE提供的API实现计算函数对输入张量进行形状和数据类型推导再把算子注册到框架中。跑通“冒烟阶段”后用profiling工具看耗的是计算时间还是搬运时间。昇腾的计算单元和GPU的CUDA Core策略差异很大。它的矩阵计算单元适合大块的乘加计算多核调度和内存搬运策略要跟着硬件走。算子写出来以后记得做多shape测试因为静态shape和动态shape下的调度路径不同不能只测一组形状就交差。4.4 一个可参考的算子融合示例用矩阵乘偏置激活的融合来做个示例这是大模型里最常见也最典型的融合场景。未融合时三个算子各自执行MatMul(A, B)输出临时张量TAdd(T, bias)输出临时张量UReLU(U)输出最终结果。中间有两个临时张量要写显存再读回来每次读写都是一次完整的内存往返。对一个超大K维度的大矩阵乘来说这些往返浪费相当可观。融合之后变成FusedMatMulBiasReLU(A, B, bias)一次算完矩阵乘在寄存器里或片上缓存里完成加偏置和激活直接把最终结果写回。中间那两次显存访问删掉了。效果就是延迟降低、显存占用减少。在推理时这样的融合往往配合内核自动调优再做一步Tactic选择找出当前GPU/输入形状下的最优kernel配置。你不用一上来就自己实现这种融合。TensorRT在转模型时会自动干这件事你只需要理解这个融合逻辑在手动优化的时候也用得上。5. 常见问题与排查技巧实录这章节写一些我在实际项目里遇到过的典型问题帮你少踩坑。5.1 算子兼容性速查表现象可能原因排查思路模型转ONNX失败某个自定义算子没有注册导出函数查看onnx导出日志定位到具体算子用onnxregistry或手动添加自定义映射推理结果和训练不一致算子在推理框架里被替换成了数学等价的低精度版本逐层输出对比保持FP32的算子优先排查量化层昇腾跑PyTorch报算子不支持torch_npu里没有对应算子实现查CANN算子支持列表考虑TBE自研算子或换等价算子组合分布式训练扩展率上不去集合通信算子和计算算子没有重叠用NCCL/通信profile定位等待时间调整通信切分和流优先级显存占用异常高某个算子申请了大块临时buffer用显存profiler观察峰值时的分配栈尝试算子融合或内存复用建议团队在项目一开始就建立算子兼容性checklist把模型里所有算子列出来逐个标注使用框架/Target平台的支持情况。这张表后面会在部署、迁移、优化三个阶段反复用。5.2 三条独家避坑心得第一不要迷信框架越大越好。完整的PyTorch算子栈在训练时确实是好东西但它在大模型推理服务里可能引入不必要的Python层开销。如果你的业务是稳定服务某个固定结构的模型可以考虑跳过PyTorch直接走TensorRT或ONNXRuntime这类执行框架效果往往更稳。第二算子优化要设置收益阈值再动手。我的习惯是先看一个算子在整体耗时里的占比。如果占不到20%先别花精力去优化它真正要动手的是把住排名前几的算子和通信环节。优化出一个占3%耗时算子的成就感弥补不了后面联调的时间成本。第三大模型上下文的公开热词里大模型微调相关的算子优化最容易忽视。微调场景和预训练场景的算子分布差异很大微调通常batch较小、激活值比例高有些算子在小batch下的kernel并不是最优实现。如果你只照着预训练的标准做算子选型微调性能可能差上一截。5.3 一点后续建议如果项目已经跑起来了但又想做算子这块的深度优化建议考虑构建一个轻量的算子基准测试集。把模型里高频使用的算子摘出来固定输入shape、精度和硬件环境做成一个benchmark脚本。后面每次换框架版本、换硬件、或者改算子实现都可以对着这个基准测一测。这个基准测试集的维护成本不高但能在版本升级的时候帮你挡住很多隐性性能回退。根据我个人的实操体会算子框架的选型从来不是一次性的决定。项目早期选什么框架可能只是影响开发效率到了优化部署阶段算子层面的技术深度才真正决定项目能做到什么程度。提前摸清楚算子层的事情比等项目出瓶颈了再临时补课省下来的时间是以周为单位的。