介绍在 PyTorch 环境下使用 C++ 和 CUDA 开发自定义算子的方法。针对标准框架性能瓶颈，通过容器化环境简化配置，演示动态加载与静态安装两种构建方式。涵盖 Kernel 编写、Python 绑定、模型集成及自动微分支持。提供最佳实践建议，如合理划分粒度、多类型支持及错误检查，助力深度学习模型优化与部署。
在深度学习模型日益复杂的今天，研究者和工程师常常面临一个共同挑战：标准框架提供的算子已经无法满足特定场景下的性能需求。比如你设计了一个全新的稀疏注意力机制，或者需要对某个小批量操作进行极致优化——此时，用 Python 写的 for 循环显然撑不住训练节奏。
PyTorch 之所以能在科研与工业界同时站稳脚跟，除了其动态图带来的灵活性外，还有一个常被低估的能力：通过 C++ 和 CUDA 编写高性能自定义算子，并无缝接入现有流程。这种能力让你既能享受 Python 的快速原型开发优势，又能在关键路径上'踩到底层'，榨干 GPU 的每一分算力。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online
本文聚焦于 如何在 PyTorch-CUDA 环境下，利用预构建容器镜像快速实现、编译并调用基于 C++/CUDA 的自定义算子。我们不走'先讲理论再贴代码'的套路，而是从实际问题切入，带你一步步打通从编写 kernel 到集成进模型的完整链路。
当你开始关心 GPU 利用率、内核启动开销或显存拷贝次数时，说明你已经走出了'能跑就行'的阶段。这时候你会发现，很多瓶颈其实来自于'组合式'操作：
# 示例：低效的逐元素加法组合
def slow_add(x, y):
z = x + y
mask = (z > 0)
return z * mask.float()
虽然这三行代码简洁明了，但背后涉及三个独立的 CUDA kernel 启动，中间还伴随着冗余内存访问。如果这个操作出现在每层网络中，累积延迟将非常可观。
而如果你把整个逻辑融合成一个 kernel，在 GPU 上一次性完成计算，就能显著减少 launch 开销和 global memory 访问频率。这就是自定义算子的核心价值所在：
更重要的是，借助现代开发工具链（尤其是容器化环境），你现在可以跳过过去令人头疼的环境配置环节，专注在算法本身。
曾经，要编译一个 CUDA 扩展，你需要确保本地安装了正确版本的：
稍有不慎就会遇到 nvcc fatal : Unsupported gpu architecture 'compute_86' 或 undefined symbol: cudnnCreate 这类问题。
而现在，使用官方或社区维护的 PyTorch-CUDA 基础镜像（如 pytorch/pytorch:2.0-cuda11.7-cudnn8-devel），一切都被封装好了。以常见的 v2.8 版本为例，这类镜像通常具备以下特性：
启动命令也很简单：
docker run --gpus all -it \
-p 8888:8888 \
-p 2222:22 \
pytorch-cuda:v2.8
进入容器后，你可以立即验证环境是否正常：
import torch
print(torch.__version__) # 应输出 2.8.x
print(torch.cuda.is_available()) # 应为 True
print(torch.cuda.get_device_name()) # 显示 GPU 型号
一旦确认无误，就可以开始编写你的第一个 CUDA 扩展了。
我们以最简单的张量加法为例，展示如何用 C++ 和 CUDA 实现一个自定义算子。
创建文件 custom_kernel.cu：
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void add_kernel(const float* A, const float* B, float* C, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
C[idx] = A[idx] + B[idx];
}
}
torch::Tensor add_tensors_cuda(torch::Tensor A, torch::Tensor B) {
// 输入检查
TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
TORCH_CHECK(A.size(0) == B.size(0), "Size mismatch between tensors");
int size = A.numel();
auto C = torch::empty_like(A);
dim3 block(256);
dim3 grid((size + block.x - 1) / block.x);
add_kernel<<<grid, block>>>(
A.data_ptr<float>(),
B.data_ptr<float>(),
C.data_ptr<float>(),
size
);
// 注意：仅用于调试！生产代码应异步执行
// cudaDeviceSynchronize();
return C;
}
几点说明：
TORCH_CHECK 是 PyTorch 提供的安全断言宏，比 raw assert 更友好；data_ptr<T>() 直接返回设备指针，无需额外拷贝；cudaDeviceSynchronize()，保持异步性。创建 bindings.cpp 文件，用于暴露接口：
#include <torch/extension.h>
torch::Tensor add_tensors_cuda(torch::Tensor A, torch::Tensor B);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("add", &add_tensors_cuda, "CUDA-accelerated tensor addition");
}
这里的关键是 PYBIND11_MODULE 宏中的 TORCH_EXTENSION_NAME，它会自动替换为 setup 脚本中指定的模块名，避免命名冲突。
对于快速迭代阶段，建议使用 torch.utils.cpp_extension.load() 动态编译并加载模块，无需打包安装。
import torch
from torch.utils.cpp_extension import load
# 动态加载，自动检测变更并重建
custom_add = load(
name="custom_add",
sources=["bindings.cpp", "custom_kernel.cu"],
verbose=True,
extra_cflags=['-O2'],
extra_cuda_cflags=['-O2', '--use_fast_math']
)
load() 函数会在首次运行时触发 JIT 编译，结果缓存在 ~/.cache/torch_extensions/ 中，下次启动时若源码未变则直接加载缓存，极大提升开发效率。
测试一下功能：
a = torch.ones(5).cuda()
b = torch.ones(5).cuda()
c = custom_add.add(a, b)
print(c) # [2., 2., 2., 2., 2.]
是不是和原生 a + b 结果一致？很好。
当你确认算子稳定后，可以将其打包为独立模块，通过 setuptools 安装。
创建 setup.py：
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='custom_add',
ext_modules=[
CUDAExtension(
name='custom_add',
sources=['bindings.cpp', 'custom_kernel.cu'],
extra_compile_args={
'cxx': ['-g', '-O2'],
'nvcc': ['-O2', '--use-fast-math']
}
)
],
cmdclass={
'build_ext': BuildExtension
}
)
然后执行：
pip install -v .
安装完成后即可像普通模块一样导入：
import custom_add
result = custom_add.add(a, b)
这种方式更适合团队协作和 CI/CD 流程。
一旦算子可用，就可以轻松嵌入到 nn.Module 中：
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return custom_add.add(x, x) # 使用自定义算子
model = MyModel().cuda()
x = torch.randn(1000).cuda()
out = model(x)
更进一步，如果你想支持自动微分，只需继承 torch.autograd.Function 并实现前向与反向传播逻辑：
// 在 C++ 中定义 Autograd Function
struct AddFunction : public torch::autograd::Function<AddFunction> {
static torch::Tensor forward(torch::autograd::AutogradContext* ctx, torch::Tensor A, torch::Tensor B) {
return add_tensors_cuda(A, B);
}
static torch::autograd::tensor_list backward(
torch::autograd::AutogradContext* ctx,
torch::autograd::tensor_list grad_outputs) {
return {grad_outputs[0].clone(), grad_outputs[0].clone()};
}
};
绑定后即可参与梯度计算：
a = torch.randn(5, requires_grad=True).cuda()
loss = custom_add.add(a, a).sum()
loss.backward() # 梯度正常回传
在真实项目中，有几个关键点值得注意：
不要试图把整个网络写成一个 kernel。应按功能拆分为可复用的小模块，例如：
这样既便于调试，也利于后期维护。
使用模板化设计，支持 float、half 甚至 bfloat16：
template<typename scalar_t>
__global__ void add_kernel_template(...) {
...
}
torch::Tensor add_tensors_cuda(torch::Tensor A, torch::Tensor B) {
return AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "add", [&] {
add_kernel_template<scalar_t><<<...>>>();
});
}
AT_DISPATCH_* 系列宏是 PyTorch 提供的类型分发工具，强烈推荐使用。
尤其是在多人协作环境中，输入校验能帮你省去大量 debug 时间：
TORCH_CHECK(A.is_cuda(), "Input A must be on GPU");
TORCH_CHECK(A.dim() == 1, "Only 1D tensors supported");
除非调试需要，否则不要在 kernel 后面加 cudaDeviceSynchronize()。它会强制主机等待设备完成，破坏流水线效率。
load() 的缓存机制默认开启，但如果修改了头文件或编译参数，可能需要手动清除：
rm -rf ~/.cache/torch_extensions/
这套基于 PyTorch-CUDA 镜像 + C++/CUDA 扩展 的开发模式，本质上提供了一条'短路径'：
新想法 → 编写 kernel → 动态加载 → 快速验证 → 封装部署
相比传统方式，它解决了几个核心痛点：
更重要的是，这种模式不仅适用于研究人员快速验证新结构，也同样适用于工程师在生产环境中优化关键路径。无论是边缘设备上的低延迟推理，还是大规模集群中的高效训练，自定义算子都已成为不可或缺的技术手段。
未来，随着 Triton、TCO 等更高层次的 DSL 工具发展，我们或许会看到更多'写 Python 语法，跑原生性能'的方案出现。但在当前阶段，掌握 C++/CUDA 扩展仍然是通往极致性能的必经之路。