做一次算子框架选型等于把整个底层计算链条重新审视了一遍。先说结论如果你正在做大模型的训练或者推理优化大概率绕不开算子和算子框架选型这件事。模型上层的架构都好说真正决定你能不能在A100/H800上把训练吞吐跑满、在推理时把首token延迟压下来往往就看那几个核心算子——矩阵乘、attention、LayerNorm、激活函数、RoPE位置编码还有各种融合后的reduction算子。算子框架选型就是在这个层面做决策用谁的算子、怎么开发缺的算子、怎么让算子在不同硬件上高效跑起来。这篇文章我想用我自己实际做选型的经验把这件事拆开来讲。从算子清单怎么盘、候选框架怎么对比到打分权重怎么定、PoC怎么做、落地时踩过哪些坑一条线捋下来。适合正在做大模型Infra、推理加速、训练框架二次开发或者准备做国产卡迁移的团队参考。哪怕你只是个刚接触大模型底层的小白看完也能知道选型这件事到底在选什么、每一步该怎么落。1. 选型之前先做算子盘点很多人都喜欢先看框架再定方案我恰恰建议反过来。算子框架选型本质上是给你手里已经存在的算子需求找一组最合适的开发/集成方式。需求都没摸清框架比一万次也是空转。1.1 为什么算子清单是选型的地基大模型的计算路径其实比大部分人想象的要窄。以GPT类Decoder-only模型为例前向过程里高频出现的算子主要就那些QKV投影和输出投影对应的GEMMAttention里的QK^T、softmax、PV这种融合/半融合算子MLP里的两个大GEMM然后是LayerNorm/RMSNorm、SiLU/GELU这类逐元素算子再加上RoPE、causal mask、KV cache相关的内存操作。这张清单听起来不长但它决定了选型的天花板。你选CUTLASS核心诉求往往是那一两个GEMM或者attention融合kernel要打到硬件理论峰值的多少你选Triton多半是想快速写一些自定义融合算子又不想碰复杂到不行的CUDA模板代码你选CANN大概率是因为昇腾卡跑不了CUDA得用NPU体系去实现同一套东西。我的建议是在做任何框架选型之前先花一周时间把线上真实模型的算子profiling做一遍把算子列表、执行时间占比、调用次数、显存占用全拉出来。这一步不完成后面所有对比都容易陷入“框架介绍文章对产品文档”的虚空较量。1.2 按性能敏感度给算子分类算子不能一刀切。我习惯把所有算子分成三类性能极度敏感的热点算子比如大shape的GEMM、attention融合kernel常常占据整体时间的60%到80%。这类算子必须追求极致性能通常优先看官方库cuBLAS、cuDNN、FlashAttention、CANN的ACL算子里现成的或者用CUTLASS这类跟硬件贴得极近的框架。中等热度、需要融合的算子链比如LayerNorm 激活函数 residual add dropout这种组合单个算子不慢但串起来启动开销大、访存次数多。这类适合用Triton快速做融合或者用“图模式” 编译器自动融合。低频或单次开销小的算子比如初始化、一些debug用的逐元素运算。这类直接用PyTorch原生算子或者随便用个简单kernel实现就行完全不必为此引入复杂的框架选型。有了这个分类选型就很简单了热点算子的技术路线决定主框架中等热度算子决定辅助工具低频算子决定要不要在框架层面兜底。1.3 一张算子盘点表该长什么样一张能支撑选型决策的算子盘点表至少要有这么几列算子名称、所属模块embedding/attention/mlp等、输入shape范围、精度要求FP32/FP16/BF16/FP8、当前耗时占比、是否已融合、现有实现方式、性能瓶颈类型访存密集还是计算密集、未来是否有shape变化。输入shape范围这一列很容易被忽略但它特别关键。比如你在训练阶段是batch size 8、序列长度4096的大GEMM到了推理阶段变成batch size 1、序列长度逐token增长的decode GEMM这两种场景对算子的要求完全不一样。前者是典型的计算密集大GEMM后者编译成M1的小GEMM瓶颈在访存和kernel启动。你选的框架必须能同时覆盖这两个场景否则上线之后会发现训练跑得飞快推理卡得离谱。2. 大模型算子框架的四个层级和主流选择把算子框架拉通来看它们其实不在同一个抽象层上。有人把CUDA和Triton放在一起比就像拿着汇编跟Python比谁更快没有意义。我通常把算子相关的技术路线分成四个层级。2.1 算子到底跑在哪一层从指令到框架的完整链路第一层是硬件指令集比如NVIDIA GPU的SASS指令、昇腾NPU的AI Core指令。这一层几乎没有开发自由度通常由芯片厂商的编译器负责生成。第二层是底层编程语言/API最典型的是CUDA C以及NVIDIA官方的CUTLASS模板库。这一层的特点是性能上限最高可以对shared memory、寄存器、warp调度做极细致的控制但开发周期长调试困难。第三层是领域专用语言DSL和编译器代表就是OpenAI推出的Triton。它把shared memory分配、线程束调度都交给编译器开发者用类似Python的语法写kernel兼顾开发效率和不错的性能。再加上PyTorch 2.0的torch.compile和Triton后端这几年Triton几乎成了自定义算子的默认首选。第四层是算子库/上层框架比如cuBLAS、cuDNN、FlashAttention以及PyTorch自带的torch.nn.functional里的大量算子。严格说这不算“框架选型”而是“库选型”但在实际项目中能用现成算子就别自己写这个原则永远有效。2.2 主流框架横向对比我直接给一张我常用的对比表参数是“在我自己的测试环境里实测 社区反馈”综合得来的大家参考时可以按自己的硬件环境重新测技术路线抽象层级性能上限开发效率可移植性适合场景CUDA C底层编程语言极高低仅NVIDIA深度调优的GEMM、融合attentionCUTLASSNVIDIA底层模板库极高中低仅NVIDIA手写高性能GEMM/卷积/attentionTritonDSL 编译器高接近CUTLASS的80-95%高可移植到部分硬件快速开发自定义融合算子cuBLAS/cuDNN官方闭源库极高不用开发仅NVIDIA标准GEMM、卷积、RNN等FlashAttention官方/社区算子库极高不用开发NVIDIA为主生态适配中attention融合与长序列优化CANN ACL昇腾底层框架高理论峰值内的优化空间大中昇腾NPU昇腾卡上的算子开发与优化MindSpore全流程AI框架中高中高多硬件全栈国产化多后端统一这张表有个关键信息要划重点Triton的性能上限虽然略低于CUTLASS但开发效率高出好几个量级。指纹一个简单融合kernelCUDA可能要三四百行还要调shared memoryTriton四五十行就能the催一个不错的版本做好了再考虑用CUTLASS重写热点。实际线上运行里“性能差5个点但两周能上线”比“性能好5个点但两个季度才上线”要务实得多。2.3 到底选CUDA/CUTLASS还是Triton这是大家问得最多的一个问题。我的判断标准是这样的如果你要开发的算子本身结构非常稳定未来不会频繁改且推测它会长时间霸占profiling榜首——比如你自研了一个比FlashAttention更适合自家稀疏结构的attention — — 那CUTLASS值得投入因为每一点性能收益都会被每天上千万次调用放大。如果你要开发的算子结构可能会随模型版本迭代而变化或者属于“链式融合”里的一个环节对绝对性能要求没那么残酷那无脑选Triton就对了。Triton代码更容易review、更容易改、更容易被新人接手。再说一个很多人没提的点框架选型还看团队里能不能招到人。CUTLASS的资 deep 工程师市场上本来就少你选CUTLASS意味着团队里必须有至少一个人能读得懂CUDA C的模板元编程和warp-level编程选Triton基本上只要熟悉Python和一点并行编程概念的人就能快速上手。对一个六到十人的Infra团队来说技能树的匹配度有时候比理论性能更重要。2.4 国产硬件生态CANN与昇腾的现状如果你的项目涉及国产卡算子框架选型的重头戏就在CANN上了。CANN的ACL算子库目前已经覆盖了PyTorch上大多数常用算子但距离cuDNN/cuBLAS那种“开箱即用且高效”的程度还有差距。实际做昇腾移植时大概率会遇到三类情况部分算子有现成实现但性能不达标需要重新设计数据排布甚至用Ascend C昇腾的算子开发语言自己写部分融合场景比如attention的整个fused kernel在ACL里没有直接对应需要用ASCEND C或MindSpore的图模式去组合算子行为与CUDA不完全一致尤其是一些“在cuda上恰好能跑通的隐式依赖”到昇腾上就得显式处理。这就要命了如果你一开始选型的目标就包含多硬件适配那从模型代码层就应该尽量用倾向于“语义清晰、易重编译”的框架少用靠底层行为巧合才能跑通的黑魔法。哪怕是Triton昇腾的适配也在逐步推进中但目前为止还没有到“写一份Triton代码直接跑昇腾零修改”的程度。多硬件不是免费的选型时必须为“重复实现”预留人力。3. 选型评估的实操方法给候选框架打一个可比较的分数框架选型最忌讳“感觉哪个好就选哪个”。我自己的做法是把选型拆成四个步骤每一步都产出可量化的东西。3.1 第1步定义评估维度与权重评估维度必须结合自家业务。我给团队常用的一套权重是这样分配的评估维度默认权重说明目标算子的理论性能达成率30%用roofline模型估出上限实测/理论算子覆盖度与生态成熟度25%现成算子能覆盖盘点的多少比例开发与联调周期20%从写代码到上线需要多少人日可维护性与团队技能匹配度10%后续能否有人持续维护跨硬件/未来扩展性10%换卡、换模型结构时能否平滑社区活跃度与技术支持5%卡住了能否找到答案权重根据项目阶段可以变如果是快速交付一个PoC demo开发周期权重上调理论性能下调如果是给长期训练集群做底座理论性能和可维护性权重上调开发速度可以打折扣。3.2 第2步对覆盖率和生态做静态调研静态调研不要只看文档要动手验证两件事。第一件拿你的算子盘点表逐一去查候选框架的算子列表。比如你高频用到FlashAttention的某些变体就要查CUTLASS/cuDNN/FlashAttention仓库版本里到底有没有对应实现不是“有没有attention”而是“有没有你需要的固定shape下的multi-head attention且支持BF16和128的head_dim”。这种细节文档里经常写不清楚要直接翻源码和测试用例。第二件查社区里对这个框架的真实吐槽。去GitHub issues、技术社群、论文复现代码里看有没有人提到某个算子性能差、某个版本有bug、某个后端长期不更新。静态调研阶段看不出来的坑都会在落地阶段加倍奉还。3.3 第3步用真实模型做基准测试静态调研只能筛掉“明显不行”的真正的差距要用公司自己的模型跑基准测试。我给一个还算标准的三步走流程造一个最小复现脚本把模型里最热的三五个算子单独抽出来构造贴近真实训练的shape和精度写成一个benchmark脚本。注意shape一定要用真实数据不要用网上demo里的默认shape。跑候选框架的同一份算子比如分别用cuBLAS的GEMM、CUTLASS的GEMM、Triton仿写的GEMM跑同一个矩阵乘记录吞吐和显存。这里有个经验测试时要固定线程数、关闭CPU频率缩放最好把GPU时钟也锁住用nvidia-smi -lgc不然测出来的数据抖动严重没法横向对比。全模型端到端跑一遍单算子快不代表全模型快因为算子之间还有launch开销、数据搬运和依赖等待。把候选框架的算子接到完整模型里用训练场景的step time或者推理场景的token/s做最终验证。这一步经常得出跟单算子测试完全相反的结论我遇到过不止一次单算子快了10%但因为引入了额外的拷贝端到端反而慢了。3.4 第4步做一个72小时概念验证PoC最后无论静态调研和基准测试结果多漂亮我都建议再做一个72小时的PoC真实模拟你团队里“最普通的一名工程师”从零开始用这个框架落地一个算子的过程。我见过太多团队卡在这一步。选型文档写得天花乱坠真到开发的时候发现环境装不上、跟现有PyTorch版本冲突、算子在边界shape下直接崩、编译一次要半小时……这些都是benchmark里看不出来的。72小时PoC的产出不是一个“又可以跑的demo”而是一份“踩坑记录”哪些问题是文档能解决的哪些问题必须靠改框架源码解决哪些问题目前根本无法解决。这份记录比任何评分表都更有决策价值。4. 落地过程中常见的坑和排查方法选型结束只是开始。真正让人头秃的是落地时的各种隐藏问题。这里挑几个我反复踩过的坑给大家做个参考。4.1 性能测试没做对没绑核、没预热、没用流很多团队算子写出来了一测性能比cuBLAS慢了三倍第一反应是框架不行其实是测试方法不对。首先要绑核。GPU kernel的host端启动也是要CPU参与的CPU调度抖动会直接影响GPU任务的提交间隔。测试时用taskset限制CPU核能极大降低噪声。然后要预热。GPU kernel第一次执行时有context初始化、模块加载、显存分配等开销这些都不是稳定态性能。一定要先跑几轮warmup再开始计时。还有一点容易忽略如果你用纯同步的torch.cuda.synchronize()计时结果会把kernel排队等待的时间也算进去。很多时候瓶颈不在算子里而在上游没发完任务。正确的做法是用CUDA事件或者profile工具如Nsight Systems、torch profiler去测每个kernel的实际执行时间而不是测“从host发起调用到kernel结束”的墙钟时间。4.2 框架版本与硬件驱动不匹配这个坑极大。CUTLASS、Triton、FlashAttention这类项目迭代极快新版本经常要求特定CUDA版本和GPU架构。比如你手头是A100sm_80装了一个给Hoppersm_90优化过的flash-attn虽然能编译但它可能根本不会走你期望的fused kernel路径而是落到一个通用低效的fallback实现。表现出来就是“下载了最新的attention库性能比之前更烂”。排查方法没什么捷径就是装好以后先跑一下官方提供的profiling脚本看看kernel name到底是不是你期望的那个。不是就说明编译时没匹配到正确架构去找对应老版本或者用TORCH_CUDA_ARCH_LIST环境变量显式指定架构重新编译。4.3 融合算子收益为负的意外情况Triton做算子融合很方便但融合不是银弹。我遇到过一次很典型的情况把LayerNorm和后面的激活函数融合后单算子时间减少了但端到端反而更慢了。原因其实不复杂融合前的LayerNorm输出一个中间张量后面激活函数直接消费它两个算子之间虽然多了一次显存读写但这个中间张量本来就是后续某些逻辑需要的。融合后不落显存了看起来省了访存但如果后续逻辑需要单独访问这个中间张量反而要再算一遍或者再多做一次拷贝。所以融合之前一定要看计算图里这个中间张量是不是只被一个消费者使用。只有一个消费者才适合融合。两个以上消费者贸然融合就是给自己埋雷。4.4 团队技能与框架复杂度错配选型阶段再理性的团队到落地阶段也可能被“技术情怀”带偏。有人觉得CUTLASS高大上有人觉得Triton不上档次还有人非要基于cuBLAS自己封装一层。我的建议是选型方案必须匹配团队的真实构成。如果团队里没有人能够基于CUDA error message定位到warp-level代码问题那我建议尽量别碰CUTLASS如果团队里全是Python背景的算法工程师Triton就是比CUDA更合适的默认选项哪怕性能差一点。这听起来不“技术”但它可能就是你的团队真正缺的那块评估维度。算子框架选型本质上是在“理想性能”和“团队可执行的现实”之间做平衡。一次选型失误消耗的不只是开发时间还有整个团队的士气——你能想象一个写Python的工程师被逼着看一个月C模板报错吗画面太惨。5. 最后分享一点个人体会选型这件事我做了很多轮最大的体会是先有清单再有方案先有基准再谈性能先有人再选技术。你手里那份算子profiling表是会说话的。热点算子的类型、shape的分布、融合的机会点都已经把答案暗示给你了。剩下要做的不过是拿候选框架一个个去碰碰完用真实模型验证再看团队能不能接得住。另外一个小建议无论选了哪个框架第一次落地新算子时一定要保留一份“基线实现”——哪怕是PyTorch原生的慢版本。因为优化版本出问题时你需要一份绝对可信的对照来确认到底是不是优化代码引入的错误。我见过太多次优化版本性能没达标、结果还对不上最后一查是基线版本本身在某种精度下就错了。没有基线这个问题根本发现不了。如果你现在正被一堆算子压得喘不过气不妨从今天开始先做那张算子盘点表。后面所有选型、优化、避坑的动作都会因为这张表而变得清晰起来。