做一次算子框架选型,等于把整个底层计算链条重新审视了一遍。
先说结论:如果你正在做大模型的训练或者推理优化,大概率绕不开算子和算子框架选型这件事。模型上层的架构都好说,真正决定你能不能在A100/H800上把训练吞吐跑满、在推理时把首token延迟压下来,往往就看那几个核心算子——矩阵乘、attention、LayerNorm、激活函数、RoPE位置编码,还有各种融合后的reduction算子。算子框架选型就是在这个层面做决策:用谁的算子、怎么开发缺的算子、怎么让算子在不同硬件上高效跑起来。
这篇文章我想用我自己实际做选型的经验,把这件事拆开来讲。从算子清单怎么盘、候选框架怎么对比,到打分权重怎么定、PoC怎么做、落地时踩过哪些坑,一条线捋下来。适合正在做大模型Infra、推理加速、训练框架二次开发,或者准备做国产卡迁移的团队参考。哪怕你只是个刚接触大模型底层的小白,看完也能知道选型这件事到底在选什么、每一步该怎么落。
1. 选型之前,先做算子盘点
很多人都喜欢先看框架再定方案,我恰恰建议反过来。算子框架选型,本质上是给你手里已经存在的算子需求找一组最合适的开发/集成方式。需求都没摸清,框架比一万次也是空转。
1.1 为什么算子清单是选型的地基
大模型的计算路径其实比大部分人想象的要窄。以GPT类Decoder-only模型为例,前向过程里高频出现的算子主要就那些:QKV投影和输出投影对应的GEMM,Attention里的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,后者编译成M=1的小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、融合attention |
| CUTLASS | NVIDIA底层模板库 | 极高 | 中低 | 仅NVIDIA | 手写高性能GEMM/卷积/attention |
| Triton | DSL + 编译器 | 高(接近CUTLASS的80-95%) | 高 | 可移植到部分硬件 | 快速开发自定义融合算子 |
| cuBLAS/cuDNN | 官方闭源库 | 极高 | 不用开发 | 仅NVIDIA | 标准GEMM、卷积、RNN等 |
| FlashAttention | 官方/社区算子库 | 极高 | 不用开发 | NVIDIA为主,生态适配中 | attention融合与长序列优化 |
| CANN + ACL | 昇腾底层框架 | 高(理论峰值内的优化空间大) | 中 | 昇腾NPU | 昇腾卡上的算子开发与优化 |
| MindSpore | 全流程AI框架 | 中高 | 中高 | 多硬件 | 全栈国产化,多后端统一 |
这张表有个关键信息要划重点:Triton的性能上限虽然略低于CUTLASS,但开发效率高出好几个量级。指纹一个简单融合kernel,CUDA可能要三四百行还要调shared memory,Triton四五十行就能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架构。比如你手头是A100(sm_80),装了一个给Hopper(sm_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原生的慢版本。因为优化版本出问题时,你需要一份绝对可信的对照来确认到底是不是优化代码引入的错误。我见过太多次,优化版本性能没达标、结果还对不上,最后一查,是基线版本本身在某种精度下就错了。没有基线,这个问题根本发现不了。
如果你现在正被一堆算子压得喘不过气,不妨从今天开始,先做那张算子盘点表。后面所有选型、优化、避坑的动作,都会因为这张表而变得清晰起来。