大模型时代做算子框架选型本质上是在给整个AI项目“选地基”。地基选错了后面模型训练跑不动、推理延迟压不下去、硬件资源吃不满每一步都会回头找你算账。我这些年从传统CV算法转型做大模型部署在算子层面积累了不少经验也踩过不少坑今天把选型思路、对比逻辑和实操经验一次性说清楚。先说明白一个常见误区很多人在聊算子框架选型时第一反应是纠结PyTorch还是MindSpore或者是TensorFlow还是PaddlePaddle。这其实只看到了最上层的东西。真正的算子框架选型至少要拆成两个层面来看一是底层的高性能算子库或算子编译器比如cuDNN、CANN、oneDNN、FlashAttention这一类二是承载这些算子、负责自动求导和计算图调度的AI训练框架比如PyTorch的ATen底层、MindSpore的MindExpress底层。两个层面互相影响但决策逻辑完全不同。这篇文章适合谁看一是要在GPU上做大模型训练和推理的算法工程师二是要针对国产硬件昇腾、寒武纪、海光做适配的部署工程师三是准备自建算子层能力的框架开发者。全文我会结合自己在NVIDIA和昇腾两块平台上的实际项目经历讲清楚选型背后的关键思考也会给出可直接复用的评估流程和排查清单。1. 算子框架选型的基本盘先看清四层结构1.1 大模型计算链路中的“算子层”到底在哪大模型跑起来的完整链路从高到低可以分成四层。最上面是模型层就是Transformer、MoE、Diffusion这些架构定义好的模块往下一层是图编译与调度层负责把模型的计算图切分、优化、安排到不同设备上执行再往下就是我们要聊的算子层它负责把每一个具体的计算操作矩阵乘、归一化、注意力、激活函数高效映射到硬件指令上最底层是硬件层也就是GPU、NPU或其它加速芯片的计算单元与显存体系。在很多人的认知里算子层不过是“一个个现成的函数调用”但实际上算子层的选型严重影响着整个大模型项目的天花板。举个例子在GPT级别的稠密大模型训练中矩阵乘GEMM算子通常会占据超过60%的算力消耗而到了MoE架构专家路由和分片操作会引入大量稀疏计算和集合通信算子计算占比结构又完全不一样。大模型的“大量使用算子对硬件性能的挑战”正是体现在这里算子种类多、调用频率高、形态变化大单靠手写优化已经跟不上模型迭代速度。所以选型第一条铁律是先看清楚你的模型结构决定的热点算子是什么再决定算子框架的侧重方向。纯稠密Transformer和MoE大模型、多模态模型、多卡训练对算子框架能力的要求完全不在一个量级上。1.2 从“三大件”看算子框架选型的边界条件算子框架选型绕不开三件事硬件平台限制、训练与推理的差异化需求、以及团队自身的软件栈能力。硬件平台是最大约束条件。如果是NVIDIA GPU你会面对的是CUDA生态体系里的cuDNN、cuBLAS、CUTLASS以及近年快速成熟起来的Triton如果是昇腾设备那就要看CANN异构计算架构NN和它配套的算子编译器如果是国产化的寒武纪、海光或天数智芯各家都有自己封闭或半封闭的算子工具链。现实情况是绝大多数开源大模型框架都是基于CUDA优化的换了硬件平台之后同样的模型性能可能出现数量级的下降。GPU和NPU之间的算子优化思维甚至是对立的GPU擅长用大量线程并行掩盖延迟而NPU通常更强调数据流驱动的静态调度。训练与推理的目标差异决定了算子框架的选型权重。训练阶段最看重算子的自动求导支持、混合精度策略、显存效率和分布式扩展能力推理阶段则更看重算子融合、内核自动调优、量化支持和低延迟。同一个项目在训练和推理阶段很可能会需要两套不同的算子方案我在下面第3章会详细展开。团队软件栈能力是最容易被低估的因素。算子层开发门槛高一个GPU上高性能的FlashAttention算子既要懂硬件指令集、共享内存管理又要懂矩阵分块和线程编排。如果团队没有专门的性能优化工程师就不宜选择需要从零手写算子的路线而是优先选择算子库丰富、社区资料多的成熟框架一旦选定了像Triton这种“半自动”方案还需要提前评估团队能否handle住Triton方言和底层硬件细节。2. 主流大模型算子框架横向对比与选型思路2.1 GPU生态三大家cuDNN/cuBLAS、CUTLASS与TritonGPU生态里现在主流的算子实现路径有三条很多项目其实是混合使用它们的。第一条是NVIDIA闭源的cuDNN和cuBLAS。cuDNN涵盖了卷积、循环神经网络、归一化、注意力等算子cuBLAS则主要覆盖矩阵乘和向量运算。它们的优点是“开箱即用”、经过大规模验证、性能稳定缺点是闭源黑盒想要做极致融合和魔改时候非常吃力。我在做推理优化时发现cuDNN的某些attention实现虽然能用但显存占用偏高比后来上FlashAttention的融合算子高了不少这就是闭源方案的一个典型副作用。第二条是CUTLASS它是NVIDIA开源的CUDA C模板库。CUTLASS把GEMM分解成“线程块级、线程束级、线程级、指令级”的层次化分块结构工程师可以像拼积木一样组装算子。它适合对性能有极致追求、愿意投入专职C工程师的团队社区里许多高性能算子比如FlashAttention的很多变体的原型都基于CUTLASS。代价是学习曲线陡峭模板参数极其复杂调试和编译时间都是按小时计的。第三条是Triton这是我在2024年以来最推荐的GPU算子开发方案。Triton的优势在于“以Python写算子”语法上很像NumPy但它会把Python代码编译成高效的GPU内核。用Triton写一个融合的attention算子可能只需要几十行代码性能可以逼近CUTLASS手写版本但开发效率高出几倍。实测下来我在GPU上做算子原型验证时Triton是性价比最高的选择一方面它自动处理了线程调度和共享内存分配的大部分脏活另一方面它保留了几乎全部的性能可调空间num_warps、num_stages、maxnreg这些参数都有实际意义。这三条路并不互斥。实际项目中我倾向于用cuDNN/cuBLAS覆盖基础算子用Triton开发融合算子只有极少数瓶颈算子才值得动用CUTLASS手写。选型时的判断标准就是这个算子是不是全局热点融合之后能不能带来实质性收益团队有没有能力长期维护C代码2.2 国产算力生态CANN算子优化与算子编译器选型如果要覆盖国产化场景昇腾CANN是绕不开的课题。很多人的第一印象是“CANN资料少、文档写得让人劝退”这个批评对前期阶段确实成立但从2023年之后的版本看CANN的工具链已经完善了不少尤其是在自定义算子开发、算子调优和迁移适配方面。CANN体系下的算子开发主要有三条路。一是使用它内置的高性能算子库比如Convolution、MatMul等标准算子可以直接调用性能已经过华为调优一般能拿到硬件理论算力的80%以上二是使用TBETensor Boost Engine算子开发框架通过Python或C定义算子的计算逻辑再由调度编译器生成NPU指令三是使用Ascend C编程接口这是更贴近硬件的开发方式适合做极致性能算子。在CANN算子优化上有几个典型的实操要点。首先NPU的算子融合思路和GPU相似但更依赖静态形状推导输入shape固定时CANN的编译优化效果会好很多所以我做昇腾部署时会尽量避免动态shape接口能padding就padding到固定形状。其次NPU上数据搬运的代价极高算子与算子之间如果频繁在GlobalMemory和片上Buffer之间搬运数据性能会被打到惨不忍睹。我实测过一个LayerNorm算子未融合前因为来回搬运数据耗时占了模型的18%用Ascend C把归一化和后面的激活函数融合进一个内核后直接降到3%左右。选型上要给团队一个明确预期昇腾场景的算子框架选型本质上是和CANN“交朋友”。如果项目是纯推理、模型结构标准就尽量基于内置算子库搭计算图如果涉及自定义结构或性能瓶颈才上Ascend C手写但一定要留足时间预算因为NPU的调试手段和工具链成熟度与CUDA生态还有差距。2.3 前缀型算子库与融合策略oneDNN、FlashAttention与权重算子的取舍除了硬件厂商绑定的算子库还有一类“前缀型”算子库值得关注典型代表是Intel的oneDNN原MKL-DNN。它面向CPU和XPU做了深度优化在大模型推理的CPU部署场景里非常有用。很多人在CPU上跑大模型自嘲“全靠内存大”实际上如果会调用oneDNN的算子融合接口同时用AMX/Tile指令加速矩阵运算CPU推理的性能可以比默认的PyTorch路径提升50%以上。如果项目有低成本CPU部署的需求oneDNN是必须纳入选型列表的。另一类融合策略的经典案例是FlashAttention系列。FlashAttention的核心思路是“在线softmax 按块遍历”把attention的QK^T计算分块进行避免把完整的注意力矩阵写回显存从而把显存占用从O(N²)降到O(N)并显著提升计算效率。在大模型场景里FlashAttention已经不是“优化技巧”而是长序列训练和推理的必需品。选型时要特别注意FlashAttention有多个版本不同版本的数值精度和硬件兼容性不一样部分老版本在FP16下存在精度风险选型时要结合模型验证一起做。热搜词里提到的halcon滤波核与权重算子放在大模型语境下可以这样理解滤波核本质上是一种规则权重矩阵算子就是在这个权重上执行的计算逻辑。在大模型里卷积算子和某些attention结构天然携带“权重局部性”算子库选型时要特别关注权重布局layout的兼容性——是NHWC还是NCHW、权重是否做了内存对齐或分块转置这些细节对性能影响极大。我在做多模态模型性能分析时就遇到过因为权重layout不匹配导致算子库频繁做隐式转置算力利用率掉了20%的案例。3. 训练与推理场景下的算子框架差异化选型3.1 训练场景算子选型的关键指标训练场景的算子选型核心指标有三个算力利用率、显存效率、分布式扩展性。算力利用率是第一个硬指标。在GPU上训练大模型时如果MFU模型算力利用率低于40%先不要急着调并行策略很可能是某个热点算子的实现有问题。这时候需要先做profile把耗时Top10的算子列出来逐个检查它们的kernel实现。我在做一次LLaMA训练优化时发现RotaryPositionEmbedding算子的实现效率极低原因是每步都在显存上重新计算cos/sin表并且频繁读写改成预计算并缓存到显存后单卡训练速度直接提升了12%。这类算子层面的小优化累积起来极其可观。显存效率方面训练场景特别要注意激活值重计算和算子融合的取舍。FlashAttention这类融合算子之所以对训练友好不只是因为它快更重要的是它省显存可以让batch size开得更大、训练更稳定。选型时如果某个候选算子能省显存但计算速度弱一点通常也值得考虑因为大batch训练往往比单算子快几毫秒更有价值。分布式扩展性考验的是算子在多卡训练时的行为。数据并行下梯度同步阶段会触发大量集合通信算子张量并行下每个算子都可能被切分成shard版本比如ColumnParallelLinear和RowParallelLinear这些并行算子的实现质量直接决定了扩展效率。选型时要看算子库对集合通信的融合程度——比如NVIDIA的NCCL是否和框架的autograd衔接顺畅或者昇腾的HCCS集合通信在AllReduce时是否能和计算算子做流水线重叠。实操中我有个习惯做选型对比时会专门做一次“多卡强扩展性测试”从1卡到8卡都跑一遍相同配置看算子的吞吐是否线性增长。很多算子库单卡性能漂亮但一上多卡就因通信优化不足出现严重瓶颈。3.2 推理场景算子选型的关键指标与量化支持推理侧的选型逻辑完全不同。这里最关心的是延迟、吞吐、显存占用和量化适配。推理延迟最大的敌人是算子调度开销和小算子放大效应。大模型生成时每个token都走一遍完整计算图如果模型里有大量小算子kernel launch的开销会被无限放大。这时候算子融合的重要性要远高于训练场景。实操上我会用TensorRT的GraphOptimization或者PyTorch的torch.compile把计算图里能融合的算子全部熔掉目标是把每次token生成的kernel启动数量压到原来的三分之一以内。量化和算子框架的连接紧密程度直接决定推理能不能用低精度跑起来。当前主流是FP8和INT8量化FP8相对简单INT8需要算子库提供完善的per-channel量化kernel和反量化逻辑。选型时注意量化感知训练QAT出来的模型与后训练量化PTQ出来的模型对算子库的要求差异很大。QAT模型通常需要算子库里自带伪量化节点支持否则部署时还得改图PTQ则更依赖calibration工具与算子的数值鲁棒性。我遇到过选型测试没做INT8推理验证上线后才发现某个LayerNorm算子的量化版本精度崩了的尴尬情况白白多耗了三天返工。推理选型中的另一个要点是动态shape支持。很多推理框架的算子库默认按静态shape优化遇到动态batch、变长输入时可能退回通用慢速实现。选型时要明确项目是否需要支持变长序列如果答案是“是”就优先选择支持动态shape且有对应优化内核的算子方案比如PyTorch的CUDA Graph落盘再回放的机制以及TensorRT对动态shape的profile机制。4. 算子框架选型的配套工具链与工程落地经验4.1 从原型验证到上线调优脚本化、自动化测试与性能剖析我见到很多团队在算子选型时直接拿着最终模型跑benchmark一旦性能不达标根本不知道瓶颈出在框架、算子库版本、硬件配置还是模型结构。正确做法是搭一套“性能基线脚本自动化测试剖析工具”的组合把选型变成可复现的实验过程。第一步是脚本化原型验证。用Triton或手写kernel快速实现候选算子的简化版本跑通函数层面的性能测试和正确性对齐。这一步不需要纠结最终性能核心目标是确认“数值正确、形态可行、性能上限有潜力”。第二步是接入自动化测试框架这里我推荐pytest。算子级测试和模型级测试要分开算子级测试关注数值精度、边界值、shape兼容性、dtype覆盖模型级测试关注端到端的loss收敛曲线或者推理精度指标。跑断言的时候要针对不同硬件和框架组合做环境隔离防止依赖污染导致测试结果不可信。第三步才是性能剖析用NVIDIA的Nsight Compute分析GPU kernel的吞吐、访存效率、占用率用Nsight Systems看端到端的CPU/GPU时间线昇腾平台上则用CANN自带的msprof工具做算子耗时分析和NPU利用率统计。我的经验是选型过程中至少要做三轮“数值正确性回归性能剖析”循环。第一轮在大致可用的算子上跑确定性能上限第二轮针对瓶颈算子做融合或参数调优第三轮做全模型级的验证。每一轮都要把结论固化到文档里不然过两周回来你根本记不住为什么选了A方案而放弃了B方案。4.2 常见选型踩坑场景与问题排查速查表我把自己和同行们踩过的坑整理成了一张速查表选型或者排查问题时可以直接对着查坑点现象根因排查与规避方法算子性能差异被环境掩盖同一模型在不同机器上性能差距极大往往是小算子被CUDA Graph或算子库版本差异影响固定CUDA/cuDNN版本与驱动版本统一基准脚本数值精度“刚好飘”训练loss收敛但推理有偶发错误FP16/FP8下部分算子数值稳定性不够增加混合精度下的数值校验换BF16或更高精度回退NPU亲和度误判模型从GPU迁移到NPU后性能掉一半算子库未选用CANN内置高性能算子走了自定义路径先用profier看算子耗时Top替换成内置算子动态shape拖垮性能推理延迟抖动严重算子库对动态shape没有优化内核退回通用实现统一padding到固定长度或开启框架的动态shape优化如CUDA Graph多卡扩展性不明2卡比1卡快4卡反而变慢集合通信算子与计算算子重叠不佳通信开销占比上升单独测试AllReduce带宽检查数据并行sharding配置量化算子精度崩INT8推理指标大幅下降部分算子对低精度量化不鲁棒逐算子量化精度评估敏感算子保留FP16算子库版本锁死新模型需要新版算子但旧版本无法升级项目过度依赖某个算子库的旧接口选型时预留“算子抽象层”隔离底层版本变化影响这张表背后有个通用方法不管踩了什么坑都要先做控制变量。一次只改一个变量——换算子库版本、切CPU/GPU后端、改shape、调精度跑完对比多试几轮总能定位到问题。别同时改一堆东西否则出了bug根本没法归因。4.3 算子框架与上层框架、向量数据库等组件的协同大模型项目里算子框架不是孤立的它要跟上层训练推理框架PyTorch、MindSpore、PaddlePaddle、部署框架TensorRT、ONNXRuntime、vLLM以及可能配套的向量数据库组件协同工作。和上层框架的协同重点看是否支持“算子级替换”。PyTorch生态里用户可以在不修改模型代码的情况下通过dispatch机制把特定算子路由到自定义实现这是FlashAttention能快速嵌入各类模型的关键能力MindSpore生态则通过算子的注册与选择策略AOE、Ascend的算子内置选择实现类似目标。选型时如果框架不允许你在算子层做“无侵入替换”后期性能优化会非常痛苦。和向量数据库的协同主要是RAG场景。热搜词里反复出现“向量数据库 选型”大模型应用里检索增强生成阶段的高频算子是向量距离计算内积、余弦、L2距离这些算子既可能在向量库里执行也可能在GPU上用自定义算子处理。如果你的应用是超大规模向量检索比如千万级以上我建议选型时把“向量检索算子在GPU上的加速能力”也纳入算子框架的考察范围比如某些框架支持把IVF或HNSW的暴力搜索部分放到GPU上并行计算收益比单纯调CPU版本大得多。4.4 团队能力与长期维护视角下的最终建议最后想聊聊选型决策中那些“工程师思维之外”的因素。一个很现实的问题是算子框架的长期维护成本通常要远高于前期的开发成本。选一个冷门闭源算子方案可能前两个月很顺利第三个月碰到文档没有覆盖的边界case就只能在论坛里大海捞针。反过来选一个开源活跃、社区文档充足的技术栈哪怕初期性能上不是最优长期来看大概率是更划算的。我的建议是分梯队做决策第一梯队是“官方整机方案”GPU上用cuDNN/cuBLAS加TritonNPU上用CANN内置算子加Ascend C这属于稳妥路线第二梯队是“深度自研”如果团队有专门的HPC或性能优化工程师可以基于CUTLASS或CANN的底层接口自研关键算子第三梯队才是“多平台探索”当业务确实需要在不同硬件间灵活迁移时考虑用Triton这类跨平台性较好的方案做统一算子层同时接受一定程度的平台适配折损。这里要特别强调做跨平台算子层时尽量不要手动搞一套抽象层去试图覆盖所有算子的语义边界C模板元编程和Python派发层写到最后往往变成维护噩梦优先用成熟框架的注册机制做降级和分发。另外也建议每年做一次“算子框架版本升级评审”因为大模型领域演进太快cuDNN、CANN、PyTorch的算子实现都在快速迭代半年前最优的方案可能在当下已经不是最佳。评审的核心是看三个东西一是新版本是否修复了影响你的bug或性能问题二是新版本是否有能简化代码的新算子三是升级的测试成本有多大。不要盲目追新但也不能一直守旧。我个人在实际项目中的感受是算子层选型项目做多了之后能明显意识到好的算子框架选型不是一次性的“点菜”而是一个持续跟进的动态过程。你要在项目里不断做实测、做对比、做归档一边熟悉硬件特性一边理解模型计算结构才能逐渐形成自己的判断框架。没有什么方案能一劳永逸但我上面这套思路至少能保证你在做选型时方向不偏、坑踩得少走完一轮之后知道下一次去哪里调优。最后再分享一个小技巧所有选型决策一定要留下“理由文档”哪怕只有几句话写清楚当时为什么选、放弃了什么、在什么条件下可能需要切换。三个月后你再回来看这份文档一定会感谢当时的自己。