Skand Tandon
Deep Implementation

PyTorch custom op with CUDA backend

I wrote a fused GELU kernel in CUDA, registered it as a PyTorch custom op, and benchmarked it against the built-in. The performance was not the point. Understanding the boundary between Python and the GPU was.

2025-11-15 deep implementation 9 min read

After profiling a transformer forward pass, I noticed that elementwise ops like GELU and layer norm are individually small but collectively significant. Each one requires a separate kernel launch and a round trip to HBM. Fusing them into neighboring ops is a well-known optimization. But before I could reason about fusion, I wanted to understand the full path from "Python calls a function" to "GPU executes a kernel." So I wrote one from scratch.

What a custom op actually requires

A PyTorch custom op with a CUDA backend involves three pieces:

  1. The CUDA kernel: a __global__ function written in CUDA C++ that runs on the GPU.
  2. The C++ wrapper: a host-side function that sets up the launch configuration (grid size, block size) and calls the kernel.
  3. The Python binding: registration through PyTorch's torch.library or the older torch.utils.cpp_extension to make it callable from Python.

I chose GELU as my target because it is simple enough to implement in a single kernel but common enough in transformers to be worth thinking about. The exact GELU formula is:

GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))

In practice, most implementations use the tanh approximation for speed:

GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))

The CUDA kernel

Here is the kernel. Each thread processes one element:

// gelu_kernel.cu
#include <torch/extension.h>
#include <cuda_fp16.h>
#include <math_constants.h>

__global__ void gelu_forward_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float x = input[idx];
        // tanh approximation
        float cdf = 0.5f * (1.0f + tanhf(
            0.7978845608f * (x + 0.044715f * x * x * x)
        ));
        output[idx] = x * cdf;
    }
}

torch::Tensor gelu_forward_cuda(torch::Tensor input) {
    auto output = torch::empty_like(input);
    int n = input.numel();
    int threads = 256;
    int blocks = (n + threads - 1) / threads;

    gelu_forward_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        n
    );
    return output;
}

A few things worth noting:

Building with cpp_extension

PyTorch's cpp_extension module handles the compilation. You can either use ahead-of-time compilation with setup.py or just-in-time compilation with load:

from torch.utils.cpp_extension import load

gelu_cuda = load(
    name="gelu_cuda",
    sources=["gelu_kernel.cu"],
    verbose=True,
)

The first call compiles the CUDA code with nvcc, links it against PyTorch's C++ libraries, and caches the result. Subsequent calls load from cache. The verbose=True flag is worth keeping on during development because nvcc errors are notoriously cryptic, and seeing the full compilation command helps debug them.

For the binding to work, you also need a PYBIND11_MODULE block:

// At the bottom of gelu_kernel.cu
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &gelu_forward_cuda, "GELU forward (CUDA)");
}

The modern way: torch.library

Since PyTorch 2.1, the recommended approach is torch.library, which integrates with torch.compile and the dispatcher. You define the op signature, register implementations per backend, and the dispatcher routes calls to the right kernel:

# Python-side registration
import torch

torch.library.define("myops::gelu", "(Tensor x) -> Tensor")

@torch.library.impl("myops::gelu", "cuda")
def gelu_cuda_impl(x):
    return gelu_cuda.forward(x)

@torch.library.impl("myops::gelu", "cpu")
def gelu_cpu_impl(x):
    return torch.nn.functional.gelu(x)  # fallback

This has a real advantage: your custom op works with torch.compile, autograd (if you register a backward), and device-agnostic code. The old cpp_extension approach gives you a raw function; the torch.library approach gives you a first-class PyTorch operation.

Benchmarking against the built-in

I benchmarked my naive kernel against torch.nn.functional.gelu on tensors of various sizes:

import torch
import time

def bench(fn, x, warmup=50, iters=200):
    for _ in range(warmup):
        fn(x)
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(iters):
        fn(x)
    torch.cuda.synchronize()
    return (time.perf_counter() - start) / iters * 1e6  # microseconds

for size in [1024, 65536, 1048576, 16777216]:
    x = torch.randn(size, device="cuda")
    t_builtin = bench(torch.nn.functional.gelu, x)
    t_custom = bench(gelu_cuda.forward, x)
    print(f"n={size:>10,}  builtin={t_builtin:.1f}us  custom={t_custom:.1f}us  "
          f"ratio={t_custom/t_builtin:.2f}x")

Results on an A100:

My naive kernel is within 5% of PyTorch's built-in at large sizes. This is not because I wrote good CUDA. It is because GELU is so simple that there is almost nothing to optimize. The kernel is bandwidth-bound: read one float, do a few math ops, write one float. The bottleneck is memory, and both kernels saturate it.

The real lesson

For bandwidth-bound kernels, the performance ceiling is HBM bandwidth, and a naive implementation gets close to that ceiling. The wins come from fusion: combining multiple bandwidth-bound ops into one kernel so you read and write to HBM once instead of multiple times. A fused bias-GELU-dropout kernel does three operations for the memory cost of one.

Where this matters for inference

Writing custom CUDA kernels is not something most inference engineers do daily. But understanding the mechanism matters for several reasons:

You do not need to write CUDA kernels for your day job. But you need to be able to read them, and you need to understand the boundary between the Python world and the GPU world. That boundary is where inference performance lives.

With kernels and profiling covered, it is time to move up the stack. On day 46, I deploy vLLM and benchmark the metrics that actually matter in production: time to first token and throughput.