做算子框架选型这个标题听起来像是一个PPT上的汇报题目但实际上它是每个做AI Infra、做推理加速、做大模型部署的团队都绕不过去的硬仗。我最早接触算子开发还是在卷积网络刚火那阵子那时候写个自定义算子主要靠CUDA硬怼现在再看大模型时代的算子框架选型逻辑已经发生了明显变化——GPU算力变强了但算子的复杂度和融合需求也上来了特别是推理侧的显存带宽瓶颈、分布式并行策略、混合精度对齐这些问题让“选哪个框架”成了直接影响项目成败的决定性因素。这篇文章我只聊一件事在大模型训练和推理这条具体的技术路径上算子框架到底该怎么选。我先给出一个结论性判断大模型算子框架选型不存在绝对最优解只存在“在特定硬件、特定模型结构、特定团队能力下的性价比最优解”。暴论一点说绕开自己的实际场景去谈算子框架优劣基本等于耍流氓。很多团队在选型阶段就把精力放在了“哪个框架更流行”上结果上了生产环境才发现问题出在算子的自动调优程度、硬件适配深度、以及团队自己能不能维护这个代码库上。这篇文章我会结合自己的实操经历把主流的算子开发方案逐一拆开谈谈它们的适用边界、性能天花板和真实维护成本再给出一套可落地的选型决策方法。1. 算子框架选型先搞清楚要解决什么问题1.1 算子到底是什么这里不搞教科书式的定义我用大白话解释一下。模型推理或者训练过程本质上就是一大批张量数据的数学运算流水线每一个运算节点矩阵乘、向量范数、Softmax、LayerNorm、残差连接、激活函数就是一个算子。你可以把算子想象成工厂流水线上的一个工位不同工位干不同的活数据在这些工位之间流转最终得到预测结果或者梯度更新。在大模型时代算子的重要性被进一步放大。核心原因有两点第一Transformer架构的计算规律性极强本质上就是那几种算子在不同维度的张量上的反复执行算子写得好不好直接决定整体性能第二大模型的显存容量和带宽极容易成为瓶颈FlashAttention、PagedAttention这类论文之所以能产生那么大的影响本质上就是在算子级别做出了巨大的性能突破——把访存密集的操作和计算密集的操作进行融合减少显存读写次数。所以算子层面做好优化收益往往是数量级的而不是百分之几。1.2 选型前必须明确的三个问题我在大量项目复盘里发现算子框架选型翻车大部分不是因为某个框架不好而是因为需求就没想清楚。我建议动手选型之前先把下面三个问题拿白纸写下来第一个问题你的算子要跑在什么硬件上。这个问题决定了选型的大方向。如果是NVIDIA GPUCUDA生态和Triton是主力选项如果目标硬件包括国产加速卡、或者需要跨多种硬件平台部署那TVM这类编译器路线或者硬件厂商的专属工具链如华为昇腾的CANN就更值得重视如果只是CPU场景那么oneDNN或者简单的OpenMP手写也能解决问题。不同硬件的指令集、内存层级、并行模型差异巨大想要一套代码打天下目前依然是理想。第二个问题你是侧重训练侧还是推理侧。训练侧算子开发更多要考虑反向传播的实现、梯度计算的正确性、以及自动微分框架的配合推理侧则更关心算子能否融合、能否利用低精度推理FP16、INT8、INT4、能否减少显存占用和带宽开销。这两类场景的最优解经常不一样。第三个问题团队真实的工程师技能树是什么。如果团队里都是经验丰富的CUDA工程师直接走CUDA/CUTLASS路线完全可行如果团队以PyTorch使用为主、系统编程经验一般那么Triton几乎是代价最小的入门路径如果团队规模小、追求短期快速交付那就尽量少自研优先依赖社区成熟算子库只在关键路径上做定制优化。这三个问题全部想清楚以后你拿到任何技术对比资料都能快速过滤出真正的候选方案。这也是下面几节内容的讨论前提。2. 主流算子开发方案的底牌与适用边界2.1 CUDA C / CUTLASS性能天花板最高开发成本也最高如果追求极致的性能上限CUDA C和CUTLASS依然是绕不开的选项。CUTLASS是NVIDIA官方开源的CUDA C模板库里面实现了GEMM矩阵乘、Convolution、FlashAttention等一系列高效计算模板设计哲学是“把通用计算模式抽象成可组合的模板组件”。用CUDA手写算子的优势非常明显你拥有对寄存器、共享内存、指令级流水线的完全控制权能做极其精细的访存调度和计算重排。FlashAttention最初的CUDA版本就是手工pipeline优化做到极致的结果那种级别的性能优化用上层语言基本不可能达到。但代价也极其明显。首先是开发周期长一个中等复杂度的融合算子从设计到调优到稳定资深工程师可能也要数周时间其次是调试地狱显存越界、bank conflict、Warp divergence这些问题报错信息往往非常晦涩新手很容易被困住最后是维护成本高硬件迭代后代码可能要跟着调整除非你把CUTLASS的抽象复用得很好。一个典型的应用模式是先用Triton或者PyTorch原生方式完成功能验证然后对验证过的核心热点算子例如注意力、MoE的Gate分支用CUDA/CUTLASS做深度优化。这个模式我实践下来效率最高。2.2 Triton易用性和性能之间的最佳平衡Triton是OpenAI开源的一套GPU编程语言和编译器它让工程师可以用接近Python的语法编写高性能GPU算子同时编译器负责生成高效的CUDA代码。我自己的观点是对于大多数团队Triton已经足够用了——这句话放在三年前我是不信的但现在Triton的发展速度已经把差距缩小了很多。Triton的核心设计思想是把GPU并行编程的复杂度部分封装掉。你不需要手动管理线程块内部的共享内存分配、不需要手动做线程同步只需要用tl.arange、tl.load、tl.store这类高层原语描述数据流和块级计算逻辑编译器自动做向量化、并行化和部分流水线优化。举个例子写一个融合了Bias、Residual和GeLU的算子Triton表达起来非常直观import triton import triton.language as tl triton.jit def fused_bias_residual_gelu_kernel( x_ptr, bias_ptr, residual_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr ): pid tl.program_id(axis0) block_start pid * BLOCK_SIZE offsets block_start tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) bias tl.load(bias_ptr offsets, maskmask) residual tl.load(residual_ptr offsets, maskmask) x x bias residual # GeLU approximate version: x * sigmoid(1.702 * x) x x * tl.sigmoid(1.702 * x) tl.store(out_ptr offsets, x, maskmask)上面这段代码放到真实环境里运行前还需要考虑向量化宽度的问题。当BLOCK_SIZE明显小于GPU线程束包含的元素数量时性能会下降而BLOCK_SIZE过大又会增加寄存器压力导致 occupancy 降低。我一般从BLOCK_SIZE1024开始做自动化参数搜索实测下来大多数elementwise类算子在这个值附近都能取得不错的效果。Triton的主要限制是不适合实现过于复杂的、依赖线程间数据交换的算法。比如某些稀疏张量操作、某些需要复杂shuffle的并行归约用Triton写会非常别扭这时候回到CUDA反而更快。另外Triton对新硬件的适配往往滞后于CUDA如果你的目标硬件是特别新的架构要留意编译器的支持状态。2.3 TVM / 编译器路线跨平台与自动优化的潜力TVM的思路和Triton不同它更像一个端到端的深度学习编译器框架。你在TVM里描述算子的计算逻辑然后TVM负责把计算图转换成适应不同硬件后端的可执行代码。理论上你写一遍算子TVM可以把它编译到GPU、CPU、NPU、甚至FPGA上运行。TVM的核心卖点在于自动调优AutoTVM / Ansor。它会在目标硬件上自动搜索代码变换方案包括循环分块大小、向量化方式、线程绑定策略等寻找在该硬件上性能最优的版本。这个思路在硬件种类多、团队人力不足的场景下很好用因为你不需要针对每款硬件手动调优。但TVM的复杂度是出了名的高。学习曲线陡峭中间表示IR层次多调试报错信息不友好。而且自动调优虽然节省了人工调优的时间但搜索本身也需要消耗大量的计算资源。我在实际项目里看到很多团队在引入TVM后陷入了“调优搜索跑得比训练还久”的尴尬。所以我的判断是TVM更适合大型平台团队他们要构建一个覆盖多硬件后端的统一推理引擎对于只针对单一GPU平台做加速的中小团队来说TVM带来的额外复杂度通常不值得。2.4 硬件厂商专属框架硬件厂商对自家芯片最了解所以专属框架的性能通常是最容易拉满的。典型的例子是华为昇腾的CANN。CANN提供了一套完整的算子开发和优化工具链工程师可以用TBETensor Boost Engine或者基于DSL的接口去开发算子最终编译映射到昇腾的AI Core上执行。这类框架的特点好坏各半。好的一面是你不需要理解非常底层的硬件细节按照厂商提供的抽象去写就行性能下限有保障同时厂商会提供大量预置算子库很多需求不需要重新开发。不好的一面是生态绑定性极强代码很难平移到其他硬件平台上另外部分框架的文档质量参差不齐社区资料匮乏遇到问题的时候排障成本很高。如果你已经确定硬件选型是某个特定厂商的设备那这个方案是绕不开的如果你还在硬件选型阶段那么厂商算子框架的成熟度应该作为一个重要评估维度华为昇腾、寒武纪、平头哥等各有各的算子生态特色建议让研发团队实际跑几个经典模型再下结论。2.5 PyTorch原生算子注册路径很多人忽略了一个事实PyTorch本身提供了完整的自定义算子注册机制。从torch.autograd.Function到torch.library库从CPU算子到CUDA算子你可以在PyTorch框架体系内完成算子的定义、注册、自动微分接入和序列化。这个路径最大的优势是“无缝集成”。你用C扩展写的自定义算子注册完成后在Python侧调用时跟原生算子没有任何区别反向传播自动接入TorchScript/torch.compile也能识别。对于团队不想引入额外依赖、只想在PyTorch框架里快速迭代的场景这是非常务实的起点。需要特别注意的是PyTorch倾向于让用户用Triton或者torch.compile去实现高性能算子原生C扩展对于简单的融合操作可能性能提升有限。这个方案最合适的使用场景是需要快速实现一个PyTorch生态里不存在的算子语义而不是做极端的性能优化。2.6 LLVM算子自发现新的自动化趋势LLVM算子自发现是热词里比较新的方向。传统上算子的接入和分发需要大量手工绑定的代码算子定义写了还要注册表、接口映射、甚至多框架适配层都要额外写一遍。LLVM算子自发现通过编译期元编程或者运行期反射机制让算子可以自动注册到框架中。比如通过LLVM的pass机制扫描内置函数列表自动辨认出哪些函数是算子入口然后自动生成调用调度表。这个思路对于算子数量大规模增长的场景很有价值。大模型时代算子种类越来越多、组合爆炸式增长如果每个算子都要手工配一套接入逻辑工程效率太低。自发现机制可以减少样板代码让开发者专注于算子本身的实现。但实话实说LLVM算子自发现目前还不是一个非常成熟的工程化方案更偏向于研究探索和大型基础设施团队的自建方向。如果你是小团队现阶段不需要在这个方向上花太多时间关注进展即可。3. 选型维度的量化对比与决策模型3.1 按场景拆选型权重我倾向于用一套权重模型来做选型决策而不是凭感觉。项目的不同阶段各维度权重完全不同。比如在快速验证阶段“开发效率”权重可以给到30%以上在生产部署阶段“性能稳定性”和“生态成熟度”的权重就必须提升。我给一个典型的权重模板大家可以根据自己的场景调整维度快速验证权重生产部署权重开发效率30%10%性能上限10%25%易调试性15%15%生态/社区活跃度20%10%长期维护成本10%20%硬件适配度15%20%这个表的意思就是你不可能所有维度都拉满。比如CUDA/CUTLASS在性能和硬件适配上几乎满分但是开发效率很低Triton正好反过来开发效率极高性能稍逊但应对90%的算子需求已经足够TVM被卡在中间哪一项都不特别突出但跨平台通用性非常好。3.2 主流方案一张表对比下面这张表我尽量把关键信息压缩了方便做快速筛选方案学习曲线开发效率性能上限跨平台性维护成本典型适用团队CUDA/CUTLASS陡峭低极高仅NVIDIA高核心算法团队、大型平台Triton平缓高高NVIDIA为主中低大多数AI训练/推理团队TVM很陡峭中中高强高多硬件平台平台团队厂商专属框架(CANN等)取决于文档中高(特定硬件)极弱中锁定硬件的业务团队PyTorch原生扩展平缓高中跟随PyTorch低快速迭代、需求明确的小团队LLVM自发现路线陡峭低(初期搭建)中高取决于实现高基础设施研发团队基于这个表我可以给出一个经验性的选择建议如果你是完全从零开始的团队目标硬件是NVIDIA GPU需要尽快交付一个可用的训练或推理系统那就选择“PyTorch Triton”的组合路径如果已经有一批CUDA工程师而且应用对Latency极其敏感比如高并发在线服务那就值得投入CUTLASS如果公司战略是“一套代码多硬件部署”再考虑TVM或者厂商专属框架的架构问题。3.3 一个真实的决策案例分析我曾经参与过一个千亿参数模型推理加速的项目最初的方案倾向是全部算子都用CUDA C重写。结果评估完发现要优化的算子总数超过50个CUDA工程师只有3个人按照每个算子平均两周的开发调优周期光是算子开发就要大半年项目根本等不起。后来我们调整了策略先把所有算子的热点ranking跑出来找出前五的热点算子用CUDA手写极致优化其余算子用Triton实现或者直接复用社区库。最终的结果是这个混合方案只用了不到三分之一的人力就达到了接近全CUDA方案的性能。这就是我说的“性价比最优解”的典型例子。4. 实操自定义融合算子的从零实现与性能对比4.1 用PyTorch原生路径实现一个融合算子先从一个最简单的融合场景开始实现y gelu(x bias)。理论上PyTorch里你可以直接写F.gelu(x bias)但这会产生一次中间张量、两次kernel launch一次加法一次gelu。我们把它融合成一个kernel减少访存和launch开销。用PyTorch C扩展方式实现需要先把C代码写好然后通过load_inline动态编译加载#include torch/extension.h #include ATen/cuda/CUDAContext.h __global__ void fused_gelu_bias_kernel( const float* __restrict__ x, const float* __restrict__ bias, float* __restrict__ out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { float val x[idx] bias[idx]; // GELU approximate out[idx] val * 0.5f * (1.0f tanhf(0.7978845608f * (val 0.044715f * val * val * val))); } } torch::Tensor fused_gelu_bias(torch::Tensor x, torch::Tensor bias) { auto out torch::empty_like(x); int n x.numel(); int threads 256; int blocks (n threads - 1) / threads; fused_gelu_bias_kernelblocks, threads, 0, at::cuda::getCurrentCUDAStream()( x.data_ptrfloat(), bias.data_ptrfloat(), out.data_ptrfloat(), n); return out; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def(fused_gelu_bias, fused_gelu_bias, Fused GELU Bias); }Python侧加载from torch.utils.cpp_extension import load_inline cpp_src # 上面那段代码 load_inline( namefused_gelu_bias_ext, cpp_sources[cpp_src], functions[fused_gelu_bias], extra_cuda_cflags[-O3], verboseFalse )这种做法的好处是所有逻辑都还在PyTorch体系内训练部署都能无缝衔接。缺点是手写CUDA kernel的代码量不小而且这块代码在不同GPU型号上的表现差异很大需要单独调block size。4.2 用Triton重写同一个算子并做自动性能调优同样的算子Triton版本短小精悍很多。前面已经贴过类似代码这里我补一个优化循环找最优BLOCK_SIZE的脚本import triton import triton.language as tl import torch triton.jit def fused_gelu_bias_triton( x_ptr, bias_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr ): pid tl.program_id(axis0) block_start pid * BLOCK_SIZE offsets block_start tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) b tl.load(bias_ptr offsets, maskmask) val x b # GELU approximate out val * 0.5 * (1.0 tl.erf(val * 0.70710678)) tl.store(out_ptr offsets, out, maskmask) # 自动搜索BLOCK_SIZE x torch.randn(1_000_000, devicecuda) bias torch.randn(1_000_000, devicecuda) out torch.empty_like(x) best_time float(inf) best_block 0 for block_size in [128, 256, 512, 1024, 2048]: # 用CUDA Event计时 start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) start.record() for _ in range(1000): fused_gelu_bias_triton[(x.numel() block_size - 1) // block_size,](x, bias, out, x.numel(), BLOCK_SIZEblock_size) end.record() torch.cuda.synchronize() elapsed start.elapsed_time(end) / 1000 if elapsed best_time: best_time elapsed best_block block_size说实话BLOCK_SIZE在原生CUDA里更多是经验值但在Triton里配合编译器自动向量化实际差距会比手写代码时小。这个搜索过程我用下来平均能带来10~20%的性能提升。4.3 性能测试怎么搞才不算白测很多人的性能测试是“测了个寂寞”因为方法不对。以下是我坚持的几条原则第一要warmup。GPU kernel第一次启动有懒加载和缓存冷启动的开销直接进入计时循环会把这部分时间算进去。我一般先跑20~50次预热再进入正式计时。第二要多次重复取中位数或者最小值而不是平均值。平均值容易受到系统其他任务干扰最小值更接近真实计算能力。第三要和朴素baseline对比。比如你这个融合算子要和PyTorch原生的F.gelu(x bias)对比而不是只测自己这个kernel的绝对时间。没有baseline的性能数字没有参考价值。第四要关注显存带宽的利用率。对于访存密集算子计算公式大概是带宽利用率 数据量 / 耗时 / 理论带宽。如果利用率不到70%说明kernel写的还有较大优化空间。4.4 算子的精度对齐与调试技巧大模型场景下算子精度问题隐藏得很深尤其是融合了多个数学操作的算子浮点运算顺序不同可能导致结果出现微小偏差。我在实际中见过因为精度对不齐导致训练Loss曲线出现异常波动的案例。调试的核心方法是数值比对。把融合算子结果和逐步骤参考实现比如用PyTorch里一步步计算的版本在相同输入下比对计算最大绝对误差和最大相对误差并结合算子本身的计算性质判断误差是否在可接受范围。如果你的目标是训练场景建议再用真正的梯度对比验证一下反向传播的正确性。Triton调试相对轻松可以直接在kernel里通过tl.device_print打印中间值这在CUDA里要麻烦得多。CUDA调试常用方法包括在host侧打印、使用compute-sanitizer检查越界和竞争条件、以及在代码中插入assert来捕获异常值。5. 常见问题与避坑实录5.1 选型阶段的常见陷阱我见过太多团队在选型上反复横跳。有个很典型的现象团队一开始选了Triton跑了一两个算子觉得很顺手然后雄心勃勃地想用Triton重写整个模型的所有算子结果写到第二个复杂算子就卡住了——因为Triton并不适合表达所有类型的算子。于是又回头去补CUDA导致前面写的Triton代码变成了半吊子遗留物。对这种问题我的建议是在选型阶段就明确哪些算子用哪种方案而不是一开始就不加区分地全盘走一个路线。先把你模型里的算子分类简单elementwise类操作、访存密集类操作优先用Triton复杂并行归约、稀疏操作、需要精细访存控制的分配给CUDA/CUTLASS。分工明确后团队协作效率会提升很多。第二个常见陷阱是忽略框架对部署环境的依赖。比如某些推理框架只支持ONNX导出那么你的Triton自定义算子就要检查能否顺利导出到ONNX或者是否有对应的IR转换路径。如果自定义算子占比太高部署时可能会有兼容性问题。这个问题我在做服务化部署时就踩过后来被迫在推理引擎侧加了一层自定义算子兼容层才算解决。第三个陷阱是低估了“算子库生态复用”的价值。很多时候你根本不需要自己写算子社区里可能已经有足够的、经过充分测试和优化的实现。比如FlashAttention、FlashDecoding这些经典算子直接用现成库比自己造轮子明智得多。选型时一定要问一句这个场景是否有成熟的社区方案如果有先拿来用再说。5.2 实操中的常见故障速查我整理了一个高频问题速查表基本都是实操中大家反复遇到的问题现象可能原因处理建议kernel执行结果报错CUDA illegal memory access显存越界或指针无效用compute-sanitizer定位越界地址检查grid/block索引计算kernel运行速度极慢可能是内存访问不合并检查是否按连续线性索引访问数据避免随机索引导致访存发散Triton编译失败报register spillBLOCK_SIZE过大减少BLOCK_SIZE或者减少kernel内的变量数量精度和参考实现差得很大浮点运算顺序差异/近似函数版本不同逐段对比中间值确认每个子操作的实现是否一致训练时Loss不收敛但前向算子结果正常可能是反向传播实现错误用torch.autograd.gradcheck验证自定义算子的梯度在不同的GPU型号上性能差异巨大未按硬件特性调优做自动化参数搜索或针对目标硬件单独调整block size以下几点是我长期用过之后沉淀出来的小经验也可以说是血泪教训。如果你在用Triton尽量保持kernel的简单性让编译器更愿意自动优化如果你在用CUDA优先把内存访问对齐到128字节这是在绝大多数GPU上都是性能收益最高的改动不管是哪种框架务必在真实负载场景、真实数据shape下做性能测试因为小shape和大shape的最优配置经常不一样团队协作时建议维护一份算子的精度基线、性能基线和典型shape清单每次改动都能自动做回归对比这套东西搭建起来以后整个团队的迭代速度完全不一样。关于大模型算子框架选型我个人的体会是不要神化任何一个框架也不要把精力耗在所谓的框架之争上。真正推动大模型性能前进的是对算子计算本质的理解、对硬件特性的掌握以及在工程上做抽象和取舍的能力。选型只是起点选完之后你依然要为精度对齐、性能调优、部署稳定性付出一大段时间。好在现在工具链成熟度已经好了很多Triton这类框架把算子开发的门槛拉低了一大截让更多工程师能参与到高性能计算的优化工作中来。这个方向很值得深耕因为越是往大模型落地的深处走算子层的长期竞争力就越重要。 
长河君 · 2026/11/2 16:34:07 · 阅读 1.2k 大模型算子框架选型：从CUDA到Triton的性价比决策
做算子框架选型这个标题听起来像是一个PPT上的汇报题目但实际上它是每个做AI Infra、做推理加速、做大模型部署的团队都绕不过去的硬仗。我最早接触算子开发还是在卷积网络刚火那阵子那时候写个自定义算子主要靠CUDA硬怼现在再看大模型时代的算子框架选型逻辑已经发生了明显变化——GPU算力变强了但算子的复杂度和融合需求也上来了特别是推理侧的显存带宽瓶颈、分布式并行策略、混合精度对齐这些问题让“选哪个框架”成了直接影响项目成败的决定性因素。这篇文章我只聊一件事在大模型训练和推理这条具体的技术路径上算子框架到底该怎么选。我先给出一个结论性判断大模型算子框架选型不存在绝对最优解只存在“在特定硬件、特定模型结构、特定团队能力下的性价比最优解”。暴论一点说绕开自己的实际场景去谈算子框架优劣基本等于耍流氓。很多团队在选型阶段就把精力放在了“哪个框架更流行”上结果上了生产环境才发现问题出在算子的自动调优程度、硬件适配深度、以及团队自己能不能维护这个代码库上。这篇文章我会结合自己的实操经历把主流的算子开发方案逐一拆开谈谈它们的适用边界、性能天花板和真实维护成本再给出一套可落地的选型决策方法。1. 算子框架选型先搞清楚要解决什么问题1.1 算子到底是什么这里不搞教科书式的定义我用大白话解释一下。模型推理或者训练过程本质上就是一大批张量数据的数学运算流水线每一个运算节点矩阵乘、向量范数、Softmax、LayerNorm、残差连接、激活函数就是一个算子。你可以把算子想象成工厂流水线上的一个工位不同工位干不同的活数据在这些工位之间流转最终得到预测结果或者梯度更新。在大模型时代算子的重要性被进一步放大。核心原因有两点第一Transformer架构的计算规律性极强本质上就是那几种算子在不同维度的张量上的反复执行算子写得好不好直接决定整体性能第二大模型的显存容量和带宽极容易成为瓶颈FlashAttention、PagedAttention这类论文之所以能产生那么大的影响本质上就是在算子级别做出了巨大的性能突破——把访存密集的操作和计算密集的操作进行融合减少显存读写次数。所以算子层面做好优化收益往往是数量级的而不是百分之几。1.2 选型前必须明确的三个问题我在大量项目复盘里发现算子框架选型翻车大部分不是因为某个框架不好而是因为需求就没想清楚。我建议动手选型之前先把下面三个问题拿白纸写下来第一个问题你的算子要跑在什么硬件上。这个问题决定了选型的大方向。如果是NVIDIA GPUCUDA生态和Triton是主力选项如果目标硬件包括国产加速卡、或者需要跨多种硬件平台部署那TVM这类编译器路线或者硬件厂商的专属工具链如华为昇腾的CANN就更值得重视如果只是CPU场景那么oneDNN或者简单的OpenMP手写也能解决问题。不同硬件的指令集、内存层级、并行模型差异巨大想要一套代码打天下目前依然是理想。第二个问题你是侧重训练侧还是推理侧。训练侧算子开发更多要考虑反向传播的实现、梯度计算的正确性、以及自动微分框架的配合推理侧则更关心算子能否融合、能否利用低精度推理FP16、INT8、INT4、能否减少显存占用和带宽开销。这两类场景的最优解经常不一样。第三个问题团队真实的工程师技能树是什么。如果团队里都是经验丰富的CUDA工程师直接走CUDA/CUTLASS路线完全可行如果团队以PyTorch使用为主、系统编程经验一般那么Triton几乎是代价最小的入门路径如果团队规模小、追求短期快速交付那就尽量少自研优先依赖社区成熟算子库只在关键路径上做定制优化。这三个问题全部想清楚以后你拿到任何技术对比资料都能快速过滤出真正的候选方案。这也是下面几节内容的讨论前提。2. 主流算子开发方案的底牌与适用边界2.1 CUDA C / CUTLASS性能天花板最高开发成本也最高如果追求极致的性能上限CUDA C和CUTLASS依然是绕不开的选项。CUTLASS是NVIDIA官方开源的CUDA C模板库里面实现了GEMM矩阵乘、Convolution、FlashAttention等一系列高效计算模板设计哲学是“把通用计算模式抽象成可组合的模板组件”。用CUDA手写算子的优势非常明显你拥有对寄存器、共享内存、指令级流水线的完全控制权能做极其精细的访存调度和计算重排。FlashAttention最初的CUDA版本就是手工pipeline优化做到极致的结果那种级别的性能优化用上层语言基本不可能达到。但代价也极其明显。首先是开发周期长一个中等复杂度的融合算子从设计到调优到稳定资深工程师可能也要数周时间其次是调试地狱显存越界、bank conflict、Warp divergence这些问题报错信息往往非常晦涩新手很容易被困住最后是维护成本高硬件迭代后代码可能要跟着调整除非你把CUTLASS的抽象复用得很好。一个典型的应用模式是先用Triton或者PyTorch原生方式完成功能验证然后对验证过的核心热点算子例如注意力、MoE的Gate分支用CUDA/CUTLASS做深度优化。这个模式我实践下来效率最高。2.2 Triton易用性和性能之间的最佳平衡Triton是OpenAI开源的一套GPU编程语言和编译器它让工程师可以用接近Python的语法编写高性能GPU算子同时编译器负责生成高效的CUDA代码。我自己的观点是对于大多数团队Triton已经足够用了——这句话放在三年前我是不信的但现在Triton的发展速度已经把差距缩小了很多。Triton的核心设计思想是把GPU并行编程的复杂度部分封装掉。你不需要手动管理线程块内部的共享内存分配、不需要手动做线程同步只需要用tl.arange、tl.load、tl.store这类高层原语描述数据流和块级计算逻辑编译器自动做向量化、并行化和部分流水线优化。举个例子写一个融合了Bias、Residual和GeLU的算子Triton表达起来非常直观import triton import triton.language as tl triton.jit def fused_bias_residual_gelu_kernel( x_ptr, bias_ptr, residual_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr ): pid tl.program_id(axis0) block_start pid * BLOCK_SIZE offsets block_start tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) bias tl.load(bias_ptr offsets, maskmask) residual tl.load(residual_ptr offsets, maskmask) x x bias residual # GeLU approximate version: x * sigmoid(1.702 * x) x x * tl.sigmoid(1.702 * x) tl.store(out_ptr offsets, x, maskmask)上面这段代码放到真实环境里运行前还需要考虑向量化宽度的问题。当BLOCK_SIZE明显小于GPU线程束包含的元素数量时性能会下降而BLOCK_SIZE过大又会增加寄存器压力导致 occupancy 降低。我一般从BLOCK_SIZE1024开始做自动化参数搜索实测下来大多数elementwise类算子在这个值附近都能取得不错的效果。Triton的主要限制是不适合实现过于复杂的、依赖线程间数据交换的算法。比如某些稀疏张量操作、某些需要复杂shuffle的并行归约用Triton写会非常别扭这时候回到CUDA反而更快。另外Triton对新硬件的适配往往滞后于CUDA如果你的目标硬件是特别新的架构要留意编译器的支持状态。2.3 TVM / 编译器路线跨平台与自动优化的潜力TVM的思路和Triton不同它更像一个端到端的深度学习编译器框架。你在TVM里描述算子的计算逻辑然后TVM负责把计算图转换成适应不同硬件后端的可执行代码。理论上你写一遍算子TVM可以把它编译到GPU、CPU、NPU、甚至FPGA上运行。TVM的核心卖点在于自动调优AutoTVM / Ansor。它会在目标硬件上自动搜索代码变换方案包括循环分块大小、向量化方式、线程绑定策略等寻找在该硬件上性能最优的版本。这个思路在硬件种类多、团队人力不足的场景下很好用因为你不需要针对每款硬件手动调优。但TVM的复杂度是出了名的高。学习曲线陡峭中间表示IR层次多调试报错信息不友好。而且自动调优虽然节省了人工调优的时间但搜索本身也需要消耗大量的计算资源。我在实际项目里看到很多团队在引入TVM后陷入了“调优搜索跑得比训练还久”的尴尬。所以我的判断是TVM更适合大型平台团队他们要构建一个覆盖多硬件后端的统一推理引擎对于只针对单一GPU平台做加速的中小团队来说TVM带来的额外复杂度通常不值得。2.4 硬件厂商专属框架硬件厂商对自家芯片最了解所以专属框架的性能通常是最容易拉满的。典型的例子是华为昇腾的CANN。CANN提供了一套完整的算子开发和优化工具链工程师可以用TBETensor Boost Engine或者基于DSL的接口去开发算子最终编译映射到昇腾的AI Core上执行。这类框架的特点好坏各半。好的一面是你不需要理解非常底层的硬件细节按照厂商提供的抽象去写就行性能下限有保障同时厂商会提供大量预置算子库很多需求不需要重新开发。不好的一面是生态绑定性极强代码很难平移到其他硬件平台上另外部分框架的文档质量参差不齐社区资料匮乏遇到问题的时候排障成本很高。如果你已经确定硬件选型是某个特定厂商的设备那这个方案是绕不开的如果你还在硬件选型阶段那么厂商算子框架的成熟度应该作为一个重要评估维度华为昇腾、寒武纪、平头哥等各有各的算子生态特色建议让研发团队实际跑几个经典模型再下结论。2.5 PyTorch原生算子注册路径很多人忽略了一个事实PyTorch本身提供了完整的自定义算子注册机制。从torch.autograd.Function到torch.library库从CPU算子到CUDA算子你可以在PyTorch框架体系内完成算子的定义、注册、自动微分接入和序列化。这个路径最大的优势是“无缝集成”。你用C扩展写的自定义算子注册完成后在Python侧调用时跟原生算子没有任何区别反向传播自动接入TorchScript/torch.compile也能识别。对于团队不想引入额外依赖、只想在PyTorch框架里快速迭代的场景这是非常务实的起点。需要特别注意的是PyTorch倾向于让用户用Triton或者torch.compile去实现高性能算子原生C扩展对于简单的融合操作可能性能提升有限。这个方案最合适的使用场景是需要快速实现一个PyTorch生态里不存在的算子语义而不是做极端的性能优化。2.6 LLVM算子自发现新的自动化趋势LLVM算子自发现是热词里比较新的方向。传统上算子的接入和分发需要大量手工绑定的代码算子定义写了还要注册表、接口映射、甚至多框架适配层都要额外写一遍。LLVM算子自发现通过编译期元编程或者运行期反射机制让算子可以自动注册到框架中。比如通过LLVM的pass机制扫描内置函数列表自动辨认出哪些函数是算子入口然后自动生成调用调度表。这个思路对于算子数量大规模增长的场景很有价值。大模型时代算子种类越来越多、组合爆炸式增长如果每个算子都要手工配一套接入逻辑工程效率太低。自发现机制可以减少样板代码让开发者专注于算子本身的实现。但实话实说LLVM算子自发现目前还不是一个非常成熟的工程化方案更偏向于研究探索和大型基础设施团队的自建方向。如果你是小团队现阶段不需要在这个方向上花太多时间关注进展即可。3. 选型维度的量化对比与决策模型3.1 按场景拆选型权重我倾向于用一套权重模型来做选型决策而不是凭感觉。项目的不同阶段各维度权重完全不同。比如在快速验证阶段“开发效率”权重可以给到30%以上在生产部署阶段“性能稳定性”和“生态成熟度”的权重就必须提升。我给一个典型的权重模板大家可以根据自己的场景调整维度快速验证权重生产部署权重开发效率30%10%性能上限10%25%易调试性15%15%生态/社区活跃度20%10%长期维护成本10%20%硬件适配度15%20%这个表的意思就是你不可能所有维度都拉满。比如CUDA/CUTLASS在性能和硬件适配上几乎满分但是开发效率很低Triton正好反过来开发效率极高性能稍逊但应对90%的算子需求已经足够TVM被卡在中间哪一项都不特别突出但跨平台通用性非常好。3.2 主流方案一张表对比下面这张表我尽量把关键信息压缩了方便做快速筛选方案学习曲线开发效率性能上限跨平台性维护成本典型适用团队CUDA/CUTLASS陡峭低极高仅NVIDIA高核心算法团队、大型平台Triton平缓高高NVIDIA为主中低大多数AI训练/推理团队TVM很陡峭中中高强高多硬件平台平台团队厂商专属框架(CANN等)取决于文档中高(特定硬件)极弱中锁定硬件的业务团队PyTorch原生扩展平缓高中跟随PyTorch低快速迭代、需求明确的小团队LLVM自发现路线陡峭低(初期搭建)中高取决于实现高基础设施研发团队基于这个表我可以给出一个经验性的选择建议如果你是完全从零开始的团队目标硬件是NVIDIA GPU需要尽快交付一个可用的训练或推理系统那就选择“PyTorch Triton”的组合路径如果已经有一批CUDA工程师而且应用对Latency极其敏感比如高并发在线服务那就值得投入CUTLASS如果公司战略是“一套代码多硬件部署”再考虑TVM或者厂商专属框架的架构问题。3.3 一个真实的决策案例分析我曾经参与过一个千亿参数模型推理加速的项目最初的方案倾向是全部算子都用CUDA C重写。结果评估完发现要优化的算子总数超过50个CUDA工程师只有3个人按照每个算子平均两周的开发调优周期光是算子开发就要大半年项目根本等不起。后来我们调整了策略先把所有算子的热点ranking跑出来找出前五的热点算子用CUDA手写极致优化其余算子用Triton实现或者直接复用社区库。最终的结果是这个混合方案只用了不到三分之一的人力就达到了接近全CUDA方案的性能。这就是我说的“性价比最优解”的典型例子。4. 实操自定义融合算子的从零实现与性能对比4.1 用PyTorch原生路径实现一个融合算子先从一个最简单的融合场景开始实现y gelu(x bias)。理论上PyTorch里你可以直接写F.gelu(x bias)但这会产生一次中间张量、两次kernel launch一次加法一次gelu。我们把它融合成一个kernel减少访存和launch开销。用PyTorch C扩展方式实现需要先把C代码写好然后通过load_inline动态编译加载#include torch/extension.h #include ATen/cuda/CUDAContext.h __global__ void fused_gelu_bias_kernel( const float* __restrict__ x, const float* __restrict__ bias, float* __restrict__ out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { float val x[idx] bias[idx]; // GELU approximate out[idx] val * 0.5f * (1.0f tanhf(0.7978845608f * (val 0.044715f * val * val * val))); } } torch::Tensor fused_gelu_bias(torch::Tensor x, torch::Tensor bias) { auto out torch::empty_like(x); int n x.numel(); int threads 256; int blocks (n threads - 1) / threads; fused_gelu_bias_kernelblocks, threads, 0, at::cuda::getCurrentCUDAStream()( x.data_ptrfloat(), bias.data_ptrfloat(), out.data_ptrfloat(), n); return out; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def(fused_gelu_bias, fused_gelu_bias, Fused GELU Bias); }Python侧加载from torch.utils.cpp_extension import load_inline cpp_src # 上面那段代码 load_inline( namefused_gelu_bias_ext, cpp_sources[cpp_src], functions[fused_gelu_bias], extra_cuda_cflags[-O3], verboseFalse )这种做法的好处是所有逻辑都还在PyTorch体系内训练部署都能无缝衔接。缺点是手写CUDA kernel的代码量不小而且这块代码在不同GPU型号上的表现差异很大需要单独调block size。4.2 用Triton重写同一个算子并做自动性能调优同样的算子Triton版本短小精悍很多。前面已经贴过类似代码这里我补一个优化循环找最优BLOCK_SIZE的脚本import triton import triton.language as tl import torch triton.jit def fused_gelu_bias_triton( x_ptr, bias_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr ): pid tl.program_id(axis0) block_start pid * BLOCK_SIZE offsets block_start tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) b tl.load(bias_ptr offsets, maskmask) val x b # GELU approximate out val * 0.5 * (1.0 tl.erf(val * 0.70710678)) tl.store(out_ptr offsets, out, maskmask) # 自动搜索BLOCK_SIZE x torch.randn(1_000_000, devicecuda) bias torch.randn(1_000_000, devicecuda) out torch.empty_like(x) best_time float(inf) best_block 0 for block_size in [128, 256, 512, 1024, 2048]: # 用CUDA Event计时 start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) start.record() for _ in range(1000): fused_gelu_bias_triton[(x.numel() block_size - 1) // block_size,](x, bias, out, x.numel(), BLOCK_SIZEblock_size) end.record() torch.cuda.synchronize() elapsed start.elapsed_time(end) / 1000 if elapsed best_time: best_time elapsed best_block block_size说实话BLOCK_SIZE在原生CUDA里更多是经验值但在Triton里配合编译器自动向量化实际差距会比手写代码时小。这个搜索过程我用下来平均能带来10~20%的性能提升。4.3 性能测试怎么搞才不算白测很多人的性能测试是“测了个寂寞”因为方法不对。以下是我坚持的几条原则第一要warmup。GPU kernel第一次启动有懒加载和缓存冷启动的开销直接进入计时循环会把这部分时间算进去。我一般先跑20~50次预热再进入正式计时。第二要多次重复取中位数或者最小值而不是平均值。平均值容易受到系统其他任务干扰最小值更接近真实计算能力。第三要和朴素baseline对比。比如你这个融合算子要和PyTorch原生的F.gelu(x bias)对比而不是只测自己这个kernel的绝对时间。没有baseline的性能数字没有参考价值。第四要关注显存带宽的利用率。对于访存密集算子计算公式大概是带宽利用率 数据量 / 耗时 / 理论带宽。如果利用率不到70%说明kernel写的还有较大优化空间。4.4 算子的精度对齐与调试技巧大模型场景下算子精度问题隐藏得很深尤其是融合了多个数学操作的算子浮点运算顺序不同可能导致结果出现微小偏差。我在实际中见过因为精度对不齐导致训练Loss曲线出现异常波动的案例。调试的核心方法是数值比对。把融合算子结果和逐步骤参考实现比如用PyTorch里一步步计算的版本在相同输入下比对计算最大绝对误差和最大相对误差并结合算子本身的计算性质判断误差是否在可接受范围。如果你的目标是训练场景建议再用真正的梯度对比验证一下反向传播的正确性。Triton调试相对轻松可以直接在kernel里通过tl.device_print打印中间值这在CUDA里要麻烦得多。CUDA调试常用方法包括在host侧打印、使用compute-sanitizer检查越界和竞争条件、以及在代码中插入assert来捕获异常值。5. 常见问题与避坑实录5.1 选型阶段的常见陷阱我见过太多团队在选型上反复横跳。有个很典型的现象团队一开始选了Triton跑了一两个算子觉得很顺手然后雄心勃勃地想用Triton重写整个模型的所有算子结果写到第二个复杂算子就卡住了——因为Triton并不适合表达所有类型的算子。于是又回头去补CUDA导致前面写的Triton代码变成了半吊子遗留物。对这种问题我的建议是在选型阶段就明确哪些算子用哪种方案而不是一开始就不加区分地全盘走一个路线。先把你模型里的算子分类简单elementwise类操作、访存密集类操作优先用Triton复杂并行归约、稀疏操作、需要精细访存控制的分配给CUDA/CUTLASS。分工明确后团队协作效率会提升很多。第二个常见陷阱是忽略框架对部署环境的依赖。比如某些推理框架只支持ONNX导出那么你的Triton自定义算子就要检查能否顺利导出到ONNX或者是否有对应的IR转换路径。如果自定义算子占比太高部署时可能会有兼容性问题。这个问题我在做服务化部署时就踩过后来被迫在推理引擎侧加了一层自定义算子兼容层才算解决。第三个陷阱是低估了“算子库生态复用”的价值。很多时候你根本不需要自己写算子社区里可能已经有足够的、经过充分测试和优化的实现。比如FlashAttention、FlashDecoding这些经典算子直接用现成库比自己造轮子明智得多。选型时一定要问一句这个场景是否有成熟的社区方案如果有先拿来用再说。5.2 实操中的常见故障速查我整理了一个高频问题速查表基本都是实操中大家反复遇到的问题现象可能原因处理建议kernel执行结果报错CUDA illegal memory access显存越界或指针无效用compute-sanitizer定位越界地址检查grid/block索引计算kernel运行速度极慢可能是内存访问不合并检查是否按连续线性索引访问数据避免随机索引导致访存发散Triton编译失败报register spillBLOCK_SIZE过大减少BLOCK_SIZE或者减少kernel内的变量数量精度和参考实现差得很大浮点运算顺序差异/近似函数版本不同逐段对比中间值确认每个子操作的实现是否一致训练时Loss不收敛但前向算子结果正常可能是反向传播实现错误用torch.autograd.gradcheck验证自定义算子的梯度在不同的GPU型号上性能差异巨大未按硬件特性调优做自动化参数搜索或针对目标硬件单独调整block size以下几点是我长期用过之后沉淀出来的小经验也可以说是血泪教训。如果你在用Triton尽量保持kernel的简单性让编译器更愿意自动优化如果你在用CUDA优先把内存访问对齐到128字节这是在绝大多数GPU上都是性能收益最高的改动不管是哪种框架务必在真实负载场景、真实数据shape下做性能测试因为小shape和大shape的最优配置经常不一样团队协作时建议维护一份算子的精度基线、性能基线和典型shape清单每次改动都能自动做回归对比这套东西搭建起来以后整个团队的迭代速度完全不一样。关于大模型算子框架选型我个人的体会是不要神化任何一个框架也不要把精力耗在所谓的框架之争上。真正推动大模型性能前进的是对算子计算本质的理解、对硬件特性的掌握以及在工程上做抽象和取舍的能力。选型只是起点选完之后你依然要为精度对齐、性能调优、部署稳定性付出一大段时间。好在现在工具链成熟度已经好了很多Triton这类框架把算子开发的门槛拉低了一大截让更多工程师能参与到高性能计算的优化工作中来。这个方向很值得深耕因为越是往大模型落地的深处走算子层的长期竞争力就越重要。