Skip to content

Custom C++ / CUDA Extensions

Overview

When the Python API can't express an op with acceptable performance (or at all), you drop to C++ extensions and, when needed, CUDA kernels. These plug straight back into PyTorch as first-class ops — autograd-compatible if you wrap them in torch.autograd.Function.

Topics

  • Building C++ Extensions - torch.utils.cpp_extension.load, CppExtension, and the module skeleton.
  • Writing Custom CUDA Kernels - The minimal kernel, launch config, and what CUDA gives you.
  • Extensions in Practice & Pitfalls - Build pain, gradients, shapes, determinism, debugging.

Key Patterns

  • cpp_extension.load() = one-file build (CPU); CUDAExtension = GPU kernels via setup.py.
  • Register with torch::Tensor in/out; use TORCH_CHECK for shape guards.
  • Wrap native fn in torch.autograd.Function for gradients (Ch 03-01).
  • Profile your kernel; naive kernels are often slower than PyTorch's own.