<<<CUDA C++ Coursegrid · block · warp · lane
EN中文
Stage 5 · Production and AI14 / 15about 26 min

Plug into PyTorch: turn your kernel into an operator

From load_inline for a quick check, to a proper setuptools extension and autograd

After this lesson you can
  • Verify a custom kernel in minutes with load_inline
  • Write a proper setuptools extension and get the tensor checks right
  • Connect autograd to a custom operator
  • Know the correct way to register a custom op in the torch.compile era

A practical use of CUDA is writing custom operators for PyTorch. The C++ extension mechanism is simpler than it looks: your kernel takes `torch::Tensor`; memory management, device placement, and dtype dispatch stay with the framework.

Fastest path: load_inline

Write it in a Python file; compile at runtime
1import torch2from torch.utils.cpp_extension import load_inline3 4cuda_src = r"""5#include <torch/extension.h>6 7__global__ void square_kernel(const float* in, float* out, int n) {8    int i = blockIdx.x * blockDim.x + threadIdx.x;9    if (i < n) out[i] = in[i] * in[i];10}11 12torch::Tensor square(torch::Tensor x) {13    TORCH_CHECK(x.is_cuda(),      "input must be on CUDA");14    TORCH_CHECK(x.is_contiguous(), "input must be contiguous");15    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "float32 only for now");16 17    auto out = torch::empty_like(x);18    int n = x.numel();19    int threads = 256, blocks = (n + threads - 1) / threads;20 21    square_kernel<<<blocks, threads>>>(22        x.data_ptr<float>(), out.data_ptr<float>(), n);23 24    // Use PyTorch's current stream so ordering with other framework ops is correct25    C10_CUDA_KERNEL_LAUNCH_CHECK();26    return out;27}28"""29 30cpp_src = "torch::Tensor square(torch::Tensor x);"31 32mod = load_inline(33    name="my_square",34    cpp_sources=cpp_src,35    cuda_sources=cuda_src,36    functions=["square"],37    extra_cuda_cflags=["-O3", "--use_fast_math"],38    verbose=True,39)40 41x = torch.randn(1_000_000, device="cuda")42torch.testing.assert_close(mod.square(x), x * x)43print("passed")

The proper path: a setuptools extension

src/ops.cpp: C++ glue
1#include <torch/extension.h>2 3// CUDA-side implementation (defined in the .cu file)4torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, double eps);5 6#define CHECK_CUDA(x)  TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")7#define CHECK_CONTIG(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")8#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIG(x)9 10torch::Tensor rmsnorm(torch::Tensor x, torch::Tensor weight, double eps) {11    CHECK_INPUT(x);12    CHECK_INPUT(weight);13    TORCH_CHECK(x.size(-1) == weight.size(0), "last dim must match weight length");14    return rmsnorm_cuda(x, weight, eps);15}16 17PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {18    m.def("rmsnorm", &rmsnorm, "RMSNorm (CUDA)",19          py::arg("x"), py::arg("weight"), py::arg("eps") = 1e-6);20}
src/rmsnorm.cu: one block per row
1#include <torch/extension.h>2#include <c10/cuda/CUDAStream.h>3 4template <typename scalar_t>5__global__ void rmsnorm_kernel(const scalar_t* __restrict__ x,6                               const scalar_t* __restrict__ w,7                               scalar_t* __restrict__ out,8                               int hidden, float eps) {9    int row = blockIdx.x;                     // one block owns one row10    const scalar_t* xr = x + (size_t)row * hidden;11    scalar_t*       orow = out + (size_t)row * hidden;12 13    // Pass 1: sum of squares14    float sum = 0.0f;15    for (int i = threadIdx.x; i < hidden; i += blockDim.x) {16        float v = static_cast<float>(xr[i]);17        sum += v * v;18    }19    // intra-warp reduce → inter-warp reduce20    for (int off = 16; off > 0; off >>= 1)21        sum += __shfl_down_sync(0xffffffff, sum, off);22 23    __shared__ float warpSums[32];24    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;25    if (lane == 0) warpSums[wid] = sum;26    __syncthreads();27 28    if (wid == 0) {29        sum = (lane < (blockDim.x + 31) / 32) ? warpSums[lane] : 0.0f;30        for (int off = 16; off > 0; off >>= 1)31            sum += __shfl_down_sync(0xffffffff, sum, off);32        if (lane == 0) warpSums[0] = rsqrtf(sum / hidden + eps);33    }34    __syncthreads();35    float scale = warpSums[0];36 37    // Pass 2: normalize and scale by weight38    for (int i = threadIdx.x; i < hidden; i += blockDim.x) {39        orow[i] = static_cast<scalar_t>(static_cast<float>(xr[i]) * scale40                                        * static_cast<float>(w[i]));41    }42}43 44torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, double eps) {45    auto out = torch::empty_like(x);46    int hidden = x.size(-1);47    int rows = x.numel() / hidden;48 49    int threads = std::min(1024, ((hidden + 31) / 32) * 32);50    auto stream = at::cuda::getCurrentCUDAStream();     // use PyTorch's stream!51 52    AT_DISPATCH_FLOATING_TYPES_AND2(53        at::ScalarType::Half, at::ScalarType::BFloat16,54        x.scalar_type(), "rmsnorm_cuda", [&] {55            rmsnorm_kernel<scalar_t><<<rows, threads, 0, stream>>>(56                x.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),57                out.data_ptr<scalar_t>(), hidden, static_cast<float>(eps));58        });59    return out;60}
setup.py
1from setuptools import setup2from torch.utils.cpp_extension import BuildExtension, CUDAExtension3 4setup(5    name="my_ops",6    ext_modules=[7        CUDAExtension(8            name="my_ops._C",9            sources=["src/ops.cpp", "src/rmsnorm.cu"],10            extra_compile_args={11                "cxx": ["-O3"],12                "nvcc": ["-O3", "--use_fast_math",13                         "-gencode", "arch=compute_80,code=sm_80",14                         "-gencode", "arch=compute_86,code=sm_86"],15            },16        )17    ],18    cmdclass={"build_ext": BuildExtension},19)

Connect autograd

A custom Function
1import torch2from my_ops import _C3 4class RMSNormFn(torch.autograd.Function):5    @staticmethod6    def forward(ctx, x, weight, eps=1e-6):7        out = _C.rmsnorm(x, weight, eps)8        ctx.save_for_backward(x, weight)9        ctx.eps = eps10        return out11 12    @staticmethod13    def backward(ctx, grad_out):14        x, weight = ctx.saved_tensors15        # In production this should also be a CUDA kernel.16        # During development, implement it with PyTorch ops first, lock the numerics, then replace.17        grad_x, grad_w = _C.rmsnorm_backward(grad_out, x, weight, ctx.eps)18        return grad_x, grad_w, None19 20def rmsnorm(x, weight, eps=1e-6):21    return RMSNormFn.apply(x, weight, eps)22 23# Numerical gradient check: required for every new op24x = torch.randn(4, 128, device="cuda", dtype=torch.double, requires_grad=True)25w = torch.randn(128, device="cuda", dtype=torch.double, requires_grad=True)26assert torch.autograd.gradcheck(rmsnorm, (x, w), eps=1e-6, atol=1e-4)

The torch.compile era: register as a custom operator

If the model goes through torch.compile, calling a pybind-exported function causes a graph break: the compiler cannot include the op in its optimizations. Register it with torch.library so it is a first-class PyTorch operator, and provide a FakeTensor implementation (output shape and dtype, used during graph capture).

Registering with torch.library
1import torch2from torch.library import custom_op, register_fake3 4@custom_op("my_ops::rmsnorm", mutates_args=())5def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:6    from my_ops import _C7    return _C.rmsnorm(x, weight, eps)8 9@register_fake("my_ops::rmsnorm")10def _(x, weight, eps=1e-6):11    # Describe output metadata only; do not compute12    return torch.empty_like(x)13 14# torch.compile can now capture it; no graph break15compiled = torch.compile(lambda a, b: torch.ops.my_ops.rmsnorm(a, b))

Self-check

Answer in your head first, then open the solution. If you cannot, reread that section.