上个月帮一个做推荐算法的朋友排查线上推理变慢的问题。他给我看模型代码前向算下来也就几十个算子调用怎么看都不该慢成那样。结果问题不出在模型结构而是落在某个自定义算子没有适配推理引擎的高效执行路径上框架兜底走了一条通用但很慢的 kernel。这件事让我想写一篇关于“AI算子”的入门文章很多算法工程师每天都在调用算子但算子本身是什么、为什么同一个数学公式会有完全不同的执行速度、以及当框架没有现成算子时该怎么自己动手写这些其实很少有人系统性讲过。这篇内容适合几类人刚开始接触深度学习框架、好奇底层实现的在校学生工作中频繁踩到性能问题但不太清楚从哪入手的算法工程师以及做过传统开发、想了解 AI 计算核心逻辑的老朋友。我会从概念边界讲起给出手写算子的完整示例也会分享我在实际项目里踩过的坑和排查思路。看完之后你至少能从“会用算子”迈入“看得懂算子、改得动算子”这个区间。1. 算子不是“算法”——先把概念边界划清楚1.1 一句话认识算子输入Tensor输出Tensor“算子”这个词源自动态数学里的 operator 翻译在深度学习领域它被用于描述 Tensor 到 Tensor 的一个计算单元。你不用把它想得多神秘一个 python 函数接收若干 Tensor经过数学变换再输出 Tensor这个函数对应的底层实现就是算子。举几个常见的例子你立刻就能对上号元素级算子ReLU、Sigmoid、Add、Mul输入和输出 shape 相同逐个元素做计算规约算子Sum、Mean、Softmax、LayerNorm往往在某个维度上聚合信息矩阵算子MatMul、Conv2D、ConvTranspose整个推荐系统、CV 模型的核心计算都靠它数据搬运算子Reshape、Transpose、Permute不改变数值本身但改变张量在内存中的排列方式。你平时写torch.relu(x)、F.linear(x, w, b)其实就是在向框架请求执行某个算子。框架做的主要事情就是把你 Python 端的“算子请求”翻译成具体硬件上的执行代码并像工厂流水线一样把一个个算子的输入输出接好。1.2 Kernel、OP、算子同一个词的三副面孔这在社区里非常容易绕晕因为不同资料用词不一样。实际上OPOperation偏上层概念表达“我要做什么”比如 PyTorch 里的torch.nn.functional.relu是一个 op。Kernel偏执行层概念指“在硬件上具体怎么算”。同一款 op有不同的 kernelCPU kernel、GPU CUDA kernel甚至 CPU 上还有针对不同指令集AVX、NEON优化的 kernel。“算子”这个词在很多时候是 OP 和 Kernel 的总称看语境偏哪一层。一个入门的诀窍当资料说“添加一个新算子”时通常要同时完成上层 OP 定义和底层 Kernel 实现以及把它们注册到框架的调度器里。研究框架源码时你会发现aten/src/ATen/native目录下存放的就是各类 CPU 算子实现而MACE/ONEAPI/TVM这类项目里讨论的 kernel 则偏向下层硬件代码。1.3 算子与计算图自动微分依赖的“砖块”另一个常见误解是把算子等同于神经网络的“层”。其实层是逻辑单元算子才是真实计算单元。例如一个nn.Linear层内部至少包含matmul与addbias两个算子反向还对应matmul的另一个隐式算子。模型训练中所谓反向传播就是沿计算图把算子之间的梯度求导规则逐一应用。框架的 autograd 引擎维护完整的计算图图上每个节点就是一个算子执行点边是 Tensor 依赖。因此算子不只是在做数学运算它还承接了两个额外职责作为可微函数提供 forward 与 backward 两个方向的实现接受框架调度器的参数选择不同 dtype、device、layout 下自动匹配不同 kernel。这就是为什么从业人员常说“算子是全栈能力的电子侦查集合点”往下看是硬件指令往上看是模型语义。2. 为什么“三行代码”的性能差几十倍算子是性能的第一现场2.1 同样一个卷积三重循环和库实现差距在哪我先说一个反直觉的事实一个数学定义上完全相同的卷积手写三重循环与高性能库实现做对比性能差距可以到几十倍或上百倍。比如一个 3x3 卷积朴素实现按 5 重循环展开每个输出像素对每个输入通道取值、乘加、写回内存而 cudnn 或 oneDNN 的实现会考虑内存布局、数据复用、SIMD 向量化、分块tiling等把同一个公式以完全不同的方式映射到硬件。用做饭做个类比数学公式是菜谱硬件是灶台而算子就是厨师。同一个酸辣土豆丝配方新手灶台前手忙脚乱是一桌菜经验丰富的师傅掂勺翻锅、火候顺着锅气走出锅完全两回事。菜谱没变问题在于“执行策略”。为什么要强调这一点因为很多入门者调试性能问题时第一反应是“模型结构是不是有问题”很少去怀疑算子本身的执行策略。但在深度学习中算子执行的效率往往比算子“数量”更影响最终速度。2.2 框架调用的实际路径Python 端的“水面之下”PyTorch 哲学的入门通常只说了 Python 用法但实际一次torch.relu(x)调用背后会经过这么几步解析 Python 参数绑定到torch.relu的 C 入口进入 ATen dispatcher根据 Tensor 的 dtype、device、layout 找到对应的 kernel 注册项如果是 CUDA tensor则从算子库 dispatch 到 CUDA kernelkernel 在 GPU 上按线程网格并行执行结果写进输出 Tensor 内存返回 Python 对象。这个调度过程里很多策略与性能强相关。比如算子融合就是“减少调度次数”的典型思路把add relu融合为add_relu一个 kernel减少一次 Tensor 读取、一次写回以及一次 kernel launch 的开销。框架层面有手动融合也有torch.compile这样的编译器自动融合。理解了调度机制你会更容易看懂框架为什么某次升级后模型快了不少——不一定是模型层面变了很可能是某个算子的 kernel 被换成了更高效版本。2.3 什么情况下你必须研究算子并不是每个项目都需要深入算子算法工程师可以先跑通再说。但下面这些场景出现时你绕不开模型里有一种标准算子无法表达的操作需要自定义实现已有算子拼出来的功能性能不理想需要把它融合成更高画面效率的单个算子从训练框架切到推理引擎TensorRT/ONNX Runtime后自定义算子不被原生支持导致整图优化被打断遇到精度问题需要对算子内部的累加顺序、归一化位置做精细化控制。提前知道这些事情至少能让你在遇到模型变慢或者部署失败时少走很多弯路。3. 手写一个 ReLU从数学到可运行实现3.1 先写一个“慢但绝对正确”的版本网上讲算子往往一上来就上 CUDA容易劝退初学者。我自己带人入门算子开发时会先要求他在 Python 侧把算子的语义写对再做性能版本。这个“慢但正确”的版本是所有后续工作的校验基准。挑选 ReLU 来实操是因为它足够简单又包含关键要素前向要对原始输入做max(0, x)反向要把上游梯度原样传给输入中大于 0 的位置等于 0 的位置梯度可以约定为 0不同框架可能微调但 PyTorch 沿用的是 0。用 NumPy 实现import numpy as np def relu_forward(x: np.ndarray) - np.ndarray: return np.maximum(x, 0.0) def relu_backward(dout: np.ndarray, x: np.ndarray) - np.ndarray: return dout * (x 0).astype(dout.dtype)别小看这几行代码。它明确了一个重要原则算子的反向不是独立定义的而是对应前向的“分段导数规则”。写算子前先把 forward 与 backward 的数学逻辑写死在注释里这比一上来就写 CUDA 值得多。3.2 用 autograd.Function 接入反向传播PyTorch 给自定义算子提供了标准入口就是继承torch.autograd.Function。我们在这里实现前向与反向而不是简单调用 Python 函数这样框架的 autograd 引擎就能识别算子的求导路径。import torch class ReLUFunc(torch.autograd.Function): staticmethod def forward(ctx, x): # ctx 用于在 forward 中保存反向需要用到的信息 ctx.save_for_backward(x) return x.clamp_min(0) staticmethod def backward(ctx, grad_output): (x,) ctx.saved_tensors # 这里一定要记得把 mask 转成与梯度一致的类型 return grad_output * (x 0).to(grad_output.dtype)之后在模型里就能这样调用y ReLUFunc.apply(x)这样算子逻辑虽然写在 Python 端却已经通过autograd.Function具备“可微分性”了。需要特别注意的是forward接收几个 Tensor 输入backward就应返回对应数量的梯度非 Tensor 的上下文性输入比如一个浮点参数通常返回None这一点很多入门者在这里被刺。3.3 小实验验证前向与反向代码写出来不代表对必须有验证实验。最直接的方式是和框架自带实现对照torch.manual_seed(42) x torch.randn(4, 4, requires_gradTrue) out_custom ReLUFunc.apply(x) x2 x.detach().clone().requires_grad_(True) out_ref torch.relu(x2) # 前向一致性 print(forward match:, torch.allclose(out_custom, out_ref)) # 反向一致性给一个任意上游梯度 grad_seed torch.randn(4, 4) out_custom.backward(grad_seed) out_ref.backward(grad_seed) print(backward match:, torch.allclose(x.grad, x2.grad))“反向一致性”是算子开源社区非常常见的 CI 测试方式。凡是写新算子我都会把这种对照测试固化到测试文件里避免后续优化时把一个看似等价的新 kernel 换进来结果把梯度悄悄“优化”错了。3.4 算子注册到框架的关键点如果你想让做出来的算子成为一个真正可被框架调度的算子而非仅在一个模型文件中可用就需要完成注册。PyTorch 的实践路径大致是用 C 扩展或torch.ops注册命名空间和算子函数注册一个或者多个 kernel每个 kernel 对应一种设备/数据类型对需要在 TorchScript 或 ONNX 里用的算子还要补充 schema输入输出类型约束与 shape 推断逻辑。C 自定义算子典型的骨架是这样这里只展示 CPU 版本#include torch/extension.h torch::Tensor relu_forward(torch::Tensor x) { return at::clamp_min(x, 0); } TORCH_LIBRARY(my_ops, m) { m.def(relu_forward, relu_forward); }写 Python 也好、写 C 也好你的目标其实相同让框架知道“有这么个算子这个算子的计算语义是什么在什么设备上该执行哪个实现”。真正难的东西不在注册代码本身而在于写出的实现在不同 dtype、不同 shape、不同内存格式下都能保持正确。4. 让算子变快之前先理解计算与存储的“瓶颈”4.1 算子的速度瓶颈往往在内存带宽而不是算力我见过很多刚接触性能优化的人总是下意识认为优化就是把“计算次数”变少于是先数 FLOPs。但实际上对大部分 elementwise 算子和不少规约算子来说瓶颈根本不在计算而在内存访问。算一下就明白了假设要处理 1GB 的 float 数据计算只是让每个元素和 0 比大小。如果内存带宽是 50GB/s读一遍数据就要 20ms你即便把“和 0 比大小”这个操作优化到零成本总时间也不会低于 20ms。这时候性能天花板叫“内存墙”有些资料也直接叫带宽瓶颈。所以对 elementwise 类算子真正有效的手段是减少内存访问次数。这就是为什么addrelu融合有意义原本读两个数、写两个数现在读两个数、写一个数节省了一次写回和一次 kernel launch。而在显存带宽受限的场景下这种融合往往比减少 FLOPs 更直接。4.2 三个最常用的优化方向融合、分块、向量化融合Fusion把多个算子的数据流合并到一个 kernel减少中间 Tensor 的读写。框架层面torch.compile已经在做手写自定义算子时也可以主动设计。分块Tiling主要针对 MatMul 这类计算量大、单个矩阵放不进缓存的问题。把输出矩阵切块让计算过程尽量复用 cache 里的数据块从而降低对内存带宽的依赖。向量化Vectorization一次处理多个元素的加载和计算。CPU 上依赖 SSE/AVX 指令GPU 上则是让线程做连续地址访问保证内存 coalescing合并访存。对一个入门项目而言我的建议流程是先确认内存访问模式是否连续 → 再考虑能否和相邻算子融合 → 然后才考虑向量化等微优化。顺序错了经常白忙活。4.3 学会测量算子的真实开销优化之前必须会测量。入门阶段做两件事就够了用torch.profiler或框架自带的 profile 工具记录各算子的耗时占比单独测一个算子的纯执行时间用多次重复加torch.cuda.synchronize()消除异步执行的影响。import torch import time x torch.randn(1024, 1024, devicecuda) # 预热 for _ in range(10): y torch.relu(x) torch.cuda.synchronize() start time.perf_counter() for _ in range(1000): y torch.relu(x) torch.cuda.synchronize() avg_ms (time.perf_counter() - start) / 1000 * 1000 print(favg: {avg_ms:.3f} ms)此时算一下数据的有效读写量再除一下显存带宽就能大概估算“理论下限”。如果实际耗时是理论上限的几倍说明还有很大优化空间如果已经贴近带宽上限那就说明这个算子本身能挖的不多了。建立这种“先估上限再对比实测”的思维比盲目套优化手段要有用得多。5. 算子开发四大坑我的踩坑记录5.1 dtype不一致输出正常梯度全乱第一次写自定义算子时我在反向里直接写了grad_output * (x 0)前向和反向都通过了我的对照测试因为当时输入也是 float32。后来换成混合精度训练输入 x 变成 float16但x 0得到的是 bool 张量和 float16 的梯度相乘时类型提升了一阵看似没报错实际上梯度值全部悄悄变成了 float32再传给下游算子时就出现了精度衰退问题。真正的教训是反向 mask 一定要显式转成与 grad_output 相同的 dtype。这也是为什么我在上文代码里特地写了.to(grad_output.dtype)。算子实现里dtype 是所有后续逻辑的地基这个地基出了偏差问题往往要到几个算子之后才爆发。5.2 backward 返回值数量必须“斤斤计较”另一个高频报错来自autograd.Function的 backward 返回值约定。写得复杂一点的算子forward 输入往往是多个staticmethod def forward(ctx, x, bias, scalar_param): ctx.save_for_backward(x, bias) return x bias scalar_param如果 backward 里写了三个 return但 scalar_param 本身又不是 Tensor 参数比如普通浮点框架会直接报错。正确写法是staticmethod def backward(ctx, grad_output): x, bias ctx.saved_tensors return grad_output, grad_output, None这条规则想通之后就很简单却真的很容易在实际算子代码里搞漏。每当我看到一串 “return grad_x, grad_w, None, None” 的代码都会想起来“有几个输入就还几个梯度”这句话真的写着写着就容易手滑。5.3 动态 shape 与 JIT 编译平台的冲突我曾经把一个小算子封装到 TorchScript 里用于生产推理。本地跑没问题一到线上容器里新传入的 batch 大小和之前不同就报找不到对应 kernel 的错误。原因在于 TorchScript 在推理时会针对特定 shape 做静态化编译而我的自定义算子没有提供足够的 shape 提示或 fallback 路径。框架默认只有在 shape 匹配的情况下才用编译缓存。解决办法对自定义算子的 JIT 使用场景补充 schema 标注并且明确告诉框架这个运行时的动态范围。更稳妥的是在部署前把动态维度固定或做 padding 对齐减少跨 shape 的编译消耗。这块暴露出的深层问题是算子开发者不能只保证“Python 里能跑”还得考虑它在其他执行引擎里的可编译性。5.4 ONNX/TorchScript 导出时 shape 推导失效把包含自定义算子的模型导出成 ONNX 时框架并不知道你这个算子的输出 shape 应该是什么。于是导出的计算图里要么多了一个未知节点导致推理引擎无法继续做图优化要么完全导出失败。正确的工程做法是为自定义算子注册 shape 推断函数。现实中很多团队图省事选择在导出前把自定义算子替换成标准算子组合比如把融合算子“展开”成 add relu。这虽牺牲推理性能但换来跨平台兼容性属于部署阶段的合理折中。当你学习算子开发从“跑通”走向“上线”这类问题就会成为主要的工作内容。6. 如果我现在重新入门算子开发会怎样安排学习路径6.1 先选一个熟悉框架吃透一个简单算子的完整链路如果让我重新入门我不会同时追 PyTorch、TensorFlow、JAX 三家生态而是挑自己最常用的一个框架用 ReLU 或 Softmax 做研究对象。具体做法在 Python 里调用这个算子用 profile 工具看耗时时长找到它对应的 C/CUDA 实现源码路径用上文讲的方法自己重写一份等价实现并做正确性对比再尝试和相邻算子做融合或向量化比较前后性能变化。这套循环下来你对“框架如何调算子、算子在硬件上如何执行”的理解会比读十篇博客更扎实。6.2 从“调用算子”到“阅读算子实现”阅读实现源码没有想象中困难。拿 PyTorch 举例aten/src/ATen/native/cpu/SoftmaxKernel.cpp这类文件中你能看到实际 kernel 的循环写法、内存布局处理以及针对最后维度优化的特殊情况。刚开始看不懂没关系按“找到入口 → 找到核心循环 → 理解数据索引 → 理解并行方式”这个顺序去读即可。值得提醒的是源码里大量存在“特例”。例如某些算子在最后一个维度上能走快速路径在中间维度上则要转置或复制数据。这些特例行代码往往才是性能的秘密所在也最容易让初学者困惑。我的建议是先忽略特例抓住通用路径理解后再回头研究为什么需要特例。6.3 站在生态上做选择最后说点我的个人感受算子开发并不是“必须手写 CUDA 才算入门”。如今很多工作可以在高层完成比如用torch.ops注册现有 C 实现、用 Triton 编写 GPU 计算、依赖编译器来自动优化。但有一点始终不会变你必须对你实现的算子在内存和计算上的行为有准确预期否则框架再智能也无法帮你兜底。把基础概念、执行流程、调试方法掌握好之后无论是转向推理引擎算子适配、编译器算子实现还是继续做模型侧算法开发这一层认知都会时刻发挥作用。这是我认为这部分知识对算法工程师最有价值的长期回报。