PyTorch 自定义算子开发：使用 C++ 与 CUDA 扩展

进入容器后，你可以立即验证环境是否正常：

import torch
print(torch.__version__)  # 应输出 2.8.x
print(torch.cuda.is_available())  # 应为 True
print(torch.cuda.get_device_name())  # 显示 GPU 型号

一旦确认无误，就可以开始编写你的第一个 CUDA 扩展了。

编写你的第一个 CUDA 扩展

我们以最简单的张量加法为例，展示如何用 C++ 和 CUDA 实现一个自定义算子。

1. 核心 CUDA Kernel

创建文件 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>() 直接返回设备指针，无需额外拷贝；
• 块大小设为 256 是常见选择，可根据具体硬件调整；
• 生产环境中应去掉 cudaDeviceSynchronize()，保持异步性。

2. 绑定到 Python 接口

创建 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 → 动态加载 → 快速验证 → 封装部署

相比传统方式，它解决了几个核心痛点：

• ❌ 环境配置复杂 → ✅ 一键启动容器
• ❌ 编译失败频繁 → ✅ 工具链预对齐
• ❌ 调试效率低下 → ✅ 支持热重载 + Jupyter 可视化
• ❌ 性能难以压榨 → ✅ 直达底层 CUDA，零解释开销

更重要的是，这种模式不仅适用于研究人员快速验证新结构，也同样适用于工程师在生产环境中优化关键路径。无论是边缘设备上的低延迟推理，还是大规模集群中的高效训练，自定义算子都已成为不可或缺的技术手段。

未来，随着 Triton、TCO 等更高层次的 DSL 工具发展，我们或许会看到更多'写 Python 语法，跑原生性能'的方案出现。但在当前阶段，掌握 C++/CUDA 扩展仍然是通往极致性能的必经之路。