把AI模型部署到推理服务器上后你盯着性能报告问的第一个问题往往是为什么这个算子这么慢从会用PyTorch搭模型到亲手写算子仿佛是隔着一条专业鸿沟——模型架构师和硬件协议栈之间的那块灰色地带大多数人一直没跨过去。这篇文章想做的就是把“AI模型到算子的翻译过程”掰开揉碎讲一遍带只有框架经验、没碰过底层开发的读者真正理解算子是什么、怎么写、怎么优化以及最值得选哪条技术路线。内容会覆盖执行全流程、CUDA/Ascend C/Triton三条开发路线的选型逻辑、一个完整小算子的手写过程、性能瓶颈的根因分析还有入门阶段最容易踩的坑。1. 从AI到算子的翻译过程一个forward算子在GPU上到底经历了什么1.1 AI框架的“层”在编译器眼里是一张算子树大多数用PyTorch的开发者写模型的时候是在和nn.Linear、nn.Conv2d、nn.ReLU这些高层模块打交道。这些模块在英语里叫Layer或者Module代表了你的模型语义。但在编译器内部比如torch.fx导出的Graph里这些Layer并不会原样保存而是会被拆成一张有向无环图图中的每个节点就是一个算子Operator。举个最简单的例子nn.Linear(128, 256)在数学上做的事情是y x * W^T b。你用torch.fx.symbolic_trace导一下会发现它其实展开成了t转置操作、matmul矩阵乘、add加法、可能是reshape等一连串节点。每个节点都是一个算子都有对应的输入输出张量。为什么要拆这么细因为框架底层的执行引擎只认算子。你写的模型再有美感进入执行阶段后都会被翻译成一个个原子的数学运算这些数学运算被串成执行图调度器再按依赖关系把它们逐个喂给硬件。理解这一点你才会明白为什么很多推理优化都是在“算子层面”做文章不是去改模型结构而是把多了去了的小算子合并成大算子减少执行图和调度的开销。这就是后面要聊的算子融合的雏形。1.2 算子 vs Kernel两者不是一回事很多人把算子Operator和Kernel混着用但严格说这是两个不同抽象层次的东西算子是框架层的概念回答“这个数学操作是什么”。它有输入输出的形状规则、有数学语义、有反向传播的梯度规则。比如“矩阵乘”是一个算子无论你跑在CPU还是GPU上它的数学定义都是一样的。Kernel是硬件层的实现回答“这个操作在具体芯片上怎么算”。它包含线程组织方式、数据指针移动方式、片上内存使用方式、循环展开层次等一大堆与硬件强相关的细节。同一台机器上一个算子往往有多套Kernel分别应对不同数据形状、不同数据类型、不同硬件架构。极端一点的例子一个add算子在CUDA上可能是某个vector_add_kernel在CPU上可能是某个avx512_add_kernel同样在GPU上处理连续内存的add和处理按通道广播的add也可能走两套不同的Kernel。所以当你准备写算子时第一件要想清楚的事是你写的是算子的数学调度逻辑还是硬件Kernel大多数情况下你真正要写的是后者而框架注册、shape推导、梯度定义这些“算子外壳”反而被很多人忽略。可是实际生产环境中这层外壳和Kernel一样重要少了一样你写的Kernel就跑不进框架里。1.3 在GPU上执行的全流程拆解假设你已经在PyTorch里写好了一个自定义算子的Kernel并且注册成功接下来你调用它一次完整流程大致是这样数据准备Python端的Tensor对象其实只是“句柄”底层数据可能还在内存里要通过cudaMemcpy之类的接口搬运到显存或者直接在显存上分配一块新空间。这一步很慢所以框架通常会复用显存尽量避免频繁分配释放。算子分派框架根据算子的名字、输入张量的dtype、shape、layout从算子库里查找到最匹配的那个Kernel。这个过程叫做Dispatch。复杂框架里还会有多级分派比如PyTorch会先按设备分派再按数据类型分派还可能有Autograd逻辑介入。Kernel Launch启动CPU端的驱动代码把Kernel函数的入口地址、调用参数输入输出指针、各维度大小打包通过驱动API提交到GPU。在CUDA里就是kernelgrid, block, sharedMem这种写法在昇腾的CANN里则类似调用或aclnn接口。这一步是异步的CPU把指令丢给GPU就立刻返回不会等GPU算完。GPU调度GPU上的硬件调度器接收Kernel后把线程块Block分发到一个个流式多处理器SM上。一个Kernel如果线程块很多SM不够剩下的区块就要排队。这是你能在Profiler里看到Kernel并行度差异的原因。线程执行每个SM内的线程按编程模型开始干活从显存加载数据到寄存器/片上缓存做计算把结果写回显存。如果一个线程块之间有数据依赖还会碰到同步原语__syncthreads()或者块间依赖的grid sync。异步同步与返回CPU端通过cudaDeviceSynchronize()或依赖后续操作触发同步确认Kernel完成后才能安全地把显存结果拷回内存。这个过程在视觉上看就像流水线车间CPU只管下订单GPU才是真正干活的工人。很多人第一次写Kernel觉得“反正异步了CPU先往下跑呗”结果在时序上踩坑——该同步的时候没同步读到了旧数据。2. 算子开发三条路线CUDA、Ascend C与Triton怎么选2.1 CUDA C性能上限最高但生态墙也高CUDA到今天依然是算子开发的“参照系”。几乎所有公开的Kernel优化案例、加速库设计思路、Profiler教程都默认你在CUDA语境下讨论。使用CUDA C写Kernel你能获得最精细的控制每个线程访问哪块显存、用不用共享内存、L2 cache命中怎么优化、异步拷贝指令用哪个版本全都可以自己定。代价也很明确学习曲线陡峭而且你必须非常熟悉硬件的线程模型、内存层次、指令调度。一个普通工程师从零开始到能写出面向生产的卷积Kernel通常要以年为单位积累。而且CUDA代码无法直接跨平台跑——你不能指望写一套CUDA代码跑到昇腾NPU或AMD GPU上。所以CUDA适合的典型场景是你在英伟达GPU上做研究、做业务性能敏感到必须手写Kernel或者你准备长期从事高性能计算工作必须先把CUDA这套基本功打扎实。2.2 Ascend C面向国产AI芯片的C算子开发语言Ascend C是华为昇腾AI处理器上、CANN软件栈里提供的一种算子开发语言。对很多开发者来说这是最近两年刚出现的名词。它本质上是C的一个扩展屏蔽了一部分底层细节你不需要直接去管每个线程的寄存器分配而是通过它封装好的“矢量计算”接口来表达并行逻辑。一个典型的Ascend C算子结构看起来大概是这样下面代码是帮助你建立感觉的示意不是整段可编译代码实际需要依赖CANN头文件和工程模板class KernelAdd { public: __aicore__ inline KernelAdd(GM_ADDR x, GM_ADDR y, GM_ADDR out, const int32_t length) { /* 初始化数据指针 */ } __aicore__ inline void Process() { // 循环分块处理数据核心思想是从全局内存拷贝到片上统一内存计算完再拷贝回去 for (int32_t offset 0; offset totalLength; offset blockLength) { CopyIn(offset); // Global Memory - Unified Buffer AddCompute(offset); // 在UB上做矢量加法 CopyOut(offset); // Unified Buffer - Global Memory } } };Ascend C最有价值的地方在于它把计算和访存两个环节做了显式区分开发者通过CopyIn/Compute/CopyOut的流水线结构来管理片上内存与全局内存的搬运。这其实是AI处理器的通用思维只是之前没人愿意用这么直白的方式教给你。为什么要关注Ascend C因为国内算力需求增大昇腾芯片在很多场景已经承担起训练和推理任务。会写CUDA的人很多但会写Ascend C的人不多。而且华为官方已经推出“Ascend C算子开发能力认证中级”把技能体系标准化了这是很好的学习路径。如果你要考虑信创平台、要做昇腾适配这条路线绕不开。2.3 Triton用Python写Kernel的折中方案Triton是OpenAI开源的一种GPU编程语言近几年在AI Infra圈子火得很。它让你用Python语法写Kernel然后编译器自动处理很多底层细节。一段最简单的Triton向量加法Kernel长这样import triton import triton.language as tl triton.jit def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): pid tl.program_id(axis0) offsets pid * BLOCK_SIZE tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) y tl.load(y_ptr offsets, maskmask) tl.store(out_ptr offsets, x y, maskmask)Triton的逻辑是你负责表达“在每个Block里对一块连续数据做什么”编译器负责决定块大小怎么切、用哪种向量化访存指令、怎么分配寄存器。对于矩阵乘这类操作Triton的tl.dot可以让编译器自动做tiling和shared memory优化性能接近手写CUDA。我刚接触Triton时觉得它是“玩具”后来发现FlashAttention在Triton里实现后性能居然不输原生CUDA版本才意识到它的编译器做得远比想象中好。Triton非常适合快速验证Kernel算法、搞研究原型、给现有算子做性能摸底。缺点是它目前主要还是面向英伟达GPU在昇腾等平台上的支持成熟度不如原生Ascend C。2.4 到底怎么选一张表说清选型逻辑维度CUDA CAscend CTriton学习曲线陡峭要理解线程模型和内存层次中等封装了部分底层但需理解流水线平缓接近Python性能上限最高几乎可做一切底层优化高面向昇腾硬件结构专门设计中高编译器承担大部分优化硬件生态英伟达GPU为主昇腾NPUAtlas系列等目前主要是英伟达社区在扩展适用人群专业高性能计算工程师国内算力平台开发者、信创适配者研究员、算法工程师、快速原型开发者调试工具Nsight Compute成熟CANN配套Profiler/MSVP在快速成长用CUDA工具链间接分析较受限生产可用度极高在昇腾平台生产级正在被越来越多推理框架采纳我的建议是如果你有明确的平台目标跟着平台走。跑英伟达GPU为主先学Triton入门、再深入到CUDA跑昇腾芯片直接从Ascend C入手配套官方认证体系学。如果你是学生或者想建立全局视野CUDA依然是必修课因为大量书籍、论文、思维方式都建立在这个语境上。3. 手写第一个kernel从向量加法建立完整的算子开发观3.1 为什么第一个算子选向量加法很多教材喜欢拿矩阵乘当入门案例我觉得这是劝退。矩阵乘涉及维度对齐、分块策略、共享内存复用、外层循环调度这些东西一上来就把你淹没在细节里。向量加法才是真正的“Hello World”每个输出元素只依赖两个输入元素的简单相加没有跨线程协作没有访存复用。但它一点都不寒酸。向量加法能让你完整看到Kernel里的三个基本动作计算线程索引决定“我这个线程负责哪个位置的数据”从显存加载数据计算结果写回显存。这三个动作是所有Kernel的骨架。你会写向量加法了后面改造elementwise类算子几乎是无脑平移就算写卷积、写FlashAttention骨架里也永远存在这三步。3.2 手写一个CUDA Kernel并解释每个参数下面这个Kernel用CUDA C实现一维向量加法// 每个线程处理一个元素注意这是个很“教科书”的写法 __global__ void vecAddKernel(float* A, float* B, float* C, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { // 边界检查 C[idx] A[idx] B[idx]; } }启动它的时候CPU端代码要算好线程组织int threads 256; int blocks (n threads - 1) / threads; vecAddKernelblocks, threads(d_A, d_B, d_C, n);这段代码的信息量其实非常大threadIdx.x是当前线程在线程块内的编号blockDim.x是线程块大小这里是256blockIdx.x是整个网格里第几个线程块。所以blockIdx.x * blockDim.x threadIdx.x就是在“所有线程排成一维长队”后你所在线程的全局编号。if (idx n)是边界检查因为blocks * threads可能比n大多出来的线程什么都不干直接退出。启动参数blocks, threads里blocks是网格大小threads是线程块大小。前者一般是几万到几百万量级后者受硬件限制通常取128到1024之间。然后你会发现一个性能隐患这个Kernel每个线程只处理一个元素线程总量巨大如果数据量很大每个线程做完一件事就结束调度和销毁线程的开销是凭空浪费。更实际的做法是用grid-stride loop让每个线程循环处理多个元素__global__ void vecAddKernelLoop(float* A, float* B, float* C, int n) { for (int idx blockIdx.x * blockDim.x threadIdx.x; idx n; idx blockDim.x * gridDim.x) { C[idx] A[idx] B[idx]; } }这样你可以固定block数量比如让blocks等于GPU上SM数量的整数倍每个线程用一个固定步长遍历完整个数组减少线程创建开销也让工作总量可控。3.3 从Kernel到框架算子还差哪几块你写完一个Kernel只是第一步。要让PyTorch调它或者让它在推理框架里跑起来至少还要补上这几块内容设备侧的内存管理显存分配、数据拷贝、内存释放。在CUDA里是cudaMalloc/cudaMemcpy在Ascend C里是Global MemoryGM上通过Context管理。报错时最常见的就是内存没对齐或者没释放干净。梯度Kernel反向传播需要另一个Kernel。对向量加法而言梯度就是原样回传所以你的grad_kernel也是类似的向量遍历。对复杂算子比如矩阵乘来说反向Kernel的数学表达和前向完全不同工作量往往是前向的好几倍。框架封装在PyTorch中你要自定义一个torch.autograd.Function定义forward和backward两个静态方法在forward里调用你的Kernel在backward里调用你的梯度Kernel。更生产级的方案还会用torch.library注册自定义符号让算子能参与torch.compile的计算图优化。Shape推导与Dispatch策略框架需要知道你这个算子的输出shape怎么算。如果你希望支持多种dtypeFP32、FP16、INT8还要写分派逻辑让算子知道在哪个硬件、哪类数据上选哪个Kernel。我见过很多入门者卡在“Kernel能跑但框架不认”。其实框架封装的代码量不比Kernel少。你现在写一个向量加法如果走PyTorch扩展的路线完整工程往往包含setup.py、CUDA源文件、C封装头文件、Python包装类加起来一大坨。这也是为什么Triton越来越受欢迎——它把框架注册和Kernel写在同一个Python文件里省掉大量工程负担。我在实际项目中用Triton封装向量加法时调用侧可以非常干净import torch import triton def add_torch(x, y): out torch.empty_like(x) n x.numel() BLOCK_SIZE 1024 grid ((n BLOCK_SIZE - 1) // BLOCK_SIZE,) add_kernel[grid](x, y, out, n_elementsn, BLOCK_SIZEBLOCK_SIZE) return out x torch.randn(10000, devicecuda) y torch.randn(10000, devicecuda) assert torch.allclose(add_torch(x, y), x y)这段代码跑通后你就掌握了一个可复用的“算子开发闭环”定义数学逻辑写Kernel做边界处理框架调用验证正确性。3.4 验证和调试的几个实用技巧新手写算子最常见的现象是“跑是跑通了但结果不对”。我常用的排查顺序先写CPU朴素参考用纯Python或PyTorch本身实现一遍标准计算作为golden output。测试时不要直接用assert x y用torch.allclose(x, y, atol1e-4, rtol1e-4)因为GPU浮点累加顺序不同可能导致极小误差。测极端shape单元素、刚好等于一个Block大小、非对齐尺寸。很多Kernel崩溃都是因为边界外的线程/Block没有mask掉。打印关键中间值别只盯着最终输出。如果算子是多阶段流水把每阶段的数据分别拷回CPU判断是加载错了还是算错了。4. 性能瓶颈的本质访存模式、算子融合与软硬件协同4.1 先分清计算密集 vs 访存密集很多人以为算子慢是因为芯片算力不够但实际上绝大部分简单算子卡在访存上。向量加法就是最典型的访存密集算子计算内容只是一个浮点加法但要从显存读两个数、写一个数你干了三趟搬运。打个比方团队要搬家搬家公司派来的是一群博士生但他们的时间全花在路上真正在办公室里动脑思考工作的时间不到十分之一。这种情况下再怎么提高博士生的智商也没用因为瓶颈是车辆运输效率——对应到GPU里就是内存带宽。反观矩阵乘这类计算密集算子每一个输出元素需要K次乘加计算量远大于访存量这时候芯片算力才会成为瓶颈。分析性能的第一步永远先算一下这个算子的“算术强度”arithmetic intensity总计算量除以总访存量。算出来很小你就应该去优化访存算出来很大才去考虑指令效率和计算流水。社区里经常讨论“大量使用算子对硬件性能的挑战”核心也在这里模型里塞满了几百个Elementwise小算子每个算子都来一次“读取-计算-写回”哪怕单个算子访存不多串起来的重复读写在墙钟时间上是很可观的。这也是为什么编译器要帮你做算子融合。4.2 优化方向一让访存更高效访存密集算子优化的第一原则是减少搬运次数提高每次搬运的数据量。具体手法包括向量化访存在CUDA里每个线程用float4类型一次读写16字节等价于用一条指令做4个浮点数的搬运。计算Kernel时要注意把总元素数除以4并且要求内存16字节对齐。__global__ void vecAddVec4(float4* A, float4* B, float4* C, int n4) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n4) { float4 a A[idx]; float4 b B[idx]; float4 c; c.x a.x b.x; c.y a.y b.y; c.z a.z b.z; c.w a.w b.w; C[idx] c; } }保证合并访问同一线程块里的相邻线程最好访问显存里相邻的地址。GPU的显存控制器会把相邻线程的请求合并成少数几条大的内存事务。如果你的线程索引设计成跨步访问每隔几百个字节读一个内存事务会碎成几十条带宽利用率直接崩。数据布局调整图像数据常见的NCHW和NHWC之争本质就是让访存模式贴合硬件的连续读取特性。对卷积来说NHWC往往更符合通道维度连续读取的需求。利用片上内存如果多个线程会重复读同一块数据应该先把它搬进共享内存或者Unified Buffer而不是每次都去全局内存。这对应Ascend C里的CopyIn到UB再计算。4.3 优化方向二算子融合把多次搬运变成一次融合算子的动机其实很简单你看一个计算链路如果它是MatMul - PReLU不融合时是这样执行的MatMul Kernel从显存读两个输入矩阵把结果矩阵完整写回显存PReLU Kernel再把这个结果矩阵从显存读出来作用以后写回显存。两次Kernel两次全量访存。而融合后的逻辑是MatMul Kernel在计算某个输出tile时顺手把PReLU应用到这个tile上然后只写回一次。中间结果根本没有离开Kernel内部的寄存器或片上内存。数据量越大融合节省的访存越明显。实际工程里MatMulPReLU这类融合在CUDA和Ascend C里都是非常经典的算子。昇腾社区里讨论的“Ascend C融合算子MatMulPReLU”正是这条思路的典型代表用Ascend C的MatMul接口先做矩阵乘紧接着用矢量接口做PReLU激活两者在一个Kernel内部协同完成。需要注意融合对精度有潜在影响。如果你的融合链路里有PReLU、SiLU这类带指数或分段逻辑的激活函数在FP16下某些中间值可能会溢出如果分开跑两个Kernel编译器有可能会在两个算子之间偷偷做一次精度提升。融合方案设计时要充分考虑数值范围必要时在Kernel内混精度计算。4.4 哪些算子值得融合、哪些不值得融合不是越多越好。我自己的经验是拿两个指标来判断访存节省占比融合能省掉多少次全量中间数据的读写如果中间结果远小于输入输出总量融合收益就小。实现复杂度融合Kernel的代码复杂度上升多少调试难度增加多少值得融合的典型场景连续Elementwise算子链比如Add - ReLU - Mul它们都是逐元素操作融合后没有任何语义损失访存从三次变一次。卷积/矩阵乘后接激活函数这是推理引擎里最常见的融合优化几乎无处不在。小算子合并成大算子某些模型里会出现大量极小的算子每个算子启动都有固定开销合并后能显著降低调度开销。不值得融合的场景跨设备通信类算子如果融合要把数据从多卡之间来回搬运那么在Kernel内部做跨卡通信代价极高不如保持原子操作。融合后寄存器溢出一个Kernel内部塞太多逻辑导致寄存器使用超过硬件限制编译器被迫把变量塞进局部内存local memory性能反而崩。后续算子几乎不访存已计算数据比如融合后马上要做一个稀疏操作而稀疏逻辑本身就是随机访存的你辛苦算好的中间结果根本复用不上那就别费劲。判断公式可以说成融合收益≈省掉的中间数据读写量×读写代价−Kernel复杂度上升导致的性能损失。前者越大、后者越小越值得融。5. 入门路线与避坑清单从能跑通到能干活5.1 建议的入门路径我自己带过几个从算法岗转到算子开发的同事整理出一条比较顺的路径大约4到8周能建立整体感觉第1周观察。不需要写任何Kernel。挑一个简单模型比如ResNet18用PyTorch的profiler跑一遍导出每个算子的耗时和访存量。你很快会找到耗时的top算子然后意识到“模型性能问题不是模型结构问题是算子实现问题”。第2周复制。用Triton写一个向量加法并成功注册给PyTorch调用。这一步主要是建立“Kernel不是神秘魔法”的感觉。第3周扰动。在Triton Kenerl里故意改BLOCK_SIZE、改循环方式、改mask条件观察性能变化和错误类型。这一步会让你直观理解线程组织和边界条件。第4周融合。尝试把Add ReLU合并成一个Kernel对比融合前后的耗时。你会发现融合效果显著由此体会“访存搬运决定上限”这句话的重量。第5周之后平台深耕。如果你是做昇腾的开始看Ascend C官方文档和认证课程了解CopyIn/CopyOut、流水线同步、MatMul接口如果你是CUDA路线去读CUDA C Programming Guide重点读内存层次和线程模型然后是Nsight Compute的优化案例。5.2 踩过的坑版本、维度、精度、边界这些坑我几乎全踩过写出来给你省时间CUDA版本和PyTorch预编译包不一致系统里装的是PyTorch默认CUDA 12.1但你自己编译扩展用的nvcc是11.8跑起来偶尔报奇怪的段错误。解决办法是强制让编译环境与PyTorch构建时的CUDA版本一致最好直接用官方Docker镜像。维度看成错拿到的是一个四维Tensor展平成一维后忘了按NCHW顺序还原算出来结果不对。这个错最坑的地方是shape校验都能过只有数值对不上。FP16精度FP16下加减法很容易溢出。对于大数值的加法在Kernel内部先把FP16转成FP32做累加最后再转回FP16输出这个习惯要养成。忘记maskTriton里如果数组长度不是BLOCK_SIZE整数倍你不加mask越界部分读取的是垃圾数据有时候不报错结果却是一个莫名其妙的数。这类问题排查时非常消耗时间建议从第一天就养成“凡是非对齐一律mask”的习惯。容器里看不到GPU用Docker跑深度学习时容器内看不到显卡多半是--gpus all没加或者NVIDIA Container Toolkit版本对不上。排查这类环境问题往往是入门阶段最耗时的部分。5.3 值得依赖的周边工具与认证工具链能帮你省掉大量瞎猜的时间。英伟达平台用Nsight Compute做Kernel性能分析它告诉你每个Kernel的访存吞吐、计算吞吐、占用率、是否存在内存合并失败昇腾平台对应的是CANN自带的Profiler工具可以看AI Core的利用率、数据和指令流水重叠情况。认证方面**华为昇腾社区推出的“Ascend C算子开发能力认证中级”**我认真看过课程大纲覆盖了算子的工程结构、矢量编程、流水线同步、性能分析这些核心内容。如果你本来就要做昇腾适配那直接跟着认证体系学比到处找散装资料效率高得多。如果你对编译器方向感兴趣可以了解一下LLVM在算子自发现和自动调优上的进展——很多新一代算子工具链都在做“声明数学语义编译器自动找Kernel实现”的事情这可能是未来算子开发范式的一个重要变化。5.4 算子概念远不止神经网络最后补一个视野层面的建议。你在热词里能看到“Sobel算子”“Canny算子”“拉普拉斯算子”这些词它们在图像处理里早已存在表达的也是“对数据做特定变换”的数学概念。而“神经算子”则是科学计算里这两年非常热的路线——用神经网络直接学习算子映射。所以说算子这个抽象并不只属于AI框架它是整个计算世界的通用语言。当你从AI框架的视角进入算子开发后再去接触图像算子、科学计算算子会发现底层的思维模式完全一致定义数学操作、设计数据流动、优化硬件映射。这也是为什么我很推荐AI工程师花时间跨出模型代码走到底层去看一看。你会获得一种“软件和硬件之间我能说了算”的掌控感。我自己在实际操作中的体会是算子开发入门最难的从来不是语法而是思维模式的切换从“我要实现什么功能”切换到“数据是怎样在硬件里流动的”。先跑通再看搬运再谈融合这条路基本不会走偏。
AI算子开发从零到性能优化：CUDA、Ascend C与Triton路线全解析
把AI模型部署到推理服务器上后你盯着性能报告问的第一个问题往往是为什么这个算子这么慢从会用PyTorch搭模型到亲手写算子仿佛是隔着一条专业鸿沟——模型架构师和硬件协议栈之间的那块灰色地带大多数人一直没跨过去。这篇文章想做的就是把“AI模型到算子的翻译过程”掰开揉碎讲一遍带只有框架经验、没碰过底层开发的读者真正理解算子是什么、怎么写、怎么优化以及最值得选哪条技术路线。内容会覆盖执行全流程、CUDA/Ascend C/Triton三条开发路线的选型逻辑、一个完整小算子的手写过程、性能瓶颈的根因分析还有入门阶段最容易踩的坑。1. 从AI到算子的翻译过程一个forward算子在GPU上到底经历了什么1.1 AI框架的“层”在编译器眼里是一张算子树大多数用PyTorch的开发者写模型的时候是在和nn.Linear、nn.Conv2d、nn.ReLU这些高层模块打交道。这些模块在英语里叫Layer或者Module代表了你的模型语义。但在编译器内部比如torch.fx导出的Graph里这些Layer并不会原样保存而是会被拆成一张有向无环图图中的每个节点就是一个算子Operator。举个最简单的例子nn.Linear(128, 256)在数学上做的事情是y x * W^T b。你用torch.fx.symbolic_trace导一下会发现它其实展开成了t转置操作、matmul矩阵乘、add加法、可能是reshape等一连串节点。每个节点都是一个算子都有对应的输入输出张量。为什么要拆这么细因为框架底层的执行引擎只认算子。你写的模型再有美感进入执行阶段后都会被翻译成一个个原子的数学运算这些数学运算被串成执行图调度器再按依赖关系把它们逐个喂给硬件。理解这一点你才会明白为什么很多推理优化都是在“算子层面”做文章不是去改模型结构而是把多了去了的小算子合并成大算子减少执行图和调度的开销。这就是后面要聊的算子融合的雏形。1.2 算子 vs Kernel两者不是一回事很多人把算子Operator和Kernel混着用但严格说这是两个不同抽象层次的东西算子是框架层的概念回答“这个数学操作是什么”。它有输入输出的形状规则、有数学语义、有反向传播的梯度规则。比如“矩阵乘”是一个算子无论你跑在CPU还是GPU上它的数学定义都是一样的。Kernel是硬件层的实现回答“这个操作在具体芯片上怎么算”。它包含线程组织方式、数据指针移动方式、片上内存使用方式、循环展开层次等一大堆与硬件强相关的细节。同一台机器上一个算子往往有多套Kernel分别应对不同数据形状、不同数据类型、不同硬件架构。极端一点的例子一个add算子在CUDA上可能是某个vector_add_kernel在CPU上可能是某个avx512_add_kernel同样在GPU上处理连续内存的add和处理按通道广播的add也可能走两套不同的Kernel。所以当你准备写算子时第一件要想清楚的事是你写的是算子的数学调度逻辑还是硬件Kernel大多数情况下你真正要写的是后者而框架注册、shape推导、梯度定义这些“算子外壳”反而被很多人忽略。可是实际生产环境中这层外壳和Kernel一样重要少了一样你写的Kernel就跑不进框架里。1.3 在GPU上执行的全流程拆解假设你已经在PyTorch里写好了一个自定义算子的Kernel并且注册成功接下来你调用它一次完整流程大致是这样数据准备Python端的Tensor对象其实只是“句柄”底层数据可能还在内存里要通过cudaMemcpy之类的接口搬运到显存或者直接在显存上分配一块新空间。这一步很慢所以框架通常会复用显存尽量避免频繁分配释放。算子分派框架根据算子的名字、输入张量的dtype、shape、layout从算子库里查找到最匹配的那个Kernel。这个过程叫做Dispatch。复杂框架里还会有多级分派比如PyTorch会先按设备分派再按数据类型分派还可能有Autograd逻辑介入。Kernel Launch启动CPU端的驱动代码把Kernel函数的入口地址、调用参数输入输出指针、各维度大小打包通过驱动API提交到GPU。在CUDA里就是kernelgrid, block, sharedMem这种写法在昇腾的CANN里则类似调用或aclnn接口。这一步是异步的CPU把指令丢给GPU就立刻返回不会等GPU算完。GPU调度GPU上的硬件调度器接收Kernel后把线程块Block分发到一个个流式多处理器SM上。一个Kernel如果线程块很多SM不够剩下的区块就要排队。这是你能在Profiler里看到Kernel并行度差异的原因。线程执行每个SM内的线程按编程模型开始干活从显存加载数据到寄存器/片上缓存做计算把结果写回显存。如果一个线程块之间有数据依赖还会碰到同步原语__syncthreads()或者块间依赖的grid sync。异步同步与返回CPU端通过cudaDeviceSynchronize()或依赖后续操作触发同步确认Kernel完成后才能安全地把显存结果拷回内存。这个过程在视觉上看就像流水线车间CPU只管下订单GPU才是真正干活的工人。很多人第一次写Kernel觉得“反正异步了CPU先往下跑呗”结果在时序上踩坑——该同步的时候没同步读到了旧数据。2. 算子开发三条路线CUDA、Ascend C与Triton怎么选2.1 CUDA C性能上限最高但生态墙也高CUDA到今天依然是算子开发的“参照系”。几乎所有公开的Kernel优化案例、加速库设计思路、Profiler教程都默认你在CUDA语境下讨论。使用CUDA C写Kernel你能获得最精细的控制每个线程访问哪块显存、用不用共享内存、L2 cache命中怎么优化、异步拷贝指令用哪个版本全都可以自己定。代价也很明确学习曲线陡峭而且你必须非常熟悉硬件的线程模型、内存层次、指令调度。一个普通工程师从零开始到能写出面向生产的卷积Kernel通常要以年为单位积累。而且CUDA代码无法直接跨平台跑——你不能指望写一套CUDA代码跑到昇腾NPU或AMD GPU上。所以CUDA适合的典型场景是你在英伟达GPU上做研究、做业务性能敏感到必须手写Kernel或者你准备长期从事高性能计算工作必须先把CUDA这套基本功打扎实。2.2 Ascend C面向国产AI芯片的C算子开发语言Ascend C是华为昇腾AI处理器上、CANN软件栈里提供的一种算子开发语言。对很多开发者来说这是最近两年刚出现的名词。它本质上是C的一个扩展屏蔽了一部分底层细节你不需要直接去管每个线程的寄存器分配而是通过它封装好的“矢量计算”接口来表达并行逻辑。一个典型的Ascend C算子结构看起来大概是这样下面代码是帮助你建立感觉的示意不是整段可编译代码实际需要依赖CANN头文件和工程模板class KernelAdd { public: __aicore__ inline KernelAdd(GM_ADDR x, GM_ADDR y, GM_ADDR out, const int32_t length) { /* 初始化数据指针 */ } __aicore__ inline void Process() { // 循环分块处理数据核心思想是从全局内存拷贝到片上统一内存计算完再拷贝回去 for (int32_t offset 0; offset totalLength; offset blockLength) { CopyIn(offset); // Global Memory - Unified Buffer AddCompute(offset); // 在UB上做矢量加法 CopyOut(offset); // Unified Buffer - Global Memory } } };Ascend C最有价值的地方在于它把计算和访存两个环节做了显式区分开发者通过CopyIn/Compute/CopyOut的流水线结构来管理片上内存与全局内存的搬运。这其实是AI处理器的通用思维只是之前没人愿意用这么直白的方式教给你。为什么要关注Ascend C因为国内算力需求增大昇腾芯片在很多场景已经承担起训练和推理任务。会写CUDA的人很多但会写Ascend C的人不多。而且华为官方已经推出“Ascend C算子开发能力认证中级”把技能体系标准化了这是很好的学习路径。如果你要考虑信创平台、要做昇腾适配这条路线绕不开。2.3 Triton用Python写Kernel的折中方案Triton是OpenAI开源的一种GPU编程语言近几年在AI Infra圈子火得很。它让你用Python语法写Kernel然后编译器自动处理很多底层细节。一段最简单的Triton向量加法Kernel长这样import triton import triton.language as tl triton.jit def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): pid tl.program_id(axis0) offsets pid * BLOCK_SIZE tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) y tl.load(y_ptr offsets, maskmask) tl.store(out_ptr offsets, x y, maskmask)Triton的逻辑是你负责表达“在每个Block里对一块连续数据做什么”编译器负责决定块大小怎么切、用哪种向量化访存指令、怎么分配寄存器。对于矩阵乘这类操作Triton的tl.dot可以让编译器自动做tiling和shared memory优化性能接近手写CUDA。我刚接触Triton时觉得它是“玩具”后来发现FlashAttention在Triton里实现后性能居然不输原生CUDA版本才意识到它的编译器做得远比想象中好。Triton非常适合快速验证Kernel算法、搞研究原型、给现有算子做性能摸底。缺点是它目前主要还是面向英伟达GPU在昇腾等平台上的支持成熟度不如原生Ascend C。2.4 到底怎么选一张表说清选型逻辑维度CUDA CAscend CTriton学习曲线陡峭要理解线程模型和内存层次中等封装了部分底层但需理解流水线平缓接近Python性能上限最高几乎可做一切底层优化高面向昇腾硬件结构专门设计中高编译器承担大部分优化硬件生态英伟达GPU为主昇腾NPUAtlas系列等目前主要是英伟达社区在扩展适用人群专业高性能计算工程师国内算力平台开发者、信创适配者研究员、算法工程师、快速原型开发者调试工具Nsight Compute成熟CANN配套Profiler/MSVP在快速成长用CUDA工具链间接分析较受限生产可用度极高在昇腾平台生产级正在被越来越多推理框架采纳我的建议是如果你有明确的平台目标跟着平台走。跑英伟达GPU为主先学Triton入门、再深入到CUDA跑昇腾芯片直接从Ascend C入手配套官方认证体系学。如果你是学生或者想建立全局视野CUDA依然是必修课因为大量书籍、论文、思维方式都建立在这个语境上。3. 手写第一个kernel从向量加法建立完整的算子开发观3.1 为什么第一个算子选向量加法很多教材喜欢拿矩阵乘当入门案例我觉得这是劝退。矩阵乘涉及维度对齐、分块策略、共享内存复用、外层循环调度这些东西一上来就把你淹没在细节里。向量加法才是真正的“Hello World”每个输出元素只依赖两个输入元素的简单相加没有跨线程协作没有访存复用。但它一点都不寒酸。向量加法能让你完整看到Kernel里的三个基本动作计算线程索引决定“我这个线程负责哪个位置的数据”从显存加载数据计算结果写回显存。这三个动作是所有Kernel的骨架。你会写向量加法了后面改造elementwise类算子几乎是无脑平移就算写卷积、写FlashAttention骨架里也永远存在这三步。3.2 手写一个CUDA Kernel并解释每个参数下面这个Kernel用CUDA C实现一维向量加法// 每个线程处理一个元素注意这是个很“教科书”的写法 __global__ void vecAddKernel(float* A, float* B, float* C, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { // 边界检查 C[idx] A[idx] B[idx]; } }启动它的时候CPU端代码要算好线程组织int threads 256; int blocks (n threads - 1) / threads; vecAddKernelblocks, threads(d_A, d_B, d_C, n);这段代码的信息量其实非常大threadIdx.x是当前线程在线程块内的编号blockDim.x是线程块大小这里是256blockIdx.x是整个网格里第几个线程块。所以blockIdx.x * blockDim.x threadIdx.x就是在“所有线程排成一维长队”后你所在线程的全局编号。if (idx n)是边界检查因为blocks * threads可能比n大多出来的线程什么都不干直接退出。启动参数blocks, threads里blocks是网格大小threads是线程块大小。前者一般是几万到几百万量级后者受硬件限制通常取128到1024之间。然后你会发现一个性能隐患这个Kernel每个线程只处理一个元素线程总量巨大如果数据量很大每个线程做完一件事就结束调度和销毁线程的开销是凭空浪费。更实际的做法是用grid-stride loop让每个线程循环处理多个元素__global__ void vecAddKernelLoop(float* A, float* B, float* C, int n) { for (int idx blockIdx.x * blockDim.x threadIdx.x; idx n; idx blockDim.x * gridDim.x) { C[idx] A[idx] B[idx]; } }这样你可以固定block数量比如让blocks等于GPU上SM数量的整数倍每个线程用一个固定步长遍历完整个数组减少线程创建开销也让工作总量可控。3.3 从Kernel到框架算子还差哪几块你写完一个Kernel只是第一步。要让PyTorch调它或者让它在推理框架里跑起来至少还要补上这几块内容设备侧的内存管理显存分配、数据拷贝、内存释放。在CUDA里是cudaMalloc/cudaMemcpy在Ascend C里是Global MemoryGM上通过Context管理。报错时最常见的就是内存没对齐或者没释放干净。梯度Kernel反向传播需要另一个Kernel。对向量加法而言梯度就是原样回传所以你的grad_kernel也是类似的向量遍历。对复杂算子比如矩阵乘来说反向Kernel的数学表达和前向完全不同工作量往往是前向的好几倍。框架封装在PyTorch中你要自定义一个torch.autograd.Function定义forward和backward两个静态方法在forward里调用你的Kernel在backward里调用你的梯度Kernel。更生产级的方案还会用torch.library注册自定义符号让算子能参与torch.compile的计算图优化。Shape推导与Dispatch策略框架需要知道你这个算子的输出shape怎么算。如果你希望支持多种dtypeFP32、FP16、INT8还要写分派逻辑让算子知道在哪个硬件、哪类数据上选哪个Kernel。我见过很多入门者卡在“Kernel能跑但框架不认”。其实框架封装的代码量不比Kernel少。你现在写一个向量加法如果走PyTorch扩展的路线完整工程往往包含setup.py、CUDA源文件、C封装头文件、Python包装类加起来一大坨。这也是为什么Triton越来越受欢迎——它把框架注册和Kernel写在同一个Python文件里省掉大量工程负担。我在实际项目中用Triton封装向量加法时调用侧可以非常干净import torch import triton def add_torch(x, y): out torch.empty_like(x) n x.numel() BLOCK_SIZE 1024 grid ((n BLOCK_SIZE - 1) // BLOCK_SIZE,) add_kernel[grid](x, y, out, n_elementsn, BLOCK_SIZEBLOCK_SIZE) return out x torch.randn(10000, devicecuda) y torch.randn(10000, devicecuda) assert torch.allclose(add_torch(x, y), x y)这段代码跑通后你就掌握了一个可复用的“算子开发闭环”定义数学逻辑写Kernel做边界处理框架调用验证正确性。3.4 验证和调试的几个实用技巧新手写算子最常见的现象是“跑是跑通了但结果不对”。我常用的排查顺序先写CPU朴素参考用纯Python或PyTorch本身实现一遍标准计算作为golden output。测试时不要直接用assert x y用torch.allclose(x, y, atol1e-4, rtol1e-4)因为GPU浮点累加顺序不同可能导致极小误差。测极端shape单元素、刚好等于一个Block大小、非对齐尺寸。很多Kernel崩溃都是因为边界外的线程/Block没有mask掉。打印关键中间值别只盯着最终输出。如果算子是多阶段流水把每阶段的数据分别拷回CPU判断是加载错了还是算错了。4. 性能瓶颈的本质访存模式、算子融合与软硬件协同4.1 先分清计算密集 vs 访存密集很多人以为算子慢是因为芯片算力不够但实际上绝大部分简单算子卡在访存上。向量加法就是最典型的访存密集算子计算内容只是一个浮点加法但要从显存读两个数、写一个数你干了三趟搬运。打个比方团队要搬家搬家公司派来的是一群博士生但他们的时间全花在路上真正在办公室里动脑思考工作的时间不到十分之一。这种情况下再怎么提高博士生的智商也没用因为瓶颈是车辆运输效率——对应到GPU里就是内存带宽。反观矩阵乘这类计算密集算子每一个输出元素需要K次乘加计算量远大于访存量这时候芯片算力才会成为瓶颈。分析性能的第一步永远先算一下这个算子的“算术强度”arithmetic intensity总计算量除以总访存量。算出来很小你就应该去优化访存算出来很大才去考虑指令效率和计算流水。社区里经常讨论“大量使用算子对硬件性能的挑战”核心也在这里模型里塞满了几百个Elementwise小算子每个算子都来一次“读取-计算-写回”哪怕单个算子访存不多串起来的重复读写在墙钟时间上是很可观的。这也是为什么编译器要帮你做算子融合。4.2 优化方向一让访存更高效访存密集算子优化的第一原则是减少搬运次数提高每次搬运的数据量。具体手法包括向量化访存在CUDA里每个线程用float4类型一次读写16字节等价于用一条指令做4个浮点数的搬运。计算Kernel时要注意把总元素数除以4并且要求内存16字节对齐。__global__ void vecAddVec4(float4* A, float4* B, float4* C, int n4) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n4) { float4 a A[idx]; float4 b B[idx]; float4 c; c.x a.x b.x; c.y a.y b.y; c.z a.z b.z; c.w a.w b.w; C[idx] c; } }保证合并访问同一线程块里的相邻线程最好访问显存里相邻的地址。GPU的显存控制器会把相邻线程的请求合并成少数几条大的内存事务。如果你的线程索引设计成跨步访问每隔几百个字节读一个内存事务会碎成几十条带宽利用率直接崩。数据布局调整图像数据常见的NCHW和NHWC之争本质就是让访存模式贴合硬件的连续读取特性。对卷积来说NHWC往往更符合通道维度连续读取的需求。利用片上内存如果多个线程会重复读同一块数据应该先把它搬进共享内存或者Unified Buffer而不是每次都去全局内存。这对应Ascend C里的CopyIn到UB再计算。4.3 优化方向二算子融合把多次搬运变成一次融合算子的动机其实很简单你看一个计算链路如果它是MatMul - PReLU不融合时是这样执行的MatMul Kernel从显存读两个输入矩阵把结果矩阵完整写回显存PReLU Kernel再把这个结果矩阵从显存读出来作用以后写回显存。两次Kernel两次全量访存。而融合后的逻辑是MatMul Kernel在计算某个输出tile时顺手把PReLU应用到这个tile上然后只写回一次。中间结果根本没有离开Kernel内部的寄存器或片上内存。数据量越大融合节省的访存越明显。实际工程里MatMulPReLU这类融合在CUDA和Ascend C里都是非常经典的算子。昇腾社区里讨论的“Ascend C融合算子MatMulPReLU”正是这条思路的典型代表用Ascend C的MatMul接口先做矩阵乘紧接着用矢量接口做PReLU激活两者在一个Kernel内部协同完成。需要注意融合对精度有潜在影响。如果你的融合链路里有PReLU、SiLU这类带指数或分段逻辑的激活函数在FP16下某些中间值可能会溢出如果分开跑两个Kernel编译器有可能会在两个算子之间偷偷做一次精度提升。融合方案设计时要充分考虑数值范围必要时在Kernel内混精度计算。4.4 哪些算子值得融合、哪些不值得融合不是越多越好。我自己的经验是拿两个指标来判断访存节省占比融合能省掉多少次全量中间数据的读写如果中间结果远小于输入输出总量融合收益就小。实现复杂度融合Kernel的代码复杂度上升多少调试难度增加多少值得融合的典型场景连续Elementwise算子链比如Add - ReLU - Mul它们都是逐元素操作融合后没有任何语义损失访存从三次变一次。卷积/矩阵乘后接激活函数这是推理引擎里最常见的融合优化几乎无处不在。小算子合并成大算子某些模型里会出现大量极小的算子每个算子启动都有固定开销合并后能显著降低调度开销。不值得融合的场景跨设备通信类算子如果融合要把数据从多卡之间来回搬运那么在Kernel内部做跨卡通信代价极高不如保持原子操作。融合后寄存器溢出一个Kernel内部塞太多逻辑导致寄存器使用超过硬件限制编译器被迫把变量塞进局部内存local memory性能反而崩。后续算子几乎不访存已计算数据比如融合后马上要做一个稀疏操作而稀疏逻辑本身就是随机访存的你辛苦算好的中间结果根本复用不上那就别费劲。判断公式可以说成融合收益≈省掉的中间数据读写量×读写代价−Kernel复杂度上升导致的性能损失。前者越大、后者越小越值得融。5. 入门路线与避坑清单从能跑通到能干活5.1 建议的入门路径我自己带过几个从算法岗转到算子开发的同事整理出一条比较顺的路径大约4到8周能建立整体感觉第1周观察。不需要写任何Kernel。挑一个简单模型比如ResNet18用PyTorch的profiler跑一遍导出每个算子的耗时和访存量。你很快会找到耗时的top算子然后意识到“模型性能问题不是模型结构问题是算子实现问题”。第2周复制。用Triton写一个向量加法并成功注册给PyTorch调用。这一步主要是建立“Kernel不是神秘魔法”的感觉。第3周扰动。在Triton Kenerl里故意改BLOCK_SIZE、改循环方式、改mask条件观察性能变化和错误类型。这一步会让你直观理解线程组织和边界条件。第4周融合。尝试把Add ReLU合并成一个Kernel对比融合前后的耗时。你会发现融合效果显著由此体会“访存搬运决定上限”这句话的重量。第5周之后平台深耕。如果你是做昇腾的开始看Ascend C官方文档和认证课程了解CopyIn/CopyOut、流水线同步、MatMul接口如果你是CUDA路线去读CUDA C Programming Guide重点读内存层次和线程模型然后是Nsight Compute的优化案例。5.2 踩过的坑版本、维度、精度、边界这些坑我几乎全踩过写出来给你省时间CUDA版本和PyTorch预编译包不一致系统里装的是PyTorch默认CUDA 12.1但你自己编译扩展用的nvcc是11.8跑起来偶尔报奇怪的段错误。解决办法是强制让编译环境与PyTorch构建时的CUDA版本一致最好直接用官方Docker镜像。维度看成错拿到的是一个四维Tensor展平成一维后忘了按NCHW顺序还原算出来结果不对。这个错最坑的地方是shape校验都能过只有数值对不上。FP16精度FP16下加减法很容易溢出。对于大数值的加法在Kernel内部先把FP16转成FP32做累加最后再转回FP16输出这个习惯要养成。忘记maskTriton里如果数组长度不是BLOCK_SIZE整数倍你不加mask越界部分读取的是垃圾数据有时候不报错结果却是一个莫名其妙的数。这类问题排查时非常消耗时间建议从第一天就养成“凡是非对齐一律mask”的习惯。容器里看不到GPU用Docker跑深度学习时容器内看不到显卡多半是--gpus all没加或者NVIDIA Container Toolkit版本对不上。排查这类环境问题往往是入门阶段最耗时的部分。5.3 值得依赖的周边工具与认证工具链能帮你省掉大量瞎猜的时间。英伟达平台用Nsight Compute做Kernel性能分析它告诉你每个Kernel的访存吞吐、计算吞吐、占用率、是否存在内存合并失败昇腾平台对应的是CANN自带的Profiler工具可以看AI Core的利用率、数据和指令流水重叠情况。认证方面**华为昇腾社区推出的“Ascend C算子开发能力认证中级”**我认真看过课程大纲覆盖了算子的工程结构、矢量编程、流水线同步、性能分析这些核心内容。如果你本来就要做昇腾适配那直接跟着认证体系学比到处找散装资料效率高得多。如果你对编译器方向感兴趣可以了解一下LLVM在算子自发现和自动调优上的进展——很多新一代算子工具链都在做“声明数学语义编译器自动找Kernel实现”的事情这可能是未来算子开发范式的一个重要变化。5.4 算子概念远不止神经网络最后补一个视野层面的建议。你在热词里能看到“Sobel算子”“Canny算子”“拉普拉斯算子”这些词它们在图像处理里早已存在表达的也是“对数据做特定变换”的数学概念。而“神经算子”则是科学计算里这两年非常热的路线——用神经网络直接学习算子映射。所以说算子这个抽象并不只属于AI框架它是整个计算世界的通用语言。当你从AI框架的视角进入算子开发后再去接触图像算子、科学计算算子会发现底层的思维模式完全一致定义数学操作、设计数据流动、优化硬件映射。这也是为什么我很推荐AI工程师花时间跨出模型代码走到底层去看一看。你会获得一种“软件和硬件之间我能说了算”的掌控感。我自己在实际操作中的体会是算子开发入门最难的从来不是语法而是思维模式的切换从“我要实现什么功能”切换到“数据是怎样在硬件里流动的”。先跑通再看搬运再谈融合这条路基本不会走偏。
