type
Post
status
Published
date
May 10, 2026
slug
summary
PyTorch 自定义 CUDA 算子编写、绑定与调用示例教程
tags
pytorch
cuda
category
技术分享
icon
password
一个常见需求是：在 PyTorch 中调用自己编写的 CUDA 算子（custom CUDA op），以便验证算子正确性并进行性能评估。下面我以一个最简单的
add 算子为例，演示如何将 CUDA 与 PyTorch Extension 结合起来，从 0 到 1 完成编写、绑定与调用。目录结构
假设所有文件都放在
operator_extension/ 目录下：add_kernel.cu：CUDA kernel 与 launcher
op.h：对外暴露的 C++ 声明（供.cpp调用）
op.cpp：PyTorch C++ 扩展入口与 pybind11 绑定
1. 编写 CUDA kernel：add_kernel.cu
这里将计算逻辑放在
__global__ kernel 中，并通过 launch_add_op 作为一个干净的 C++ 接口对外暴露，方便后续从 PyTorch 的 C++ 扩展层调用。2. 声明 launcher 接口：op.h
op.h 主要用于给 PyTorch C++ 侧（.cpp）提供函数声明：3. 编写 PyTorch 扩展入口与绑定：op.cpp
在
op.cpp 中完成三件事：- 校验输入 Tensor 是否为 CUDA 且连续
- 分配输出 Tensor 并调用 CUDA launcher
- 通过 pybind11 将函数暴露给 Python
这样分层设计的好处
- CUDA kernel（.cu）与 C++/PyTorch 绑定（.cpp）解耦
Kernel 只关注并行计算；接口层只关注 Tensor 校验、形状推导与 Python 绑定。后续更换绑定方式（pybind / ATen / TorchScript）时，kernel 基本无需改动。
- 编译与调试更清晰
.cu 由 NVCC 编译，.cpp 由 C++ 编译器编译，出错信息与编译选项边界更明确；也更便于对 kernel 单独做性能分析（nsys/nvprof）。- 复用与扩展性更强
同一个 launcher（如
launch_add_op）可以复用到更多算子或更多前端（Python / C++）。新增算子通常只需要在接口层扩展即可。4. Python 侧编译与调用
使用
torch.utils.cpp_extension.load 进行即时编译（JIT build）：运行后如果输出
Test passed!，说明自定义 CUDA 算子与 PyTorch 的结果一致，你就可以在此基础上继续扩展更复杂的算子逻辑了。