企业官网建设流程全解析

1. 为什么需要自定义算子

1.1 PyTorch 算子的两种来源

我一直觉得，很多人对“自定义算子”这个词的第一反应是“这是那些做框架的人才会碰的东西”。实际上 PyTorch 每天在用的所有功能，无论是torch.add还是conv2d，本质上都能归成两类：一类是官方已经编译进 C++ 核心的算子，另一类是用 Python 组合出来的高层模块。官方算子覆盖了 95% 的日常需求，可一旦你开始做模型部署、性能优化、或者写一篇新论文里的特殊函数，官方算子经常不够用，这时候你就会想自己动手写一个。

自定义算子的含义其实很朴素：我们自己写一段 C++ 代码，或者 CUDA 代码，把它编译成一个 Python 可以调用的函数。这个函数和torch.sigmoid长得一模一样，能接收 Tensor、返回 Tensor，差别在于它的计算逻辑由你来定义。这意味着你可以把 Python 层很难表达的循环、内存访问模式、或者某个第三方 C++ 库，原封不动地接进 PyTorch 的计算图里。

我这里说的“吃透”，不是只讲 API 怎么调，而是想把背后的机制拆出来：PyTorch 怎么把你的 C++ 函数变成 Python 能导入的模块？CUDA kernel 怎么在 GPU 上被调度？反向传播怎么接到 autograd 上？只有把这三件事弄明白，你才能在一张白纸上设计自己的算子，而不是照着官方示例改个名字就完事。

1.2 什么场景必须自己写算子

判断依据其实很简单，遇到下面四种情况你就要认真考虑上自定义算子。

第一种是“计算逻辑没法用 PyTorch 原生算子拼出来”。比如说你要做某种自定义的稀疏索引操作，或者实现一个拉普拉斯算子的变体，需要逐个遍历邻域像素并做条件更新，用 Python 写for循环在 GPU 上会直接卡死，用torch操作东拼西凑又很别扭，这种情况写成 CUDA kernel 才是正路。

第二种是“能拼出来但性能差得离谱”。最典型的就是多个逐元素算子串联，比如y = a * b + c，如果拆成mul和add两个 kernel，GPU 要启动两次、读写两遍内存，带宽成本直接翻倍。如果把它融合成一个 CUDA kernel，一次遍历全做完，速度可能提升 20% 甚至 50% 以上。我在做图像滤波的时候遇到过类似问题，一个简单的权重滤波算子拆开写，处理一张 4K 图像每次都要多等两个 kernel launch，融合之后肉眼可见的流畅。

第三种是“要复用一个已有的 C++/CUDA 库”。比如某些工业软件里的 Halcon 滤波核或者自研相机 SDK 已经提供了高效的 C++ 函数，你不想把它翻译成 Python，更不想用 ctypes 一点点抠内存，直接从 C++ 扩展里包一层 Tensor 封装是最高效的接法。

第四种是“需要自定义反向传播逻辑”。PyTorch 的 autograd 已经把所有标准算子都定义了梯度，但如果你的前向过程有分段不可导点，或者你想在反向里注入某种额外的梯度变换，就必须自己控制反向函数。自定义算子里的backward让你拥有绝对主动权。

这里我给一个实用建议：不要为了自定义而自定义。如果你的需求能用torch.compile或者torch.vmap解决，优先用这些高层工具。它们的开发成本低，而且能自动做算子级优化。只有当你需要极致的控制权时，才决定走 C++/CUDA 这条路。这个判断会帮你省掉大量调试时间。

2. 动手前的工具链准备

2.1 理解 PyTorch 内置的扩展机制

PyTorch 官方早就把“自定义算子”的通道给铺好了，核心就是torch.utils.cpp_extension这个模块。它本质上是一个编译器包装器：拿到你的 C++ 源文件和 CUDA 源文件之后，会调用编译器把它们编译成动态库，再用 pybind11 把函数导出成 Python 可调用的对象。你不需要手动写完整的setup.py和 pybind11 绑定，虽然自己写也可以，但官方模块已经处理了大部分坑。

这套机制里其实还有底层注册逻辑，叫TORCH_LIBRARY。通过它可以在 PyTorch 内部注册一个带命名空间的算子，让 PyTorch 的调度器认识你。官方推荐的现代写法是TORCH_LIBRARY(my_ops, m) { m.def("custom_add", &custom_add); }，然后用torch.ops.my_ops.custom_add在 Python 端访问。很多老教程喜欢用PYBIND11_MODULE直接导出普通函数，差别在于，PYBIND11_MODULE导出的函数不参与 PyTorch 的 op 调度，也没有办法被torch.jit或者torch.compile完整追踪。如果你只是给自己跑实验，怎么搞都行；如果要进正式模型或者部署链路，建议一开始就用TORCH_LIBRARY的姿势。

2.2 版本匹配与依赖安装

PyTorch 自定义 C++/CUDA 算子，最怕的就是版本错位。CUDA 编译器必须和 PyTorch 编译时使用的 CUDA 版本兼容，不然会链接到一半报undefined reference to cudaMemcpy，或者显卡驱动不支持，跑起来直接报CUDA driver version is insufficient。建议先跑一下python -c "import torch; print(torch.version.cuda)"，拿到当前 PyTorch 对应的 CUDA 版本号。如果你使用的是 conda 安装的 PyTorch，通常自带配套 CUDA 库，安装 CUDA Toolkit 时选择相近版本即可。

Windows 下还有一个隐藏依赖：Microsoft Visual C++ Redistributable。很多人奇怪为什么装了个 Python 包还要管 VC++ 运行库，因为 PyTorch 的 C++ 扩展在 Windows 上依赖 MSVC 编译器和运行时库。建议提前安装 Visual Studio 2019 或 2022，尤其要注意勾选“使用 C++ 的桌面开发”工作负载，否则命令行里找不到cl.exe。至于 Redistributable，它一般会随 VS 一起装好，单独下载也有用，但真正编译时你需要的是编译器，不是运行库。把这两者分开理解，能少翻很多帖子。

我在 Ubuntu 上工作的时间比较长，那边主要用 g++。但说实话，Ubuntu 上的坑也不比 Windows 少，最典型的是系统 gcc 和 PyTorch 官方预编译用的 gcc 版本不一致，后面编译时你会看到一大堆莫名其妙的符号找不到。我的建议是尽量用系统默认的编译器版本，或者参考 PyTorch 官方文档里的支持矩阵。避坑的核心就是一句话：让你的编译环境和 PyTorch 的编译环境保持同频。

2.3 从零搭建一个最小可编译工程

我平时写自定义算子，文件夹结构非常简单。无论你是要用setup.py还是load_inline，都有几样东西是必备的：一个.cpp文件，负责定义 C++ 函数和绑定；一个.cu文件，负责写 CUDA kernel；一个setup.py，负责告诉编译器怎么组织和输出。目录大致长这样：

my_op/ ├── setup.py ├── my_add.cpp ├── my_add_kernel.cu └── test.py

setup.py最简单的写法如下：

from torch.utils.cpp_extension import CUDAExtension, BuildExtension from setuptools import setup setup( name='my_add', ext_modules=[ CUDAExtension( name='my_add', sources=['my_add.cpp', 'my_add_kernel.cu'], extra_compile_args={ 'cxx': ['-O3'], 'nvcc': ['-O3'] } ) ], cmdclass={'build_ext': BuildExtension} )

然后执行python setup.py build_ext --inplace。如果顺利，会在文件夹里生成一个.pyd或.so文件，之后就能import my_add。这里我想强调一下extra_compile_args：如果你是第一次编译，先别加花里胡哨的-arch标志，让 PyTorch 自己根据torch.cuda.get_device_capability()决定。编译通过以后再慢慢调-gencode=arch=compute_80,code=sm_80这类参数，这样能少踩很多“编译错误一半是熄火一半是警告”的坑。

3. 核心实战：C++/CUDA 算子怎么落地

3.1 用 C++ 扩展写一个带反向的算子

我拿一个最简单但完整的例子说明：实现z = (x + y) * w。这在 PyTorch 里当然一行就写完了，但它的反向要回传三个梯度，很适合用来展示 autograd 对接。先看 C++ 端定义：

#include <torch/extension.h> torch::Tensor my_forward(torch::Tensor x, torch::Tensor y, torch::Tensor w) { return (x + y) * w; } std::vector<torch::Tensor> my_backward(torch::Tensor grad_output, torch::Tensor x, torch::Tensor y, torch::Tensor w) { auto grad_w = grad_output * (x + y); return {grad_output * w, grad_output * w, grad_w}; }

然后需要注册成一个 autograd 函数。在 PyTorch 的 C++ API 里，可以使用torch::autograd::Function<派生类>：

class MyCombinedOp : public torch::autograd::Function<MyCombinedOp> { public: static torch::Tensor forward(AutogradContext* ctx, torch::Tensor x, torch::Tensor y, torch::Tensor w) { ctx->save_for_backward({x, y, w}); return my_forward(x, y, w); } static std::vector<torch::Tensor> backward(AutogradContext* ctx, std::vector<torch::Tensor> grad_output) { auto saved = ctx->get_saved_variables(); auto x = saved[0]; auto y = saved[1]; auto w = saved[2]; return my_backward(grad_output[0], x, y, w); } }; torch::Tensor my_combined_op(torch::Tensor x, torch::Tensor y, torch::Tensor w) { return MyCombinedOp::apply(x, y, w); } TORCH_LIBRARY(my_ops, m) { m.def("my_combined_op", my_combined_op); }

注意forward和backward都是静态成员函数，apply负责把参数传进去并触发 autograd 记录。save_for_backward保存前向用过的变量，反向时调用get_saved_variables取出来。可能你会问：为什么不直接在 Python 里写torch.autograd.Function？原因很简单，因为这里我们在演示 C++ 侧完成整个闭环，这样编译后的算子不依赖 Python 解释器的 Function 机制，调度性能更高，也更接近原生 op。

写好这段之后，最重要的就是注册。上面代码用了TORCH_LIBRARY，这是 PyTorch 1.10+ 推荐的方式。需要注意my_ops是命名空间，Python 端调用时使用torch.ops.my_ops.my_combined_op。如果你希望让 Python 直接import my_add然后调用普通函数，也可以用PYBIND11_MODULE，但正如前面说的，那样 autograd 不会自动帮你记录反向。通常的做法是：要么用Function包一层并导出普通函数，要么用TORCH_LIBRARY注册，Python 端再包一个torch.autograd.Function定义反向。我用上面的写法是为了展示纯 C++ autograd，实际项目中你完全可以根据需求混用。

3.2 CUDA kernel 的写法与线程分配

真正大规模计算时，我们希望算子跑在 GPU 上。CUDA 端的基本模式是：先写一个__global__函数作为 kernel，再在 host 函数里设置 grid 和 block 大小，最后把 Tensor 的 data_ptr 传给 kernel。以最简单的逐元素加法为例：

__global__ void vector_add_kernel(const float* a, const float* b, float* out, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { out[idx] = a[idx] + b[idx]; } } torch::Tensor vector_add_cuda(torch::Tensor a, torch::Tensor b) { auto out = torch::empty_like(a); auto n = a.numel(); const int threads = 256; const int blocks = (n + threads - 1) / threads; vector_add_kernel<<<blocks, threads>>>( a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), n); return out; }

这段代码里最关键的两个参数是threads和blocks。threads我固定为 256，是因为对大多数 GPU 来说，一个 block 里 128 到 512 个线程通常能获得不错的占用率。blocks的计算方法是一个经典公式：(元素总数 + 线程数 - 1) / 线程数，这样能保证即使最后一个 block 不满也不会漏掉元素。kernel 内部通过if (idx < n)做边界检查，防止越界访问。

真实工程里，我强烈建议再加上 CUDA 错误检查。比如内核启动后立刻调用cudaGetLastError()或者AT_CUDA_CHECK(cudaDeviceSynchronize())，不然一个越界访问可能不会立刻报错，而是等到下一帧才给你一个莫名其妙的illegal memory access。另外，a.data_ptr<float>()返回的是指针，必须确保传入的 tensor 在 GPU 上、类型是 float、内存连续。一般我会在函数开头写三个检查：

TORCH_CHECK(a.is_cuda(), "input must be CUDA tensor"); TORCH_CHECK(a.scalar_type() == torch::kFloat, "input must be float32"); TORCH_CHECK(a.is_contiguous(), "input must be contiguous");

这三个检查会避免 90% 的奇怪崩溃。如果发现输入不连续，你可以直接用.contiguous()拷贝一份再传进去，只是要记住这会增加一次内存拷贝，性能敏感时要提前在 Python 侧保证数据布局。

3.3 Autograd 怎么接管梯度

有人问：前向写完了，反向怎么自动接上？其实原理很简单。torch::autograd::Function<派生类>重写了forward和backward两个接口。当你调用apply的时候，PyTorch 的 autograd 引擎会创建一个节点放在计算图里，节点里保存了你通过save_for_backward留下的变量。反向传播时，引擎从最终输出往输入方向走，每经过一个节点就调用对应节点的backward，把上游梯度传进去，拿到回传给前向输入的梯度。

我上面那个例子，反向实现的是链式法则。z = (x + y) * w，设g是上游梯度，那么对x的梯度是g * w，对y的梯度也是g * w，对w的梯度是g * (x + y)。可以看到我在my_backward里就按这个公式算。这里有个非常重要的细节：backward返回的std::vector<torch::Tensor>顺序必须和forward的输入顺序一致。也就是说，forward接受的第一个参数 x、第二个 y、第三个 w，那么backward返回的三元组就分别对应 x、y、w 的梯度，顺序千万不能乱。我刚开始写的时候就在这儿翻过车，梯度不报错，但对不上号，最后只能靠打印每个变量的梯度检查。

另外需要注意，自动微分对传进去的 Tensor 默认会做版本控制，如果前向过程中你对保存的变量做了原地修改，梯度计算会报错。如果你是优化内存的狂热爱好者，用了set_data或者copy_这类操作，一定要想清楚会不会破坏 autograd 的版本计数器。

3.4 Python 端加载与调用

如果你不想搞一整个 setup.py 工程，PyTorch 提供了load_inline这个神器，可以直接在 Python 脚本里把 C++/CUDA 源码编译成一个模块。它的典型用法是：

from torch.utils.cpp_extension import load_inline cuda_source = """ #include <torch/extension.h> #include <cuda_runtime.h> __global__ void add_kernel(const float* a, const float* b, float* out, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) out[idx] = a[idx] + b[idx]; } torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) { auto out = torch::empty_like(a); const int threads = 256; const int blocks = (a.numel() + threads - 1) / threads; add_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), a.numel()); return out; } """ module = load_inline( name='inline_add_op', cuda_sources=[cuda_source], functions=['add_cuda'], with_cuda=True, extra_cuda_cflags=['-O3'], verbose=False, ) result = module.add_cuda(torch.randn(1000, device='cuda'), torch.randn(1000, device='cuda'))

load_inline会临时在系统缓存目录里生成一个工程并编译，好处是适合快速试验，坏处是每次编译都要等几十秒，而且不好控制版本依赖。如果你想长期维护一个算子，我建议还是用标准 setup.py，把源码放在项目里，方便版本管理。我在实验阶段用load_inline，验证通过以后立刻迁移到独立工程，这样后面加测试和 benchmark 都很方便。

4. 性能优化与调试

4.1 从正确到高效：访存优化三板斧

很多人第一次写完自定义算子，发现比 PyTorch 原生还慢，就开始怀疑人生。其实绝大多数时候不是 C++ 的问题，而是访存模式没写好。GPU 上的计算能力远超内存带宽，所以逐元素算子基本上都是内存密集型。你写的 kernel 每多读一遍内存，性能就掉一截，优化访存是自定义算子性能提升的大头。

第一板斧是“减少全局内存读写”。这个很好理解，比如y = (a * b) + (a * c)，如果写成两次 kernel，就注定要把a读两遍、中间结果写一遍。正确做法是融合成一个 kernel，只读一次a，算完直接写出去。这种融合只靠 Python 层很难控制，但在自定义 operator 里你可以自由调度。

第二板斧是“向量化访问”。现代 GPU 支持float4这种 128 位宽度的加载，一次能拿 4 个 float。把普通的标量循环改成float4读入，可以有效减少指令数和请求次数。常见做法是把float*转成float4*，不过要注意总线程数和向量宽度的整除关系，通常要处理余数部分。这个优化对带宽受限的算子提升非常可观，有时能从 60% 的利用率涨到 85% 以上。

第三板斧是“使用 grid-stride loop”。当你的 tensor 元素数远超 GPU 能同时启动的线程数时，与其让一个线程只处理一个元素，不如让每个线程循环处理多个间隔的元素，形成一个网格跨步循环。这样做可以减少 block 数量，分摊 kernel 启动开销，同时也有助于更好地利用 GPU 上不同的 SM 调度。最简单的形式是把idx += gridDim.x * blockDim.x写进循环，直到越界。这个模式我几乎在所有 elementwise kernel 里都用，代码量没增加多少，实测却很稳定。

4.2 调试工具与常见坑

自定义算子调试比纯 Python 痛苦，主要是我们突然从 Python 的友好报错切换到了原生世界的空指针、段错误和 CUDA 错误码。我调试的固定套路是：先加cudaDeviceSynchronize()兜底，再开CUDA_LAUNCH_BLOCKING=1。前者确保 kernel 同步，后者让 CUDA 的异步错误变成同步错误，这样崩溃位置能更准确。然后在关键位置用std::cerr <<打印变量形状、数值、指针，确认数据流没断。

如果出现illegal memory access，首选工具是compute-sanitizer。在终端跑：

compute-sanitizer --tool memcheck python test.py

它能报告具体是哪个 kernel 哪一行访问非法。虽然跑起来会慢很多，但找内存问题比盲猜效率高多了。还有一个容易被忽略的是“kernel 启动失败但不报错”，比如blocks参数过大，超过机器上限，或者threads不是 32 的倍数。CUDA 有个规则就是 blockDim.x 最大 1024，且应该是 warp size 的整数倍，你用 100 或者 300 这种值，看起来能跑，实际上可能直接在启动阶段挂掉。所以安全做法是TORCH_CHECK(threads <= 1024)，并且固定用 128、256 或 512。

4.3 和现有 PyTorch 生态的融合

自定义算子最终是要被人调用的。在 Python 层你可以做一个torch.nn.Module把它包起来，也可以做一个普通函数。如果希望它参与 autograd，就用Function包一层，或者依靠自定义 forward 里返回能触发 C++ autograd 的节点。更复杂的场景是torch.compile：它会尝试把 Python 操作图降低到计算图，如果你自定义的算子没有注册进 dispatch 体系，很可能被当成黑盒保留下，这样计算图优化就不能穿越这个算子。所以如果要塞进一个需要被compile的模型，尽量走TORCH_LIBRARY注册路线。

不过按我实际踩坑的经验，torch.compile对自定义 C++ 算子的支持还在快速演进。如果你只是想在一个稳定模型里跑推理，更稳妥的方案是直接用torch.no_grad加普通 function 调用，不要强行交给torch.compile去改。等确认框架版本的支持程度之后再升级集成方式也不迟。

5. 常见问题与排查技巧实录

5.1 版本不匹配与 ABI 问题

自定义算子最让人头疼的，就是“换个环境就编译不过”。我遇到过最典型的是 PyTorch 用 conda 装，CUDA Toolkit 用系统装，两个版本相差一个主版本。编译的时候 nvcc 用了 Toolkit 11.8，但运行时 PyTorch 动态库链接的是 CUDA 12 的 runtime，结果一执行就报CUDA driver version is insufficient，或者干脆符号找不到。我的经验是：torch.version.cuda显示什么，你的 CUDA Toolkit 就尽量用什么，哪怕你的显卡驱动已经支持更新的版本，也别轻易混用。混合版本带来的收益远小于排查成本。

还有 ABI 问题。如果你自己编译的 PyTorch 和预编译的 PyTorch 编译器版本不同，一个用 GCC 9，一个用 GCC 11，链接时会出现一堆undefined symbol: _ZN2at4TensorC1ERKS_这类符号找不到的报错。解决办法也很简单：尽量使用和 PyTorch 官方预编译环境相同的编译器。Linux 下官方一般用 GCC 7.3 到 9.3 的版本范围，Windows 下要求 MSVC 版本不低于 2015。所以在折腾自定义算子之前，先检查编译日志里提示的编译器版本，再决定要不要升级。

5.2 Windows 下 MSVC 的幺蛾子

Windows 下的坑比 Linux 多得多。最常见的是error: identifier "cl" is undefined，或者cl.exe not found。这多半是因为你在普通 cmd 或者 PowerShell 里直接跑python setup.py，而cl.exe只在 Visual Studio 的开发者命令提示符里才有。解决方法是开始菜单里搜索“x64 Native Tools Command Prompt”，打开后进入项目目录再执行编译命令。需要注意的是，Python 默认是 64 位，所以必须用 x64 的命令行，如果用 x86 的命令行，链接时会出现平台冲突。

另一个坑是 Redistributable 的版本。经常有人把 VC++ 运行库装好，但编译器还是缺失，因为 Redistributable 只提供 DLL，编译期工具链需要完整的 Visual Studio Build Tools。如果你用的是 VS Code 或者 JetBrains，千万别以为编辑器能写代码就能编译，C++ 扩展始终需要 MSVC 编译器。我在 Windows 上还遇到过.pyd生成到子目录、import找不到模块的情况，检查一下build_ext --inplace的日志，确认输出路径是否和脚本位置一致。实在不行就手动把生成的.pyd复制到项目根目录。

5.3 多卡分布式里的自坑

当你在DataParallel或者DistributedDataParallel里使用自定义 CUDA 算子时，风险主要来自设备指针和 device guard。CUDA 的 kernel 是按 device 区分的，如果你的函数在某个设备上分配了 Tensor，却没把当前设备切换过去，就可能访问错误的内存。解决方法是使用const c10::cuda::OptionalCUDAGuard device_guard(device_of(x));，保证 kernel 在当前 tensor 所在的设备上运行。

另一个坑是 stream。PyTorch 默认使用各自的 stream 来做异步操作，如果你的自定义 kernel 是在 CUDA default stream 上启动，而前后操作都在当前 PyTorch stream 上，就可能出现未定义顺序，表现为结果时对时错。安全起见，在 host 函数里调用 CUDAGuard 的同时，还需要用at::cuda::CUDAStream将当前 kernel 调度到与 PyTorch 一致的 stream 上。简单做法是：

auto stream = at::cuda::getCurrentCUDAStream(); vector_add_kernel<<<blocks, threads, 0, stream.stream()>>>(...);

这个细节如果漏掉，单卡跑得好好的，多卡一上就开始随机错误，排查起来非常磨人。所以我建议大家把 device guard 和 stream 两件事写进自己的 kernel 模板里，不要等出了问题再补。下面是一个常见问题速查表，整理了我经常遇到的现象和对应的解法：

5.4 ONNX 导出自定义算子（进阶）

如果你的自定义算子最终要部署到 ONNX Runtime，或者要转换成其他推理引擎，直接导出是导不出来的，因为 ONNX 没有你的自定义节点。PyTorch 提供了一个 symbolic 注册机制：把自定义算子映射到既有 ONNX op，或者自定义 domain 下的 op。你需要写一个 symbolic 函数，告诉 TorchScript 在导出时如何表示这个节点。

以最简单的加法为例，如果自定义算子本质上就是 Add，你可以这么写：

from torch.onnx import register_custom_op_symbolic def my_add_symbolic(g, x, y): return g.op('Add', x, y) register_custom_op_symbolic('my_ops::my_combined_op', my_add_symbolic, opset_version=17)

更复杂的算子，比如拉普拉斯算子这种本地滤波，通常会在自定义 domainai.onnx.contrib下生成一个CustomOp节点，然后由推理后端去实现对应 kernel。导出时一般需要传入custom_opsets={"ai.onnx.contrib": 1}。这里要特别注意算子版本的约定：onnx op 版本、PyTorch 版本、后端实现版本三者的对应关系。我在实践中建议：导出前先在小模型上试跑，确认节点格式和 ONNX Runtime 支持的 opset 匹配，否则部署时还是容易踩到unsupported operator的坑。

6. 进阶战术：融合算子与性能基线

6.1 为什么要融合

前面讲了很多理论，最后我用一个具体例子说清楚“融合”的收益。假设我们有三个输入a、b、c，要计算d = a * b + c。如果拆成两条 PyTorch kernel，第一步tmp = a * b，第二步d = tmp + c，内存读写次数是：读a、读b、写tmp，再读tmp、读c、写d，一共至少 6 次全局内存操作。如果写一个融合 kernel，一次循环里同时读a、b、c，算出a*b+c后直接写d，全局内存操作是 4 次，少了三分之一。这种差距在超大 tensor 上尤其明显，因为内存带宽就是瓶颈。

写法其实很简单，就是把两个 kernel 的逻辑合并到一个__global__函数里：

__global__ void fused_mul_add_kernel(const float* a, const float* b, const float* c, float* d, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { d[idx] = a[idx] * b[idx] + c[idx]; } }

编译器通常会把它优化成一条 FMA 指令，计算吞吐也更高。更复杂的融合，例如 softmax 融合 scale 和 masked fill，收益会更明显。不过要注意：融合后不能再从中间结果做别的分支，因为你省掉的就是那份中间结果。这是设计层面需要权衡的。

6.2 性能基线的正确打开方式

写算子一定要测准，不然根本不知道优化有没有效果。我推荐用 PyTorch 自带的torch.profiler配合torch.cuda.Event做 baseline。torch.cuda.Event的用法很简单：

start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() result = fused_op(a, b, c) end.record() torch.cuda.synchronize() print(f"fused: {start.elapsed_time(end):.3f} ms")

注意一定要torch.cuda.synchronize()，不然异步执行会报一个极不合理的时间。除了测总耗时，我还会用torch.profiler看 kernel 的 achieved occupancy、memory throughput、dram bytes read 这类指标，判断瓶颈究竟在计算还是访存。工具输出的数据里，memory throughput 接近峰值而 compute throughput 很低，说明算子已经带宽饱和，再优化计算没有意义。这时候能优化的方向只有减少访存，比如用更低精度、压缩数据，或者改用 shared memory 做数据复用。

测 baseline 时还要注意排除掉第一次调用时的 lazy initialization。PyTorch 的 CUDA context 启动、算子第一次加载都有一笔一次性开销，所以正式测试前先跑几十次 warm-up，让 GPU 内核都准备好，再记录后面若干次的时间。这个过程没人爱做，但只有这么做才能得到可信的数字，否则你会被第一轮的 200ms 欺骗，误以为自己的算子写得很差。我还会把每次测到的耗时都存进一个列表，三次取中位数而不是平均值，因为 GPU 上偶尔会有背景任务干扰，中位数更稳定。

我个人折腾下来最深的感受是：自定义算子的难点从来不在写 kernel，而在于你愿不愿意把 Python 以外的工具链弄明白。当你第一次绕过 GIL、直接在 C++ 层把 Tensor 逻辑跑通，再回头看nn.Sequential就像看说明书。下一步你可以试试把我这里的最小示例改成原子操作的累加，或者给算子加一个 workspace 分配，能力就是这么慢慢长出来的。