不能直接用torch.cuda写CUDA算子，因其Python API不提供裸指针访问或kernel启动能力，所有操作均经ATen抽象层；必须通过C++/CUDA扩展机制，在C++ host code中调用kernel并正确绑定张量指针、stream和设备上下文。
为什么不能直接用 torch.cuda 写 CUDA 算子
PyTorch 的 Python API（比如 torch.cuda）不提供裸设备指针访问或 kernel 启动能力，所有张量操作都走 ATen 抽象层。想写自定义 CUDA kernel，必须绕过 Python 层，进入 C++/CUDA 原生上下文——也就是通过 PyTorch 的 C++ 扩展机制（torch.utils.cpp_extension 提供的构建支持），把编译好的 device code 挂到 Python 可调用的函数上。
常见错误是试图在 Python 里用 ctypes 加载 .cu 编译的 so，或误以为 torch.cuda.stream 能直接 launch kernel：都不行。CUDA kernel 必须由 C++ host code 显式调用，且需正确绑定张量内存、流、设备索引。
如何组织 setup.py 和 .cu 文件结构
PyTorch 官方推荐用 load（开发调试）或 setup.py（发布）方式编译 C++/CUDA 扩展。关键点不是“能不能编”，而是“能不能让 PyTorch 正确识别 CUDA 设备上下文并传递 stream”。
-
setup.py中必须显式启用 CUDA 支持：CUDAExtension替代CppExtension，并传入sources=['my_op.cu'] -
.cu文件里不能直接写cudaMemcpy或cudaMalloc—— 张量内存已由 PyTorch 分配，只需用tensor.data_ptr<float>()获取指针 - 必须包含
#include <ATen/cuda/CUDAContext.h>，并在 kernel launch 前调用at::cuda::getCurrentCUDAStream()获取当前 stream - kernel 调用必须带
stream参数，例如：my_kernel<float><<<grid, block, 0, stream>>>(…)
示例片段（my_op.cu）：
立即学习“Python免费学习笔记（深入）”；
__global__ void add_kernel(float* a, float* b, float* out, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) out[idx] = a[idx] + b[idx];
}
<p>torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
auto output = torch::empty_like(a);
const int n = a.numel();
const int block = 256;
const int grid = (n + block - 1) / block;
auto stream = at::cuda::getCurrentCUDAStream();
add_kernel<<<grid, block, 0, stream>>>(
a.data_ptr<float>(),
b.data_ptr<float>(),
output.data_ptr<float>(),
n
);
return output;
}Python 端调用时怎么保证设备和 stream 对齐
最常踩的坑：输入张量在 GPU 上，但没指定 device 或 non_blocking=True，导致 host-device 同步阻塞；或多个算子间 stream 不一致，引发 race condition。
- 所有输入
torch.Tensor必须已位于 CUDA 设备上（a.is_cuda == True），否则data_ptr返回的是无效地址 - 调用前建议显式同步一次：
torch.cuda.synchronize()（仅调试用），正式代码应依赖 stream 顺序，而非全局同步 - 若需跨 stream 协作（如前一个算子输出进下一个自定义 kernel），务必用
torch.cuda.Stream显式管理，并在 C++ 侧用at::cuda::setCurrentCUDAStream(stream)切换上下文 - PyTorch 默认使用 per-thread default stream，所以只要 Python 端没手动切 stream，C++ 侧
getCurrentCUDAStream()就能拿到正确值
编译失败常见报错及对应修复
报错往往不是语法问题，而是环境链路断裂。下面几个错误出现频率极高：
-
nvcc fatal : Unsupported gpu architecture 'compute_86'：CUDA 版本太旧，不支持你 GPU 的 compute capability。查nvidia-smi→cuda-smi→ 查表确认 arch，然后在CUDAExtension中加extra_compile_args={'nvcc': ['-gencode', 'arch=compute_80,code=sm_80']} -
undefined symbol: _ZN2at6native13get_cudnn_mode...：链接时没找到 ATen CUDA 符号。确保from torch.utils.cpp_extension import CUDAExtension被导入，且 setup.py 中ext_modules=[CUDAExtension(...)]正确注册 - Python 导入后调用报
RuntimeError: Expected all tensors to be on the same device：C++ 函数里没检查输入设备一致性，建议开头加AT_ASSERTM(a.is_cuda() && b.is_cuda(), "all inputs must be CUDA tensors") - kernel 结果全为 0 或乱码：大概率是
data_ptr<T>()类型与 kernel 模板类型不匹配（如 tensor 是torch.float64，但 kernel 用float*），务必用a.scalar_type()校验
复杂点不在 kernel 逻辑本身，而在于 device context、stream scope、scalar type、memory layout（尤其是 stride 不为 1 的情况）四者必须同时对齐。少一个，结果就不可预测。
